|
@@ -0,0 +1,101 @@
|
|
|
+import os
|
|
|
+import ffmpeg
|
|
|
+import whisper
|
|
|
+import argparse
|
|
|
+import warnings
|
|
|
+import tempfile
|
|
|
+from .utils import filename, str2bool, write_srt
|
|
|
+
|
|
|
+
|
|
|
+def main():
|
|
|
+ parser = argparse.ArgumentParser(
|
|
|
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
|
+ parser.add_argument("video", nargs="+", type=str,
|
|
|
+ help="paths to video files to transcribe")
|
|
|
+ parser.add_argument("--model", default="small",
|
|
|
+ choices=whisper.available_models(), help="name of the Whisper model to use")
|
|
|
+ parser.add_argument("--output_dir", "-o", type=str,
|
|
|
+ default=".", help="directory to save the outputs")
|
|
|
+ parser.add_argument("--verbose", type=str2bool, default=False,
|
|
|
+ help="Whether to print out the progress and debug messages")
|
|
|
+
|
|
|
+ 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')")
|
|
|
+
|
|
|
+ args = parser.parse_args().__dict__
|
|
|
+ model_name: str = args.pop("model")
|
|
|
+ output_dir: str = args.pop("output_dir")
|
|
|
+ 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"
|
|
|
+
|
|
|
+ model = whisper.load_model(model_name)
|
|
|
+ audios = get_audio(args.pop("video"))
|
|
|
+ subtitles = get_subtitles(
|
|
|
+ audios, lambda audio_path: model.transcribe(audio_path, **args)
|
|
|
+ )
|
|
|
+ # 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
|
|
|
+
|
|
|
+ for path, srt_path in subtitles.items():
|
|
|
+ out_path = os.path.join(output_dir, f"{filename(path)}.mp4")
|
|
|
+
|
|
|
+ print(f"Adding subtitles to {filename(path)}...")
|
|
|
+
|
|
|
+ video = ffmpeg.input(path)
|
|
|
+ audio = video.audio
|
|
|
+
|
|
|
+ stderr = ffmpeg.concat(
|
|
|
+ video.filter('subtitles', srt_path, force_style="OutlineColour=&H40000000,BorderStyle=3"), audio, v=1, a=1
|
|
|
+ ).output(out_path).run(quiet=True, overwrite_output=True)
|
|
|
+
|
|
|
+ print(f"Saved subtitled video to {os.path.abspath(out_path)}.")
|
|
|
+
|
|
|
+
|
|
|
+def get_audio(paths):
|
|
|
+ temp_dir = tempfile.gettempdir()
|
|
|
+
|
|
|
+ audio_paths = {}
|
|
|
+
|
|
|
+ for path in paths:
|
|
|
+ print(f"Extracting audio from {filename(path)}...")
|
|
|
+ output_path = os.path.join(temp_dir, f"{filename(path)}.wav")
|
|
|
+
|
|
|
+ ffmpeg.input(path).output(
|
|
|
+ output_path,
|
|
|
+ acodec="pcm_s16le", ac=1, ar="16k"
|
|
|
+ ).run(quiet=True, overwrite_output=True)
|
|
|
+
|
|
|
+ audio_paths[path] = output_path
|
|
|
+
|
|
|
+ return audio_paths
|
|
|
+
|
|
|
+
|
|
|
+def get_subtitles(audio_paths: list, transcribe: callable):
|
|
|
+ temp_dir = tempfile.gettempdir()
|
|
|
+ subtitles_path = {}
|
|
|
+
|
|
|
+ for path, audio_path in audio_paths.items():
|
|
|
+ srt_path = os.path.join(temp_dir, f"{filename(path)}.srt")
|
|
|
+
|
|
|
+ print(
|
|
|
+ f"Generating subtitles for {filename(path)}... This might take a while."
|
|
|
+ )
|
|
|
+
|
|
|
+ warnings.filterwarnings("ignore")
|
|
|
+ result = transcribe(audio_path)
|
|
|
+ warnings.filterwarnings("default")
|
|
|
+
|
|
|
+ with open(srt_path, "w", encoding="utf-8") as srt:
|
|
|
+ write_srt(result["segments"], file=srt)
|
|
|
+
|
|
|
+ subtitles_path[path] = srt_path
|
|
|
+
|
|
|
+ return subtitles_path
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == '__main__':
|
|
|
+ main()
|