kernel982/Youtube-Transcriber

1import whisper2from pytube import YouTube3import requests, io4from urllib.request import urlopen5from PIL import Image6import time7import streamlit as st8from streamlit_lottie import st_lottie9import numpy as np10import os11from typing import Iterator12from io import StringIO13from utils import write_vtt, write_srt14 15st.set_page_config(page_title="YouTube Transcriber", page_icon="🗣", layout="wide")16 17# Define a function that we can use to load lottie files from a link.18@st.cache(allow_output_mutation=True)19def load_lottieurl(url: str):20    r = requests.get(url)21    if r.status_code != 200:22        return None23    return r.json()24 25col1, col2 = st.columns([1, 3])26with col1:27    lottie = load_lottieurl("https://assets9.lottiefiles.com/private_files/lf30_bntlaz7t.json")28    st_lottie(lottie, speed=1, height=200, width=200)29 30with col2:31    st.write("""32    ## Youtube Transcriber 33    ##### This is an app that transcribes YouTube videos into text.""")34 35 36#def load_model(size):37    #default_size = size38    #if size == default_size:39        #return None40    #else:41        #loaded_model = whisper.load_model(size)42        #return loaded_model43 44@st.cache(allow_output_mutation=True)45def inference(link):46    yt = YouTube(link)47    print(yt.title)48    author = yt.author49    title = yt.title50    description = yt.description51    thumbnail = yt.thumbnail_url52    length = yt.length53    views = yt.views54    path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4")55    results = loaded_model.transcribe(path)56    vtt = getSubs(results["segments"], "vtt", 80)57    srt = getSubs(results["segments"], "srt", 80)58    return author, title, description, thumbnail, length, views, results["text"], vtt, srt59 60 61# Uncomment if you want to fetch the thumbnails as well.62# def fetch_thumbnail(thumbnail):63#     tnail = urlopen(thumbnail)64#     raw_data = tnail.read()65#     image = Image.open(io.BytesIO(raw_data))66#     st.image(image, use_column_width=True)67 68 69def convert(seconds):70    return time.strftime("%H:%M:%S", time.gmtime(seconds))71 72 73loaded_model = whisper.load_model("base")74current_size = "None"75size = st.selectbox("Model Size", ["tiny.en", "tiny", "base", "small", "medium", "large"], index=1)76 77 78def change_model(current_size, size):79    if current_size != size:80        loaded_model = whisper.load_model(size)81        st.write(f"Model is {'multilingual' if loaded_model.is_multilingual else 'English-only'} "82        f"and has {sum(np.prod(p.shape) for p in loaded_model.parameters()):,} parameters.")83        return loaded_model84    else:85        return None86 87 88def getSubs(segments: Iterator[dict], format: str, maxLineWidth: int) -> str:89    segmentStream = StringIO()90 91    if format == 'vtt':92        write_vtt(segments, file=segmentStream, maxLineWidth=maxLineWidth)93    elif format == 'srt':94        write_srt(segments, file=segmentStream, maxLineWidth=maxLineWidth)95    else:96        raise Exception("Unknown format " + format)97 98    segmentStream.seek(0)99    return segmentStream.read()100 101 102def main():103    change_model(current_size, size)104    link = st.text_input("YouTube Link")105    if st.button("Transcribe"):106        author, title, description, thumbnail, length, views, text, vtt, srt = inference(link)107        results = (text, vtt, srt)108 109        col3, col4 = st.columns(2)110        with col3:111            #fetch_thumbnail(thumbnail)112            st.video(link)113            st.markdown(f"**Channel**: {author}")114            st.markdown(f"**Title**: {title}")115            st.markdown(f"**Length**: {convert(length)}")116            st.markdown(f"**Views**: {views:,}")117 118        with col4:119            with st.expander("Video Description"):120                st.write(description)121            #st.markdown(f"**Video Description**: {description}")122            with st.expander("Video Transcript"):123                st.write(results[0])124            # Write the results to a .txt file and download it.125            with open("transcript.txt", "w+") as f:126                f.writelines(results[0])127                f.close()128            with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:129                datatxt = f.read()130 131            with open("transcript.vtt", "w+") as f:132                f.writelines(results[1])133                f.close()134            with open(os.path.join(os.getcwd(), "transcript.vtt"), "rb") as f:135                datavtt = f.read()136 137            with open("transcript.srt", "w+") as f:138                f.writelines(results[2])139                f.close()140            with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f:141                datasrt = f.read()142 143            if st.download_button(label="Download Transcript (.txt)",144                                data=datatxt,145                                file_name=f"{title}.txt"):146                st.success("Downloaded Successfully!")147 148            elif st.download_button(label="Download Transcript (.vtt)",149                                data=datavtt,150                                file_name=f"{title}.vtt"):151                st.success("Downloaded Successfully!")152 153            elif st.download_button(label="Download Transcript (.srt)",154                                data=datasrt,155                                file_name=f"{title}.srt"):156                st.success("Downloaded Successfully!")157            else:158                st.success("You can download the transcript in .srt format and upload it to YouTube to create subtitles for your video.")159                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.")160 161if __name__ == "__main__":162    main()