diff --git a/.gitignore b/.gitignore index 1c79a48..4ba0a43 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,7 @@ dist .DS_Store *.egg-info build -__pycache__ \ No newline at end of file +__pycache__ +venv/ +test/ +.vscode/launch.json diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..d1f22a4 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,30 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Python: Current File", + "type": "python", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal", + "justMyCode": false, + "args": [ + // "Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv", + "--model", + "base", + // "--srt_only", + // "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 6e030f5..e1ec10b 100644 --- a/auto_subtitle/cli.py +++ b/auto_subtitle/cli.py @@ -1,8 +1,8 @@ import argparse from faster_whisper import available_models -from .utils.constants import LANGUAGE_CODES -from .main import process -from .utils.convert import str2bool, str2timeinterval +from utils.constants import LANGUAGE_CODES +from main import process +from utils.convert import str2bool, str2timeinterval def main(): @@ -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 cad112f..3102ce6 100644 --- a/auto_subtitle/main.py +++ b/auto_subtitle/main.py @@ -1,9 +1,12 @@ import os import warnings import tempfile -from .utils.files import filename, write_srt -from .utils.ffmpeg import get_audio, overlay_subtitles -from .utils.whisper import WhisperAI +import time +from utils.files import filename, write_srt +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 def process(args: dict): @@ -13,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,21 +28,28 @@ def process(args: dict): elif language != "auto": args["language"] = language - 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 - - overlay_subtitles(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 new file mode 100644 index 0000000..0e8b08b --- /dev/null +++ b/auto_subtitle/utils/bazarr.py @@ -0,0 +1,42 @@ +import requests +import os +token = os.getenv('bazarr_token') + +def get_wanted_episodes(): + url = "http://192.168.4.23/api/episodes/wanted" + + payload={} + headers = { + 'accept': 'application/json', + 'X-API-KEY': token + } + + response = requests.request("GET", url, headers=headers, data=payload) + + return response.json() + + +def get_episode_details(episode_id: str): + url = f"http://192.168.4.23/api/episodes?episodeid%5B%5D={episode_id}" + + payload={} + headers = { + 'accept': 'application/json', + 'X-API-KEY': token + } + + 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 9f6fdd4..a7c2c31 100644 --- a/auto_subtitle/utils/ffmpeg.py +++ b/auto_subtitle/utils/ffmpeg.py @@ -37,32 +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") +def add_subs_new(subtitles: dict): + + input_file = list(subtitles.keys())[0] + subtitle_file = subtitles[input_file] + output_file = input_file + os.rename(input_file, input_file+'_edit') - print(f"Adding subtitles to {filename(path)}...") + input_stream = ffmpeg.input(input_file+'_edit') + subtitle_stream = ffmpeg.input(subtitle_file) - 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)}.") + # Combine input video and subtitle + 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..e6282ae --- /dev/null +++ b/auto_subtitle/utils/sonarr.py @@ -0,0 +1,21 @@ +import requests +import json +import os +token = os.getenv('sonarr_token') + +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': token, + } + + 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