cleanup and add ability to process specifc show only

.vscode/launch.json

@@ -14,6 +14,8 @@
             "args": [
                 "--model",
                 "base",
+                "--show",
+                "Gary Neville's Soccerbox"
             ],
         }
     ]

bazarr-ai-sub-generator/cli.py

@@ -15,16 +15,16 @@ def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
-    parser.add_argument(
-        "--audio_channel", default="0", type=int, help="audio channel index to use"
-    )
-    parser.add_argument(
-        "--sample_interval",
-        type=str2timeinterval,
-        default=None,
-        help="generate subtitles for a specific \
-                              fragment of the video (e.g. 01:02:05-01:03:45)",
-    )
+    # parser.add_argument(
+    #     "--audio_channel", default="0", type=int, help="audio channel index to use"
+    # )
+    # parser.add_argument(
+    #     "--sample_interval",
+    #     type=str2timeinterval,
+    #     default=None,
+    #     help="generate subtitles for a specific \
+    #                           fragment of the video (e.g. 01:02:05-01:03:45)",
+    # )
     parser.add_argument(
         "--model",
         default="small",
@@ -38,46 +38,46 @@ def main():
         choices=["cpu", "cuda", "auto"],
         help='Device to use for computation ("cpu", "cuda", "auto")',
     )
+    # parser.add_argument(
+    #     "--compute_type",
+    #     type=str,
+    #     default="default",
+    #     choices=[
+    #         "int8",
+    #         "int8_float32",
+    #         "int8_float16",
+    #         "int8_bfloat16",
+    #         "int16",
+    #         "float16",
+    #         "bfloat16",
+    #         "float32",
+    #     ],
+    #     help="Type to use for computation. \
+    #                           See https://opennmt.net/CTranslate2/quantization.html.",
+    # )
+    # 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,
+    #     help="model parameter, tweak to increase accuracy",
+    # )
+    # parser.add_argument(
+    #     "--condition_on_previous_text",
+    #     type=str2bool,
+    #     default=True,
+    #     help="model parameter, tweak to increase accuracy",
+    # )
     parser.add_argument(
-        "--compute_type",
+        "--show",
         type=str,
-        default="default",
-        choices=[
-            "int8",
-            "int8_float32",
-            "int8_float16",
-            "int8_bfloat16",
-            "int16",
-            "float16",
-            "bfloat16",
-            "float32",
-        ],
-        help="Type to use for computation. \
-                              See https://opennmt.net/CTranslate2/quantization.html.",
-    )
-    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,
-        help="model parameter, tweak to increase accuracy",
-    )
-    parser.add_argument(
-        "--condition_on_previous_text",
-        type=str2bool,
-        default=True,
-        help="model parameter, tweak to increase accuracy",
-    )
-    parser.add_argument(
-        "--task",
-        type=str,
-        default="transcribe",
-        choices=["transcribe", "translate"],
+        default=None,
+        #choices=["transcribe", "translate"],
         help="whether to perform X->X speech recognition ('transcribe') \
                               or X->English translation ('translate')",
     )

bazarr-ai-sub-generator/main.py

@@ -24,8 +24,9 @@ def process(args: dict):
     
     model_name: str = args.pop("model")
     language: str = args.pop("language")
-    sample_interval: str = args.pop("sample_interval")
-    audio_channel: str = args.pop("audio_channel")
+    show: str = args.pop("show")
+    # sample_interval: str = args.pop("sample_interval")
+    # audio_channel: str = args.pop("audio_channel")
 
     if model_name.endswith(".en"):
         warnings.warn(
@@ -37,18 +38,18 @@ def process(args: dict):
         args["language"] = language
 
     model_args = {}
-    model_args["model_size_or_path"] = model_name
+    # model_args["model_size_or_path"] = model_name
     model_args["device"] = args.pop("device")
-    model_args["compute_type"] = args.pop("compute_type")
+    # model_args["compute_type"] = args.pop("compute_type")
 
-    list_of_episodes_needing_subtitles = get_wanted_episodes()
+    list_of_episodes_needing_subtitles = get_wanted_episodes(show)
     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)
+        audios = get_audio([episode_data["path"]], 0, None)
         subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
 
         add_subtitles_to_mp4(subtitles)

bazarr-ai-sub-generator/utils/bazarr.py

@@ -8,15 +8,19 @@ token = config._sections["bazarr"]["token"]
 base_url = config._sections["bazarr"]["url"]
 
 
-def get_wanted_episodes():
+def get_wanted_episodes(show: str=None):
     url = f"{base_url}/api/episodes/wanted"
 
     payload = {}
     headers = {"accept": "application/json", "X-API-KEY": token}
 
     response = requests.request("GET", url, headers=headers, data=payload)
-
-    return response.json()
+    
+    data = response.json()
+    if show != None:
+        data['data'] = [item for item in data['data'] if item['seriesTitle'] == show]
+        data['total'] = len(data['data'])
+    return data
 
 
 def get_episode_details(episode_id: str):