Karl/bazarr-ai-sub-generator
1
0
Fork 0
mirror of https://github.com/karl0ss/bazarr-ai-sub-generator.git synced 2025-05-08 07:53:39 +01:00
Code Issues Packages Projects Releases Wiki Activity

add argument to enable output of .srt file

Browse Source 
enable srt only output
case insensitive string to bool conversion for arg
This commit is contained in:
Kevin Le 2022-10-05 23:55:41 +08:00 committed by Miguel Piedrafita
parent 9f692d4d60
commit cfafce9ccf
2 changed files with 18 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
auto_subtitle/cli.py
View File

@ -16,8 +16,12 @@ def main():
choices=whisper.available_models(), help="name of the Whisper model to use") choices=whisper.available_models(), help="name of the Whisper model to use")
parser.add_argument("--output_dir", "-o", type=str, parser.add_argument("--output_dir", "-o", type=str,
default=".", help="directory to save the outputs") 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, 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=[ 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')") "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__ args = parser.parse_args().__dict__
model_name: str = args.pop("model") model_name: str = args.pop("model")
output_dir: str = args.pop("output_dir") 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) os.makedirs(output_dir, exist_ok=True)
if model_name.endswith(".en"): if model_name.endswith(".en"):
@ -35,8 +41,12 @@ def main():
model = whisper.load_model(model_name) model = whisper.load_model(model_name)
audios = get_audio(args.pop("video")) audios = get_audio(args.pop("video"))
subtitles = get_subtitles( 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`: # 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 # 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 return audio_paths
def get_subtitles(audio_paths: list, transcribe: callable): def get_subtitles(audio_paths: list, output_srt: bool, output_dir: str, transcribe: callable):
temp_dir = tempfile.gettempdir() srt_path = output_dir if output_srt else tempfile.gettempdir()
subtitles_path = {} subtitles_path = {}
for path, audio_path in audio_paths.items(): 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( print(
f"Generating subtitles for {filename(path)}... This might take a while." f"Generating subtitles for {filename(path)}... This might take a while."

4
auto_subtitle/utils.py
View File

@ -3,7 +3,9 @@ from typing import Iterator, TextIO
def str2bool(string): def str2bool(string):
str2val = {"True": True, "False": False} string = string.lower()
str2val = {"true": True, "false": False}
if string in str2val: if string in str2val:
return str2val[string] return str2val[string]
else: else: