import whisper from pytube import YouTube import requests, io from urllib.request import urlopen from PIL import Image import time 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") # Define a function that we can use to load lottie files from a link. @st.cache(allow_output_mutation=True) def load_lottieurl(url: str): r = requests.get(url) if r.status_code != 200: return None return r.json() col1, col2 = st.columns([1, 3]) with col1: lottie = load_lottieurl("https://assets9.lottiefiles.com/private_files/lf30_bntlaz7t.json") st_lottie(lottie, speed=1, height=200, width=200) with col2: st.write(""" ## Youtube Transcriber ##### This is an app that transcribes YouTube videos into text.""") #def load_model(size): #default_size = size #if size == default_size: #return None #else: #loaded_model = whisper.load_model(size) #return loaded_model @st.cache(allow_output_mutation=True) def populate_metadata(link): yt = YouTube(link) author = yt.author title = yt.title description = yt.description thumbnail = yt.thumbnail_url length = yt.length views = yt.views return author, title, description, thumbnail, length, views # Uncomment if you want to fetch the thumbnails as well. #def fetch_thumbnail(thumbnail): #tnail = urlopen(thumbnail) #raw_data = tnail.read() #image = Image.open(io.BytesIO(raw_data)) #st.image(image, use_column_width=True) def convert(seconds): return time.strftime("%H:%M:%S", time.gmtime(seconds)) loaded_model = whisper.load_model("base") current_size = "None" size = st.selectbox("Model Size", ["tiny", "base", "small", "medium", "large"], index=1) def change_model(current_size, size): if current_size != size: loaded_model = whisper.load_model(size) st.write(f"Model is {'multilingual' if loaded_model.is_multilingual else 'English-only'} " f"and has {sum(np.prod(p.shape) for p in loaded_model.parameters()):,} parameters.") return loaded_model else: return None @st.cache(allow_output_mutation=True) def inference(link): yt = YouTube(link) path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4") results = loaded_model.transcribe(path) vtt = getSubs(results["segments"], "vtt", 80) srt = getSubs(results["segments"], "srt", 80) return results["text"], vtt, srt 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(): change_model(current_size, size) link = st.text_input("YouTube Link") if st.button("Transcribe"): author, title, description, thumbnail, length, views = populate_metadata(link) results = inference(link) col3, col4 = st.columns(2) with col3: #fetch_thumbnail(thumbnail) st.video(link) st.markdown(f"**Channel**: {author}") st.markdown(f"**Title**: {title}") st.markdown(f"**Length**: {convert(length)}") st.markdown(f"**Views**: {views:,}") with col4: with st.expander("Video Description"): st.write(description) #st.markdown(f"**Video Description**: {description}") with st.expander("Video Transcript"): st.write(results[0]) # Write the results to a .txt file and download it. with open("transcript.txt", "w+") as f: f.writelines(results[0]) f.close() with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f: datatxt = f.read() with open("transcript.vtt", "w+") as f: f.writelines(results[1]) f.close() with open(os.path.join(os.getcwd(), "transcript.vtt"), "rb") as f: datavtt = f.read() with open("transcript.srt", "w+") as f: f.writelines(results[2]) f.close() with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f: datasrt = f.read() if st.download_button(label="Download Transcript (.txt) ", data=datatxt, file_name=f"{title}.txt"): st.success("Downloaded Successfully!") elif st.download_button(label="Download Transcript (.vtt)", data=datavtt, file_name=f"{title}.vtt"): st.success("Downloaded Successfully!") elif st.download_button(label="Download Transcript (.srt)", data=datasrt, file_name=f"{title}.srt"): st.success("Downloaded Successfully! ") else: st.success("You can download the transcript in .srt format and upload it to YouTube to create subtitles for your video.") st.info("Streamlit refreshes after the download button is clicked. The data is cached so you can download the transcript again without having to transcribe the video again.") if __name__ == "__main__": main()