Update app.py

app.py

@@ -8,8 +8,11 @@ import streamlit as st
 from streamlit_lottie import st_lottie
 import numpy as np
 import os
+from typing import Iterator
+from io import StringIO
+from utils import write_vtt, write_srt
 
-st.set_page_config(page_title="Youtube Transcriber", page_icon="🗣", layout="wide")
+st.set_page_config(page_title="YouTube Transcriber", page_icon="🗣", layout="wide")
 
 
 # Define a function that we can use to load lottie files from a link.
@@ -83,7 +86,22 @@ def inference(link):
  yt = YouTube(link)
  path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4")
  results = loaded_model.transcribe(path)
- return results["text"]
+ vtt = getSubs(results["segments"], "vtt", 80)
+ return results["text"], vtt
+
+def getSubs(segments: Iterator[dict], format: str, maxLineWidth: int) -> str:
+ segmentStream = StringIO()
+
+ if format == 'vtt':
+ write_vtt(segments, file=segmentStream, maxLineWidth=maxLineWidth)
+ elif format == 'srt':
+ write_srt(segments, file=segmentStream, maxLineWidth=maxLineWidth)
+ else:
+ raise Exception("Unknown format " + format)
+
+ segmentStream.seek(0)
+ return segmentStream.read()
+
 
 
 def main():
@@ -107,10 +125,10 @@ def main():
  st.write(description)
  #st.markdown(f"**Video Description**: {description}")
  with st.expander("Video Transcript"):
- st.write(results)
+ st.write(results[0])
  # Write the results to a .txt file and download it.
  with open("transcript.txt", "w+") as f:
- f.writelines(results)
+ f.writelines(results[1])
  f.close()
  with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:
  data = f.read()