diff --git a/.gitignore b/.gitignore index 4ba0a43..689e0d9 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ venv/ test/ .vscode/launch.json +config.cfg diff --git a/.vscode/launch.json b/.vscode/launch.json index d1f22a4..94c9cc5 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,19 +12,9 @@ "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/README.md b/README.md index f14f278..59eeb8e 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ Clunky, and slow, but works. ## Usage -The following command will generate a `subtitled/video.mp4` file contained the input video with overlayed subtitles. + diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py index e1ec10b..85e2f2b 100644 --- a/auto_subtitle/cli.py +++ b/auto_subtitle/cli.py @@ -29,12 +29,6 @@ def main(): "int16", "float16", "bfloat16", "float32"], help="Type to use for computation. \ See https://opennmt.net/CTranslate2/quantization.html.") - parser.add_argument("--output_dir", "-o", type=str, - default=".", help="directory to save the outputs") - parser.add_argument("--output_srt", type=str2bool, default=False, - help="whether to output the .srt file along with the video files") - parser.add_argument("--srt_only", type=str2bool, default=False, - help="only generate the .srt file and not create overlayed video") parser.add_argument("--beam_size", type=int, default=5, help="model parameter, tweak to increase accuracy") parser.add_argument("--no_speech_threshold", type=float, default=0.6, diff --git a/auto_subtitle/main.py b/auto_subtitle/main.py index 3102ce6..90e4bc2 100644 --- a/auto_subtitle/main.py +++ b/auto_subtitle/main.py @@ -3,7 +3,7 @@ import warnings import tempfile import time from utils.files import filename, write_srt -from utils.ffmpeg import get_audio, add_subs_new +from utils.ffmpeg import get_audio, add_subtitles_to_mp4 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 @@ -11,15 +11,10 @@ from utils.whisper import WhisperAI def process(args: dict): model_name: str = args.pop("model") - output_dir: str = args.pop("output_dir") - output_srt: bool = args.pop("output_srt") - 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) - if model_name.endswith(".en"): warnings.warn( f"{model_name} is an English-only model, forcing English detection.") @@ -39,19 +34,14 @@ def process(args: dict): 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) + subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args) - if srt_only: - return - - add_subs_new(subtitles) + add_subtitles_to_mp4(subtitles) update_show_in_soarr(episode['sonarrSeriesId']) time.sleep(5) sync_series() - def get_subtitles(audio_paths: list, output_dir: str, model_args: dict, transcribe_args: dict): model = WhisperAI(model_args, transcribe_args) diff --git a/auto_subtitle/utils/bazarr.py b/auto_subtitle/utils/bazarr.py index 0e8b08b..e557c0e 100644 --- a/auto_subtitle/utils/bazarr.py +++ b/auto_subtitle/utils/bazarr.py @@ -1,9 +1,13 @@ import requests -import os -token = os.getenv('bazarr_token') +import configparser +config = configparser.RawConfigParser() +config.read('config.cfg') + +token = config._sections['bazarr']['token'] +base_url = config._sections['bazarr']['url'] def get_wanted_episodes(): - url = "http://192.168.4.23/api/episodes/wanted" + url = f"{base_url}/api/episodes/wanted" payload={} headers = { @@ -17,7 +21,7 @@ def get_wanted_episodes(): def get_episode_details(episode_id: str): - url = f"http://192.168.4.23/api/episodes?episodeid%5B%5D={episode_id}" + url = f"{base_url}/api/episodes?episodeid%5B%5D={episode_id}" payload={} headers = { @@ -30,7 +34,7 @@ def get_episode_details(episode_id: str): def sync_series(): - url = f"http://192.168.4.23/api/system/tasks?taskid=update_series" + url = f"{base_url}/api/system/tasks?taskid=update_series" payload={} headers = { @@ -39,4 +43,5 @@ def sync_series(): } response = requests.request("POST", url, headers=headers, data=payload) - return response.json()['data'][0] \ No newline at end of file + if response.status_code == 204: + print('Updated Bazarr') diff --git a/auto_subtitle/utils/ffmpeg.py b/auto_subtitle/utils/ffmpeg.py index a7c2c31..2ce317f 100644 --- a/auto_subtitle/utils/ffmpeg.py +++ b/auto_subtitle/utils/ffmpeg.py @@ -1,7 +1,6 @@ import os import tempfile import ffmpeg -from .mytempfile import MyTempFile from .files import filename @@ -37,7 +36,7 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list): return audio_paths -def add_subs_new(subtitles: dict): +def add_subtitles_to_mp4(subtitles: dict): input_file = list(subtitles.keys())[0] subtitle_file = subtitles[input_file] @@ -51,5 +50,6 @@ def add_subs_new(subtitles: dict): 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 + # remove tempfiles + os.remove(subtitle_file) + os.remove(subtitle_file.replace(".srt",".wav")) \ No newline at end of file diff --git a/auto_subtitle/utils/mytempfile.py b/auto_subtitle/utils/mytempfile.py deleted file mode 100644 index 372c74d..0000000 --- a/auto_subtitle/utils/mytempfile.py +++ /dev/null @@ -1,35 +0,0 @@ -import tempfile -import os -import shutil - -class MyTempFile: - """ - A context manager for creating a temporary file in current directory, copying the content from - a specified file, and handling cleanup operations upon exiting the context. - - Usage: - ```python - with MyTempFile(file_path) as temp_file_manager: - # Access the temporary file using temp_file_manager.tmp_file - # ... - # The temporary file is automatically closed and removed upon exiting the context. - ``` - - Args: - - file_path (str): The path to the file whose content will be copied to the temporary file. - """ - def __init__(self, file_path): - self.file_path = file_path - self.tmp_file = None - self.tmp_file_path = None - - def __enter__(self): - self.tmp_file = tempfile.NamedTemporaryFile('w', dir='.', delete=False) - self.tmp_file_path = os.path.relpath(self.tmp_file.name, '.') - shutil.copyfile(self.file_path, self.tmp_file_path) - return self - - def __exit__(self, exc_type, exc_value, exc_traceback): - self.tmp_file.close() - if os.path.isfile(self.tmp_file_path): - os.remove(self.tmp_file_path) diff --git a/auto_subtitle/utils/sonarr.py b/auto_subtitle/utils/sonarr.py index e6282ae..4bbb394 100644 --- a/auto_subtitle/utils/sonarr.py +++ b/auto_subtitle/utils/sonarr.py @@ -1,10 +1,14 @@ import requests import json -import os -token = os.getenv('sonarr_token') +import configparser +config = configparser.RawConfigParser() +config.read('config.cfg') + +token = config._sections['sonarr']['token'] +base_url = config._sections['sonarr']['url'] def update_show_in_soarr(show_id): - url = "http://192.168.4.9:8989/api/v3/command" + url = f"{base_url}/api/v3/command" payload = json.dumps({ "name": "RefreshSeries", diff --git a/config.cfg.example b/config.cfg.example new file mode 100644 index 0000000..c51ce0f --- /dev/null +++ b/config.cfg.example @@ -0,0 +1,6 @@ +[bazarr] +url = http://1.1.1.1 +token = djfkjadncdfjkanvfjkvandfj +[sonarr] +url = http://2.2.2.2:8989 +token = dfifdmnajcdnjcvaldnjlk \ No newline at end of file