From cfafce9ccf158df6ae08f1bdadb88922c780c91e Mon Sep 17 00:00:00 2001
From: Kevin Le <github@kevinle.com.au>
Date: Wed, 5 Oct 2022 23:55:41 +0800
Subject: [PATCH] add argument to enable output of .srt file

enable srt only output
case insensitive string to bool conversion for arg
---
 auto_subtitle/cli.py   | 20 +++++++++++++++-----
 auto_subtitle/utils.py |  4 +++-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py
index 5f43118..3edc26b 100644
--- a/auto_subtitle/cli.py
+++ b/auto_subtitle/cli.py
@@ -16,8 +16,12 @@ def main():
                         choices=whisper.available_models(), help="name of the Whisper model to use")
     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("--verbose", type=str2bool, default=False,
-                        help="Whether to print out the progress and debug messages")
+                        help="whether to print out the progress and debug messages")
 
     parser.add_argument("--task", type=str, default="transcribe", choices=[
                         "transcribe", "translate"], help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')")
@@ -25,6 +29,8 @@ def main():
     args = parser.parse_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")
     os.makedirs(output_dir, exist_ok=True)
 
     if model_name.endswith(".en"):
@@ -35,8 +41,12 @@ def main():
     model = whisper.load_model(model_name)
     audios = get_audio(args.pop("video"))
     subtitles = get_subtitles(
-        audios, lambda audio_path: model.transcribe(audio_path, **args)
+        audios, output_srt or srt_only, output_dir, lambda audio_path: model.transcribe(audio_path, **args)
     )
+
+    if srt_only:
+        return
+
     # bash command to download a youtube video with `youtube-dl` and save it as `video.mp4`:
     # youtube-dl -f 22 -o video.mp4 https://www.youtube.com/watch?v=QH2-TGUlwu4
 
@@ -74,12 +84,12 @@ def get_audio(paths):
     return audio_paths
 
 
-def get_subtitles(audio_paths: list, transcribe: callable):
-    temp_dir = tempfile.gettempdir()
+def get_subtitles(audio_paths: list, output_srt: bool, output_dir: str, transcribe: callable):
+    srt_path = output_dir if output_srt else tempfile.gettempdir()
     subtitles_path = {}
 
     for path, audio_path in audio_paths.items():
-        srt_path = os.path.join(temp_dir, f"{filename(path)}.srt")
+        srt_path = os.path.join(srt_path, f"{filename(path)}.srt")
 
         print(
             f"Generating subtitles for {filename(path)}... This might take a while."
diff --git a/auto_subtitle/utils.py b/auto_subtitle/utils.py
index c63c3d5..ee5515b 100644
--- a/auto_subtitle/utils.py
+++ b/auto_subtitle/utils.py
@@ -3,7 +3,9 @@ from typing import Iterator, TextIO
 
 
 def str2bool(string):
-    str2val = {"True": True, "False": False}
+    string = string.lower()
+    str2val = {"true": True, "false": False}
+
     if string in str2val:
         return str2val[string]
     else: