From 6462cd4d564158efafb7f9e9367278afc93ced0a Mon Sep 17 00:00:00 2001 From: Alistair Bahr Date: Fri, 10 Mar 2023 19:55:06 +0100 Subject: [PATCH] add custom language support --- auto_subtitle/cli.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/auto_subtitle/cli.py b/auto_subtitle/cli.py index a58d14f..21cdc16 100644 --- a/auto_subtitle/cli.py +++ b/auto_subtitle/cli.py @@ -25,19 +25,26 @@ def main(): 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')") + parser.add_argument("--language", type=str, default="auto", choices=["auto","af","am","ar","as","az","ba","be","bg","bn","bo","br","bs","ca","cs","cy","da","de","el","en","es","et","eu","fa","fi","fo","fr","gl","gu","ha","haw","he","hi","hr","ht","hu","hy","id","is","it","ja","jw","ka","kk","km","kn","ko","la","lb","ln","lo","lt","lv","mg","mi","mk","ml","mn","mr","ms","mt","my","ne","nl","nn","no","oc","pa","pl","ps","pt","ro","ru","sa","sd","si","sk","sl","sn","so","sq","sr","su","sv","sw","ta","te","tg","th","tk","tl","tr","tt","uk","ur","uz","vi","yi","yo","zh"], + help="What is the origin language of the video? If unset, it is detected automatically.") 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") + language: str = args.pop("language") + os.makedirs(output_dir, exist_ok=True) if model_name.endswith(".en"): 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 = whisper.load_model(model_name) audios = get_audio(args.pop("video")) subtitles = get_subtitles(