Alter logic to serve video on static endpoint

Dockerfile

@@ -29,7 +29,7 @@ COPY --chown=user youtubeaudio.py $HOME/root
 ARG model=base
 ENV model=$model
 RUN mkdir $HOME/root/models
-RUN mkdir -p -m 777 /tmp/holosubs/results
+RUN mkdir -p -m 775 /tmp/holosubs/tracks /tmp/holosubs/videos
 COPY --from=build "/whisper/models/ggml-${model}.bin" "${HOME}/root/models/ggml-${model}.bin"
 COPY --from=build /whisper/main /usr/local/bin/whisper
 

main.py

@@ -1,12 +1,13 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
-from youtubeaudio import YoutubeAudio
+from youtubeaudio import YoutubeMedia
 import subprocess
 import os
 import uuid
 import logging
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
+import ffmpeg
 
 class Url(BaseModel):
     url: str
@@ -15,9 +16,11 @@ format = "%(asctime)s: %(message)s"
 logging.basicConfig(format=format, level=logging.DEBUG,
                     datefmt="%H:%M:%S")
 MODEL=os.environ['model']
+BASE_DIR="/tmp/holosubs"
+VIDEO_DIR=os.path.join(BASE_DIR,"videos")
 
 app = FastAPI()
-app.mount("/tracks", StaticFiles(directory="/tmp/holosubs/results"), name="tracks")
+app.mount("/videos", StaticFiles(directory=VIDEO_DIR), name="videos")
 
 # CORS
 origins = ["*"]
@@ -30,27 +33,32 @@ app.add_middleware(
 )
 
 @app.post("/captions/")
-def read_root(url: Url):
+def read_root(url: Url, request: Request):
+    requestID=str(uuid.uuid4())
     # Download wav file and get filename
-    ytaudio=YoutubeAudio(url)
-    ytaudio.download_audio()
-    filename=ytaudio.filename
+    ytaudio=YoutubeMedia(url, os.path.join(BASE_DIR,requestID))
+    ytaudio.download()
+    audio_filename, video_filename = ytaudio.audio_filename, ytaudio.video_filename
     # Resample file
     ytaudio.resample('16k')
+    
     # Generate subtitles
-    captionFilename=str(uuid.uuid4())
-    output_file=os.path.join("/tmp/holosubs/results", captionFilename)
-    logging.info(f'Writing to file {output_file}.vtt')
+    subtitle_filename=os.path.join("/tmp/holosubs/tracks", requestID)
+    logging.info(f'Output will be writen to {subtitle_filename}.srt')
     cmd=['/usr/local/bin/whisper','-m',f'/home/user/root/models/ggml-{MODEL}.bin'
-         ,'-f',filename, '-di', '-of', output_file, '-tr', '-ovtt', '-t', '8']
+         ,'-f',audio_filename, '-of', subtitle_filename, '-tr', '-osrt', '-t', '8']
     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
     (output, err) = p.communicate()
     p_status = p.wait()
-    logging.info(output)
     if err:
-        logging.error("Whisper translation failed with error",err)
+        logging.error("Translation failed with error",err)
         raise HTTPException(status_code=500, detail="Whisper translation failed")
-    with open(output_file+".vtt", 'r') as f:
-        raw_vtt=f.read()
-    return {"captions": f"{captionFilename}.vtt"}
+    media=ffmpeg.input(video_filename)
+    ffmpeg.concat(media.video. \
+        filter('subtitles', subtitle_filename+".srt"), \
+        media.audio, v=1, a=1). \
+        output(os.path.join(VIDEO_DIR,requestID+".mp4")).run()
+    
+    # Return video url delivered
+    return {"video": f'/videos/{requestID+".mp4"}'}
 

youtubeaudio.py

@@ -14,45 +14,49 @@ format = "%(asctime)s: %(message)s"
 logging.basicConfig(format=format, level=logging.DEBUG,
                     datefmt="%H:%M:%S")
 
-class YoutubeAudio:
-    def __init__(self, url, dir="/tmp/holosubs/audio"):
+class YoutubeMedia:
+    def __init__(self, url, dir, audio_format="wav", video_format="mp4"):
         self.url=url
         self.dir=dir
+        self.audio_format=audio_format
+        self.video_format=video_format
 
-    def download_audio(self):
+    def download(self):
         ydl_opts = {
             'outtmpl': os.path.join(self.dir, "%(id)s_%(epoch)s.%(ext)s"),
             'logger': logging,
             'progress_hooks': [self.progress_hook],
-            'format': 'm4a/bestaudio/best',
+            'format': f'{self.video_format}/bestvideo',
+            'keepvideo': True,
             'postprocessors': [{  # Extract audio using ffmpeg
                 'key': 'FFmpegExtractAudio',
-                'preferredcodec': 'wav'
+                'preferredcodec': self.audio_format
             }]
         }
         with YoutubeDL(ydl_opts) as ydl: 
             error_code = ydl.download([self.url.url])
     
-    def resample(self,sr='16k'):
-        tmp_filename=os.path.join(self.dir,str(uuid.uuid4()))+".wav"
-        ffmpeg.input(self.filename).output(tmp_filename,ar=sr).run()
-        shutil.move(tmp_filename, self.filename)
-        logging.info(f"Succesfuly resampled {self.filename}")
+    def resample(self,sr='16k'): #defaults to 16k sampling rate
+        tmp_filename=os.path.join(self.dir,str(uuid.uuid4()))+"."+self.audio_format
+        ffmpeg.input(self.audio_filename).output(tmp_filename,ar=sr).run()
+        shutil.move(tmp_filename, self.audio_filename)
+        logging.info(f"Succesfuly resampled {self.audio_filename}")
 
     def clean(self):
-        if not self.filename:
+        if not self.audio_filename:
             logging.error("Audio not downloaded")
             return
-        location=os.path.join(self.dir, self.filename)
-        if os.path.exists(self.filename):
-            os.remove(self.filename)
-            logging.info(f"File {self.filename} successfully removed")
-            self.filename=None
+        location=os.path.join(self.dir, self.audio_filename)
+        if os.path.exists(self.audio_filename):
+            os.remove(self.audio_filename)
+            logging.info(f"File {self.audio_filename} successfully removed")
+            self.audio_filename=None
         else:
-            print(f"File {self.filename} does not exist")
+            print(f"File {self.audio_filename} does not exist")
 
     def progress_hook(self, d):
         if d['status'] == 'finished':
-            self.filename=os.path.join(self.dir, Path(d.get('info_dict').get('_filename')).stem + ".wav") 
-            print(f'Done downloading {self.filename}, now post-processing ...')
+            self.audio_filename=os.path.join(self.dir, Path(d.get('info_dict').get('_filename')).stem + "."+self.audio_format) 
+            self.video_filename=os.path.join(self.dir, Path(d.get('info_dict').get('_filename')).stem + "."+self.video_format) 
+            print(f'Done downloading {self.audio_filename}, now post-processing ...')