diff --git a/bazarr-ai-sub-generator/cli.py b/bazarr-ai-sub-generator/cli.py index e73cee8..2786fe9 100644 --- a/bazarr-ai-sub-generator/cli.py +++ b/bazarr-ai-sub-generator/cli.py @@ -2,8 +2,6 @@ import argparse from faster_whisper import available_models from utils.constants import LANGUAGE_CODES from main import process -from utils.convert import str2bool, str2timeinterval - def main(): """ @@ -12,15 +10,20 @@ def main(): Parses command line arguments, processes the inputs using the specified options, and performs transcription or translation based on the specified task. """ + # Create an ArgumentParser object with a specific formatter for default values parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) + + # Add argument for selecting the Whisper model parser.add_argument( "--model", default="small", choices=available_models(), help="name of the Whisper model to use", ) + + # Add argument for specifying the device to use (CPU, CUDA, or auto-detect) parser.add_argument( "--device", type=str, @@ -28,35 +31,24 @@ 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.", - # ) + + # Add argument for processing a single file parser.add_argument( "--file", type=str, default=None, help="Process a single file" ) + + # Add argument for processing all videos in a folder parser.add_argument( "--folder", type=str, default=None, help="Process all videos in folder" ) + + # Add argument for specifying the task: transcribe or translate parser.add_argument( "--show", type=str, @@ -64,6 +56,8 @@ def main(): help="whether to perform X->X speech recognition ('transcribe') \ or X->English translation ('translate')", ) + + # Add argument for setting the origin language of the video, with auto-detection as default parser.add_argument( "--language", type=str, @@ -72,16 +66,20 @@ def main(): help="What is the origin language of the video? \ If unset, it is detected automatically.", ) + + # Add argument for selecting the backend: whisper or faster_whisper parser.add_argument( "--backend", type=str, default="whisper", choices=["whisper", "faster_whisper"], ) + + # Parse the command line arguments into a dictionary args = parser.parse_args().__dict__ - + + # Call the process function with the parsed arguments process(args) - if __name__ == "__main__": main() diff --git a/bazarr-ai-sub-generator/main.py b/bazarr-ai-sub-generator/main.py index a876d21..8738faf 100644 --- a/bazarr-ai-sub-generator/main.py +++ b/bazarr-ai-sub-generator/main.py @@ -2,8 +2,9 @@ import os import warnings import tempfile import time +from typing import List, Dict, Any from utils.files import filename, write_srt -from utils.ffmpeg import get_audio, add_subtitles_to_mp4 +from utils.ffmpeg import get_audio, add_subtitles_to_mp4, check_for_subtitles from utils.bazarr import get_wanted_episodes, get_episode_details, sync_series from utils.sonarr import update_show_in_sonarr from utils.faster_whisper import WhisperAI as fasterWhisperAI @@ -11,57 +12,101 @@ from utils.whisper import WhisperAI from utils.decorator import measure_time +def process_audio_and_subtitles(file_path: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes audio extraction and subtitle generation for a given file. -def folder_flow(folder, model_args, args, backend): - print(f"Processing {folder}") - files = os.listdir(folder) - for file in files: - print(f"processing {file}") - path = folder+file - try: - audios = get_audio([path], 0, None) - subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend) + Args: + file_path (str): Path to the video file. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). - add_subtitles_to_mp4(subtitles) - time.sleep(5) - except Exception as ex: - print(f"skipping file due to - {ex}") + Returns: + None + """ + try: + audios = get_audio([file_path], 0, None) + subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend) + add_subtitles_to_mp4(subtitles) + time.sleep(5) + except Exception as ex: + print(f"Skipping file {file_path} due to - {ex}") -def file_flow(show, model_args, args, backend): - print(f"Processing {show}") - try: - audios = get_audio([show], 0, None) - subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend) - add_subtitles_to_mp4(subtitles) - time.sleep(5) - except Exception as ex: - print(f"skipping file due to - {ex}") +def folder_flow(folder: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes all files within a specified folder. -def bazzar_flow(show, model_args, args, backend): + Args: + folder (str): Path to the folder containing video files. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ + print(f"Processing folder {folder}") + files = os.listdir(folder) + for file in files: + path = os.path.join(folder, file) + print(f"Processing file {path}") + if not check_for_subtitles(path): + process_audio_and_subtitles(path, model_args, args, backend) + + +def file_flow(file_path: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes a single specified file. + + Args: + file_path (str): Path to the video file. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ + print(f"Processing file {file_path}") + if not check_for_subtitles(file_path): + process_audio_and_subtitles(file_path, model_args, args, backend) + + +def bazzar_flow(show: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes episodes needing subtitles from Bazarr API. + + Args: + show (str): The show name. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ list_of_episodes_needing_subtitles = get_wanted_episodes(show) - print( - f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles." - ) + 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"]) - try: - audios = get_audio([episode_data["path"]], 0, None) - subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend) - - add_subtitles_to_mp4(subtitles) - update_show_in_sonarr(episode["sonarrSeriesId"]) - time.sleep(5) - sync_series() - except Exception as ex: - print(f"skipping file due to - {ex}") + process_audio_and_subtitles(episode_data["path"], model_args, args, backend) + update_show_in_sonarr(episode["sonarrSeriesId"]) + sync_series() @measure_time -def get_subtitles( - audio_paths: list, output_dir: str, model_args: dict, transcribe_args: dict, backend: str -): +def get_subtitles(audio_paths: List[str], output_dir: str, model_args: Dict[str, Any], transcribe_args: Dict[str, Any], backend: str) -> Dict[str, str]: + """Generates subtitles for given audio files using the specified model. + + Args: + audio_paths (List[str]): List of paths to the audio files. + output_dir (str): Directory to save the generated subtitle files. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + transcribe_args (Dict[str, Any]): Transcription arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + Dict[str, str]: A dictionary mapping audio file paths to generated subtitle file paths. + """ if backend == 'whisper': model = WhisperAI(model_args, transcribe_args) else: @@ -82,8 +127,15 @@ def get_subtitles( return subtitles_path -def process(args: dict): +def process(args: Dict[str, Any]) -> None: + """Main entry point to determine which processing flow to use. + Args: + args (Dict[str, Any]): Dictionary of arguments including model, language, show, file, folder, and backend. + + Returns: + None + """ model_name: str = args.pop("model") language: str = args.pop("language") show: str = args.pop("show") @@ -92,16 +144,12 @@ def process(args: dict): backend: str = args.pop("backend") if model_name.endswith(".en"): - warnings.warn( - f"{model_name} is an English-only model, forcing English detection." - ) + warnings.warn(f"{model_name} is an English-only model, forcing English detection.") args["language"] = "en" - # if translate task used and language argument is set, then use it elif language != "auto": args["language"] = language - model_args = {} - model_args["device"] = args.pop("device") + model_args = {"device": args.pop("device")} if file: file_flow(file, model_args, args, backend) diff --git a/bazarr-ai-sub-generator/utils/ffmpeg.py b/bazarr-ai-sub-generator/utils/ffmpeg.py index 190b8fe..50664c3 100644 --- a/bazarr-ai-sub-generator/utils/ffmpeg.py +++ b/bazarr-ai-sub-generator/utils/ffmpeg.py @@ -4,6 +4,18 @@ import ffmpeg from .files import filename +def check_for_subtitles(video_path:str): + # Probe the video file to get information about its streams + probe = ffmpeg.probe(video_path) + + # Check if there are any subtitle streams + for stream in probe['streams']: + if stream['codec_type'] == 'subtitle': + print("File has subtitles") + return True + + return False + def get_audio(paths: list, audio_channel_index: int, sample_interval: list): temp_dir = tempfile.gettempdir()