From 4cc5c3e39e44724b5d5035cc5c01d206626a061d Mon Sep 17 00:00:00 2001 From: Karl Date: Mon, 8 Jan 2024 09:20:36 +0000 Subject: [PATCH] cwc --- .gitignore | 3 ++- .vscode/launch.json | 15 ++++++----- auto_subtitle/cli.py | 2 -- auto_subtitle/main.py | 37 +++++++++++++++----------- auto_subtitle/utils/bazarr.py | 17 ++++++++++-- auto_subtitle/utils/ffmpeg.py | 50 ++++++++--------------------------- auto_subtitle/utils/sonarr.py | 20 ++++++++++++++ 7 files changed, 78 insertions(+), 66 deletions(-) create mode 100644 auto_subtitle/utils/sonarr.py diff --git a/.gitignore b/.gitignore index b5b2650..4ba0a43 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ dist build __pycache__ venv/ -test/ \ No newline at end of file +test/ +.vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json index 150004a..d1f22a4 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,16 +12,19 @@ "console": "integratedTerminal", "justMyCode": false, "args": [ - "Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv", - // "--model", - // "base", + // "Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv", + "--model", + "base", // "--srt_only", // "TRUE", - "--output_srt", - "TRUE", + // "--output_srt", + // "TRUE", "-o", "./test" - ] + ], + "env": { + "token": "" + } } ] } \ No newline at end of file diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py index 2c0c8b7..e1ec10b 100644 --- a/auto_subtitle/cli.py +++ b/auto_subtitle/cli.py @@ -14,8 +14,6 @@ def main(): """ parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("video", nargs="+", type=str, - help="paths to video files to transcribe") parser.add_argument("--audio_channel", default="0", type=int, help="audio channel index to use") parser.add_argument("--sample_interval", type=str2timeinterval, default=None, diff --git a/auto_subtitle/main.py b/auto_subtitle/main.py index 5fa438e..3102ce6 100644 --- a/auto_subtitle/main.py +++ b/auto_subtitle/main.py @@ -1,9 +1,11 @@ import os import warnings import tempfile +import time from utils.files import filename, write_srt -from utils.ffmpeg import get_audio, overlay_subtitles, add_subs_new -from utils.bazarr import get_wanted_episodes, get_episode_details +from utils.ffmpeg import get_audio, add_subs_new +from utils.bazarr import get_wanted_episodes, get_episode_details, sync_series +from utils.sonarr import update_show_in_soarr from utils.whisper import WhisperAI @@ -14,6 +16,7 @@ def process(args: dict): srt_only: bool = args.pop("srt_only") language: str = args.pop("language") sample_interval: str = args.pop("sample_interval") + audio_channel: str = args.pop('audio_channel') os.makedirs(output_dir, exist_ok=True) @@ -24,27 +27,29 @@ def process(args: dict): # if translate task used and language argument is set, then use it elif language != "auto": args["language"] = language - - a = get_wanted_episodes() - print(f"Found {a['total']} episodes needing subtitles.") - for episode in a['data']: - episode_data = get_episode_details(episode['sonarrEpisodeId']) - print(episode_data) - audios = get_audio(args.pop("video"), args.pop( - 'audio_channel'), sample_interval) model_args = {} model_args["model_size_or_path"] = model_name model_args["device"] = args.pop("device") model_args["compute_type"] = args.pop("compute_type") + + list_of_episodes_needing_subtitles = get_wanted_episodes() + print(f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles.") + for episode in list_of_episodes_needing_subtitles['data']: + print(f"Processing {episode['seriesTitle']} - {episode['episode_number']}") + episode_data = get_episode_details(episode['sonarrEpisodeId']) + audios = get_audio([episode_data['path']], audio_channel, sample_interval) + srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir() + subtitles = get_subtitles(audios, srt_output_dir, model_args, args) - srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir() - subtitles = get_subtitles(audios, srt_output_dir, model_args, args) + if srt_only: + return - if srt_only: - return - - add_subs_new(subtitles, output_dir, sample_interval) + add_subs_new(subtitles) + update_show_in_soarr(episode['sonarrSeriesId']) + time.sleep(5) + sync_series() + def get_subtitles(audio_paths: list, output_dir: str, diff --git a/auto_subtitle/utils/bazarr.py b/auto_subtitle/utils/bazarr.py index e7c643c..0e8b08b 100644 --- a/auto_subtitle/utils/bazarr.py +++ b/auto_subtitle/utils/bazarr.py @@ -1,6 +1,6 @@ import requests - -token = '' +import os +token = os.getenv('bazarr_token') def get_wanted_episodes(): url = "http://192.168.4.23/api/episodes/wanted" @@ -26,4 +26,17 @@ def get_episode_details(episode_id: str): } response = requests.request("GET", url, headers=headers, data=payload) + return response.json()['data'][0] + + +def sync_series(): + url = f"http://192.168.4.23/api/system/tasks?taskid=update_series" + + payload={} + headers = { + 'accept': 'application/json', + 'X-API-KEY': token + } + + response = requests.request("POST", url, headers=headers, data=payload) return response.json()['data'][0] \ No newline at end of file diff --git a/auto_subtitle/utils/ffmpeg.py b/auto_subtitle/utils/ffmpeg.py index 0c9f91f..a7c2c31 100644 --- a/auto_subtitle/utils/ffmpeg.py +++ b/auto_subtitle/utils/ffmpeg.py @@ -37,47 +37,19 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list): return audio_paths -def overlay_subtitles(subtitles: dict, output_dir: str, sample_interval: list): - for path, srt_path in subtitles.items(): - out_path = os.path.join(output_dir, f"{filename(path)}.mp4") - - print(f"Adding subtitles to {filename(path)}...") - - ffmpeg_input_args = {} - if sample_interval is not None: - ffmpeg_input_args['ss'] = str(sample_interval[0]) - - ffmpeg_output_args = {} - if sample_interval is not None: - ffmpeg_output_args['t'] = str( - sample_interval[1] - sample_interval[0]) - - # HACK: On Windows it's impossible to use absolute subtitle file path with ffmpeg - # so we use temp copy instead - # see: https://github.com/kkroening/ffmpeg-python/issues/745 - with MyTempFile(srt_path) as srt_temp: - video = ffmpeg.input(path, **ffmpeg_input_args) - audio = video.audio - - ffmpeg.concat( - video.filter( - 'subtitles', srt_temp.tmp_file_path, - force_style="OutlineColour=&H40000000,BorderStyle=3"), audio, v=1, a=1 - ).output(out_path, **ffmpeg_output_args).run(quiet=True, overwrite_output=True) - - print(f"Saved subtitled video to {os.path.abspath(out_path)}.") - - -def add_subs_new(subtitles: dict, output_dir: str, sample_interval: list): - import ffmpeg as fmp +def add_subs_new(subtitles: dict): + input_file = list(subtitles.keys())[0] subtitle_file = subtitles[input_file] - output_file = 'class.mp4' + output_file = input_file + os.rename(input_file, input_file+'_edit') - input_stream = fmp.input(input_file) - subtitle_stream = fmp.input(subtitle_file) + input_stream = ffmpeg.input(input_file+'_edit') + subtitle_stream = ffmpeg.input(subtitle_file) # Combine input video and subtitle - output = fmp.output(input_stream, subtitle_stream, output_dir + '/' + output_file.replace('.mkv','.mp4'), c='copy', **{'c:s': 'mov_text'}, **{'metadata:s:s:0': 'language=eng'}) - - fmp.run(output) + output = ffmpeg.output(input_stream, subtitle_stream, output_file.replace('.mkv','.mp4'), c='copy', **{'c:s': 'mov_text'}, **{'metadata:s:s:0': 'language=eng'}) + ffmpeg.run(output, quiet=True, overwrite_output=True) + os.remove(input_file+'_edit') + if '.mkv' in output_file: + os.remove(output_file) \ No newline at end of file diff --git a/auto_subtitle/utils/sonarr.py b/auto_subtitle/utils/sonarr.py new file mode 100644 index 0000000..245edea --- /dev/null +++ b/auto_subtitle/utils/sonarr.py @@ -0,0 +1,20 @@ +import requests +import json + + +def update_show_in_soarr(show_id): + url = "http://192.168.4.9:8989/api/v3/command" + + payload = json.dumps({ + "name": "RefreshSeries", + "seriesId": show_id + }) + headers = { + 'Content-Type': 'application/json', + 'X-Api-Key': 'f6ea49a75a44469daec03969bdf6764d', + } + + response = requests.request("POST", url, headers=headers, data=payload) + + if response.status_code != 404: + print("Updated show in Sonarr") \ No newline at end of file