diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py index 21cdc16..6ef7393 100644 --- a/auto_subtitle/cli.py +++ b/auto_subtitle/cli.py @@ -18,6 +18,8 @@ def main(): 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("--output_txt", type=str2bool, default=False, + help="whether to save the subtitles in a txt file") parser.add_argument("--srt_only", type=str2bool, default=False, help="only generate the .srt file and not create overlayed video") parser.add_argument("--verbose", type=str2bool, default=False, @@ -32,6 +34,7 @@ def main(): model_name: str = args.pop("model") output_dir: str = args.pop("output_dir") output_srt: bool = args.pop("output_srt") + output_txt: bool = args.pop("output_txt") srt_only: bool = args.pop("srt_only") language: str = args.pop("language") @@ -48,7 +51,7 @@ def main(): model = whisper.load_model(model_name) audios = get_audio(args.pop("video")) subtitles = get_subtitles( - audios, output_srt or srt_only, output_dir, lambda audio_path: model.transcribe(audio_path, **args) + audios, output_srt or srt_only, output_txt, output_dir, lambda audio_path: model.transcribe(audio_path, **args) ) if srt_only: @@ -88,7 +91,7 @@ def get_audio(paths): return audio_paths -def get_subtitles(audio_paths: list, output_srt: bool, output_dir: str, transcribe: callable): +def get_subtitles(audio_paths: list, output_srt: bool, output_txt: bool, output_dir: str, transcribe: callable): subtitles_path = {} for path, audio_path in audio_paths.items(): @@ -108,6 +111,13 @@ def get_subtitles(audio_paths: list, output_srt: bool, output_dir: str, transcri subtitles_path[path] = srt_path + if output_txt: + text_path = os.path.join(output_dir, f"{filename(path)}.txt") + with open(text_path, "w", encoding="utf-8") as text_file: + for segment in result["segments"]: + print(segment["text"], file=text_file) + print(f"Saving subtitles to text file: {text_path}") + return subtitles_path