adding translation

.gitignore

@@ -0,0 +1,2 @@
+local
+.idea

app.py

@@ -1,23 +1,129 @@
+import os.path
+import pathlib
+from io import StringIO
 import gradio as gr
 import whisper
 import numpy as np
 import pytube as pt
+import ffmpeg
+import textwrap
+import cv2
+from transformers import pipeline
 
-model = whisper.load_model("medium")
+APP_DIR = pathlib.Path(__file__).parent.absolute()
+
+LOCAL_DIR = APP_DIR / "local"
+LOCAL_DIR.mkdir(exist_ok=True)
+save_dir = LOCAL_DIR / "output"
+save_dir.mkdir(exist_ok=True)
+
+transcriber = whisper.load_model("base")
 print(
-    f"Model is {'multilingual' if model.is_multilingual else 'English-only'} "
-    f"and has {sum(np.prod(p.shape) for p in model.parameters()):,} parameters."
+    f"Model is {'multilingual' if transcriber.is_multilingual else 'English-only'} "
+    f"and has {sum(np.prod(p.shape) for p in transcriber.parameters()):,} parameters."
 )
 
-options = dict(language='es', beam_size=5, best_of=5)
-transcribe_options = dict(task="transcribe", **options)
-translate_options = dict(task="translate", **options)
+options = dict(language='en', beam_size=5, best_of=5)
+transcribe_options = dict(task="translate", **options)
+
+translator = pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es")
+
+
+def translate(text):
+    return translator(text)[0]["translation_text"]
+
+
+def format_timestamp(seconds, always_include_hours=False, fractional_seperator='.'):
+    assert seconds >= 0, "non-negative timestamp expected"
+    milliseconds = round(seconds * 1000.0)
+
+    hours = milliseconds // 3_600_000
+    milliseconds -= hours * 3_600_000
+
+    minutes = milliseconds // 60_000
+    milliseconds -= minutes * 60_000
+
+    seconds = milliseconds // 1_000
+    milliseconds -= seconds * 1_000
+
+    hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
+    return f"{hours_marker}{minutes:02d}:{seconds:02d}{fractional_seperator}{milliseconds:03d}"
+
+
+def process_text(text: str, max_line_width=None):
+    if max_line_width is None or max_line_width < 0:
+        return text
+
+    lines = textwrap.wrap(text, width=max_line_width, tabsize=4)
+    return '\n'.join(lines)
+
+
+def write_srt(transcript, file, max_line_width=None):
+    for i, segment in enumerate(transcript, start=1):
+        text = process_text(segment['text'].strip(), max_line_width).replace('-->', '->')
+
+        # write srt lines
+        print(
+            f"{i}\n"
+            f"{format_timestamp(segment['start'], always_include_hours=True, fractional_seperator=',')} --> "
+            f"{format_timestamp(segment['end'], always_include_hours=True, fractional_seperator=',')}\n"
+            f"{text}\n",
+            file=file,
+            flush=True,
+        )
+
+
+def get_subs(segments, max_line_width):
+    for sentence in segments:
+        sentence['text'] = translate(sentence["text"])
+    segment_stream = StringIO()
+    write_srt(segments, file=segment_stream, max_line_width=max_line_width)
+    segment_stream.seek(0)
+    return segment_stream.read()
+
+
+def generate_subtitled_video(video, audio, transcript):
+    video_file = ffmpeg.input(video)
+    audio_file = ffmpeg.input(audio)
+    ffmpeg.concat(video_file.filter(
+        "subtitles", transcript
+    ), audio_file, v=1, a=1).output(f"{save_dir}/final.mp4").run(quiet=True, overwrite_output=True)
+    return f"{save_dir}/final.mp4"
+
+
+def generate_subtitled_audio(audio, transcript):
+    if not os.path.exists(f'{save_dir}/cover.jpg'):
+        cover = np.zeros([320, 640, 3], dtype=np.uint8)
+        cv2.imwrite(f'{save_dir}/cover.jpg', cover)
+    os.system(f'ffmpeg -y -loop 1 -i {save_dir}/cover.jpg '
+              f'-i {audio} -c:v libx264 '
+              f'-tune stillimage -c:a aac -b:a 192k -pix_fmt yuv420p -shortest '
+              f'-vf "subtitles={transcript}" {save_dir}/final.mp4')
+    return f"{save_dir}/final.mp4"
 
 
 def transcribe(audio):
