diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py index 8c5f0dd..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 efdcfeb..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, 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/ffmpeg.py b/auto_subtitle/utils/ffmpeg.py index 223cb02..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] 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)