Code enhancement

app.py

@@ -3,6 +3,7 @@ from process_media import MediaProcessor
 import streamlit as st
 from summarizer import BARTSummarizer
 from  Utils import fetch_article_text, get_text_from_youtube_url
+from pytube import YouTube
 
 st.markdown(
 """
@@ -26,7 +27,89 @@ with st.sidebar:
     input_type = st.radio("", ["Text", "Media"], label_visibility = "hidden")
 
 
-text_to_summarize = None
+def generate_summary(overall_summary, auto_chapters_summarize, text_to_summarize, show_text = False):
+    if overall_summary:
+        with st.spinner("Generating overall summary..."):
+            summarizer = BARTSummarizer()
+            summary = summarizer.chunk_summarize(text_to_summarize)
+    elif auto_chapters_summarize:
+        with st.spinner("Generating auto chapters summary..."):
+            summarizer = BARTSummarizer()
+            summary = summarizer.auto_chapters_summarize(text_to_summarize)
+    if show_text:
+        st.markdown("#### Text before summarization:")
+        st.markdown('<div style="height: 500px; overflow: auto; margin-bottom: 20px;">' + text_to_summarize + '</div>', unsafe_allow_html=True)
+    st.markdown("#### Summary:")
+    st.write(summary)
+
+def show_buttons(type, data = None):
+    text_to_summarize = None
+    overall_summary = st.button("Overall summary")
+    auto_chapters_summary = st.button("Auto Chapters summary...")
+
+    if type == "raw_text" or type == "url" and data:
+        if overall_summary or auto_chapters_summary:
+            if type == "raw_text":
+                text_to_summarize = data
+                generate_summary(overall_summary, auto_chapters_summary, text_to_summarize)
+            elif type == "url":
+                text_to_summarize = fetch_article_text(data)
+                generate_summary(overall_summary, auto_chapters_summary, text_to_summarize, show_text = True)
+
+    elif type == "audio_file" and data:
+        if overall_summary or auto_chapters_summary:
+            with st.spinner("Fetching text from audio..."):
+                text_to_summarize = process_audio_file(data)
+            generate_summary(overall_summary, auto_chapters_summary, text_to_summarize, show_text = True)
+
+    elif type == "video_file" and data:
+        if overall_summary or auto_chapters_summary:
+            with st.spinner("Fetching text from video..."):
+                text_to_summarize = process_video_file(data)
+            generate_summary(overall_summary, auto_chapters_summary, text_to_summarize, show_text = True)
+
+    elif type == "youtube_url" and data:
+        if overall_summary or auto_chapters_summary:
+            try:
+                try:
+                    with st.spinner("Fetching text from video..."):
+                        text_to_summarize = get_text_from_youtube_url(data)
+                except:
+                    with st.spinner("Captions not available. Downloading video..."):
+                        text_to_summarize = get_yt_video(data)
+            except:
+                st.error("Unable to fetch text from video. Please try a different video.")
+                text_to_summarize = None
+            if text_to_summarize:
+                generate_summary(overall_summary, auto_chapters_summarize, text_to_summarize, show_text = True)
+    elif type == 'document':
+        # TODO: Add document summary
+        pass
+
+
+def process_video_file(video_file):
+    media_processor = MediaProcessor()
+    text = media_processor.process_video(video_file.read())
+    return text
+
+def process_audio_file(audio_file):
+    media_processor = MediaProcessor()
+    if audio_file.type == "audio/mpeg":
+        wav_bytes = media_processor.get_wav_from_audio(audio_file.read())
+    else:
+        wav_bytes = audio_file.read()
+    text = media_processor.process_audio(wav_bytes)
+    return text
+
+def get_yt_video(youtube_url):
+    yt = YouTube(youtube_url)
+    video_stream = yt.streams.first()
+    video_buffer = io.BytesIO()
+    video_stream.stream_to_buffer(video_buffer)
+    media_processor = MediaProcessor()
+    text = media_processor.process_video(video_buffer.getvalue())
+    return text
+    
 
 if input_type == "Text":
 
@@ -37,16 +120,12 @@ if input_type == "Text":
     if text_type == "Raw Text":
         text = st.text_area("Enter raw text here", height=240, max_chars=10000, placeholder="Enter a paragraph to summarize")
         if text:
-            text_to_summarize = text
+            show_buttons("raw_text", text)
         
     elif text_type == "URL":
         url = st.text_input("Enter URL here", placeholder="Enter URL to an article, blog post, etc.")
         if url:
-            article_text = fetch_article_text(url)
-            if article_text:
-                st.markdown("#### Text from url:")
-                st.write(article_text)
-                text_to_summarize = article_text
+            show_buttons("url", url)
     else:
         ## TODO: Add file upload option
         pass
@@ -60,53 +139,30 @@ elif input_type == "Media":
     if media_type == "Audio file":
         audio_file = st.file_uploader("Upload an audio file", type=["mp3", "wav"], label_visibility="visible")
         if audio_file is not None:
-            with st.spinner("Fetching text from audio..."):
-                # print(audio_file.read())
-                wav_bytes = None
-                media_processor = MediaProcessor()
-                if audio_file.type == "audio/mpeg":
-                    wav_bytes = media_processor.get_wav_from_audio(audio_file.read())
-                else:
-                    wav_bytes = audio_file.read()
-                text = media_processor.process_audio(wav_bytes)
-                st.markdown("#### Text from audio:")
-                st.write(text)
+            show_buttons("audio_file", audio_file)
+
     elif media_type == "Video file":
         video_file = st.file_uploader("Upload a video file", type=["mp4"], label_visibility="visible")
         if video_file is not None:
-            with st.spinner("Fetching text from video..."):
-                media_processor = MediaProcessor()
-                text = media_processor.process_video(video_file.read())
-                st.markdown("#### Text from video:")
-                st.write(text)
+            show_buttons("video_file", video_file)
+
     else:
         youtube_url = st.text_input("Enter YouTube URL here", placeholder="Enter URL to an YouTube video", label_visibility="visible")
         if youtube_url:
-            with st.spinner("Fetching text from video..."):
+            if not get_text(url):
                 try:
-                    text_to_summarize = get_text_from_youtube_url(youtube_url)
-                    st.markdown("#### Text from video:")
-                    st.markdown('<div style="height: 300px; overflow: auto; margin-bottom: 20px;">' + text_to_summarize + '</div>', unsafe_allow_html=True)
+                    try:
+                        with st.spinner("Fetching text from video..."):
+                            text_to_summarize = get_text_from_youtube_url(youtube_url)
+                    except:
+                        with st.spinner("Captions not available. Downloading video..."):
+                            get_yt_video(youtube_url)
+                    text_to_summarize = get_text(url)
                 except:
                     st.error("Unable to fetch text from video. Please try a different video.")
                     text_to_summarize = None
-
-if text_to_summarize is not None:
-    overall_summary = st.button("Overall summary")
-    auto_chapters_summary = st.button("Auto Chapters summary")
-    if overall_summary:
-        with st.spinner("Summarizing..."):
-            # time.sleep(2)
-            # st.write(text_to_summarize)
-            summarizer = BARTSummarizer()
-            summary = summarizer.chunk_summarize(text_to_summarize)
-            st.markdown("#### Summary:")
-            st.write(summary)
-    elif auto_chapters_summary:
-        with st.spinner("Summarizing..."):
-            # time.sleep(2)
-            # st.write(text_to_summarize)
-            summarizer = BARTSummarizer()
-            summary = summarizer.auto_chapters_summarize(text_to_summarize)
-            st.markdown("#### Summary:")
-            st.write(summary)
+            else:
+                text_to_summarize = get_text(url)
+            st.markdown("#### Text from video:")
+            
+