diff --git a/.vscode/launch.json b/.vscode/launch.json index a5131f3..6cacf66 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -14,6 +14,8 @@ "args": [ "--model", "base", + "--show", + "Gary Neville's Soccerbox" ], } ] diff --git a/bazarr-ai-sub-generator/cli.py b/bazarr-ai-sub-generator/cli.py index d1436b3..8780f82 100644 --- a/bazarr-ai-sub-generator/cli.py +++ b/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')", ) diff --git a/bazarr-ai-sub-generator/main.py b/bazarr-ai-sub-generator/main.py index d5202fa..2fe1ba3 100644 --- a/bazarr-ai-sub-generator/main.py +++ b/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) diff --git a/bazarr-ai-sub-generator/utils/bazarr.py b/bazarr-ai-sub-generator/utils/bazarr.py index dc110a7..722f53d 100644 --- a/bazarr-ai-sub-generator/utils/bazarr.py +++ b/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):