|
@@ -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):
|
|
|
|
- temp_dir = tempfile.gettempdir()
|
|
|
|
|
|
+def get_subtitles(audio_paths: list, output_srt: bool, output_dir: str, transcribe: callable):
|
|
|
|
+ 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."
|