-    transcription = model.transcribe(audio, **transcribe_options)
-    translation = model.transcribe(audio, **translate_options)
-    return transcription["text"], translation["text"]
+    transcription = transcriber.transcribe(audio, **transcribe_options)
+    srt = get_subs(transcription["segments"], 80)
+    with open(f"{save_dir}/transcript.srt", "w+", encoding='utf8') as f:
+        f.writelines(srt)
+        f.close()
+    return transcription["text"], srt
+
+
+def transcribe_audio(audio):
+    transcription, translation = transcribe(audio)
+    return generate_subtitled_audio(audio, f"{save_dir}/transcript.srt"), transcription, translation
+
+
+def transcribe_video(video):
+    audio = ffmpeg.input(video)
+    audio = ffmpeg.output(audio, f"{save_dir}/output.wav", acodec="pcm_s16le", ac=1, ar="16k")
+    ffmpeg.run(audio, overwrite_output=True)
+    audio = whisper.load_audio(f"{save_dir}/output.wav")
+    transcription, translation = transcribe(audio)
+    return generate_subtitled_video(video, f"{save_dir}/output.wav",
+                                    f"{save_dir}/transcript.srt"), transcription, translation
 
 
 def youtube_transcribe(url):
@@ -36,16 +142,16 @@ mic_interface = gr.Interface(
 )
 
 audio_interface = gr.Interface(
-    fn=transcribe,
+    fn=transcribe_audio,
     inputs=gr.Audio(type="filepath"),
-    outputs=["text", "text"],
+    outputs=["video", "text", "text"],
     title="Transcribir y traducir audio",
 )
 
 video_interface = gr.Interface(
-    fn=transcribe,
+    fn=transcribe_video,
     inputs="video",
-    outputs=["text", "text"],
+    outputs=["video", "text", "text"],
     title="Transcribir y traducir audio",
 )
 

requirements.txt

@@ -1,16 +1,23 @@
 aiohttp==3.8.3
 aiosignal==1.3.1
+altair==4.2.0
 anyio==3.6.2
 async-timeout==4.0.2
 attrs==22.1.0
+backports.zoneinfo==0.2.1
 bcrypt==4.0.1
+blinker==1.5
+cachetools==5.2.0
 certifi==2022.9.24
 cffi==1.15.1
 charset-normalizer==2.1.1
 click==8.1.3
+commonmark==0.9.1
 contourpy==1.0.6
 cryptography==38.0.4
 cycler==0.11.0
+decorator==5.1.1
+entrypoints==0.4
 fastapi==0.88.0
 ffmpeg-python==0.2.0
 ffmpy==0.3.0
@@ -19,13 +26,19 @@ fonttools==4.38.0
 frozenlist==1.3.3
 fsspec==2022.11.0
 future==0.18.2
+gitdb==4.0.10
+GitPython==3.1.29
 gradio==3.12.0
 h11==0.12.0
 httpcore==0.15.0
 httpx==0.23.1
 huggingface-hub==0.11.1
 idna==3.4
+importlib-metadata==5.1.0
+importlib-resources==5.10.1
 Jinja2==3.1.2
+joblib==1.2.0
+jsonschema==4.17.3
 kiwisolver==1.4.4
 linkify-it-py==1.0.3
 markdown-it-py==2.1.0
@@ -36,39 +49,60 @@ mdurl==0.1.2
 more-itertools==9.0.0
 multidict==6.0.3
 numpy==1.23.5
+opencv-python==4.6.0.66
 orjson==3.8.3
 packaging==21.3
 pandas==1.5.2
 paramiko==2.12.0
 Pillow==9.3.0
+pkgutil_resolve_name==1.3.10
+protobuf==3.20.3
+pyarrow==10.0.1
 pycparser==2.21
 pycryptodome==3.16.0
 pydantic==1.10.2
+pydeck==0.8.0
 pydub==0.25.1
+Pygments==2.13.0
+Pympler==1.0.1
 PyNaCl==1.5.0
 pyparsing==3.0.9
+pyrsistent==0.19.2
 python-dateutil==2.8.2
 python-multipart==0.0.5
+pytube==12.1.0
 pytz==2022.6
+pytz-deprecation-shim==0.1.0.post0
 PyYAML==6.0
 regex==2022.10.31
 requests==2.28.1
 rfc3986==1.5.0
+rich==12.6.0
+sacremoses==0.0.53
 semantic-version==2.10.0
+semver==2.13.0
+sentencepiece==0.1.97
 setuptools-rust==1.5.2
 six==1.16.0
+smmap==5.0.0
 sniffio==1.3.0
 starlette==0.22.0
+streamlit==1.15.2
 tokenizers==0.13.2
+toml==0.10.2
+toolz==0.12.0
 torch==1.13.0
+tornado==6.2
 tqdm==4.64.1
 transformers==4.25.1
 typing_extensions==4.4.0
+tzdata==2022.7
+tzlocal==4.2
 uc-micro-py==1.0.1
 urllib3==1.26.13
 uvicorn==0.20.0
+validators==0.20.0
 websockets==10.4
 whisper @ git+https://github.com/openai/whisper.git@fd8f80c8b880dd7c284c109ca7f03dbe978bc532
 yarl==1.8.2
-
-pytube~=12.1.0
+zipp==3.11.0