Update app.py

app.py

@@ -1,157 +1,89 @@
-# import streamlit as st
-# import whisper
-# import ffmpeg
-# import pandas as pd
-# import pickle
-# import os
-# from chromadb.utils import embedding_functions
-# from chromadb import Client
-# from chromadb.config import Settings
-
-# def extract_audio(video_file):
-#     """Extracts audio using ffmpeg."""
-#     audio_path = "temp_audio.wav"
-#     ffmpeg.input(video_file).output(audio_path).run(overwrite_output=True)
-#     return audio_path
-
-# def transcribe_audio(audio_path):
-#     """Transcribes audio to text using Whisper."""
-#     model = whisper.load_model("base")
-#     result = model.transcribe(audio_path)
-#     return result['text']
-
-# def load_embeddings():
-#     """Loads subtitle embeddings from pkl file."""
-#     with open('subtitle_embeddings.pkl', 'rb') as f:
-#         embeddings = pickle.load(f)
-#     return embeddings
-
-# def save_to_chroma(embeddings):
-#     """Stores embeddings in Chroma DB."""
-#     client = Client(Settings())
-#     collection = client.create_collection(name="subtitles")
-#     for idx, row in embeddings.iterrows():
-#         collection.add(
-#             documents=[row['subtitle']],
-#             ids=[str(idx)],
-#             embeddings=[row['embedding']]
-#         )
-#     return collection
-
-# def search_subtitles(query, collection):
-#     """Searches for subtitles in Chroma DB."""
-#     results = collection.query(query_texts=[query], n_results=5)
-#     return results['documents']
-
-# def main():
-#     st.set_page_config(page_title="Video Subtitle Generator", layout="wide")
-#     st.title("🎥 Video Subtitle Generator")
-
-#     with st.sidebar:
-#         uploaded_file = st.file_uploader("Upload Video", type=["mp4", "mkv"])
-#         query = st.text_input("Search Subtitles")
-#         download_btn = st.button("Download Subtitles")
-        
-#     if uploaded_file:
-#         with st.spinner("Extracting audio..."):
-#             audio_path = extract_audio(uploaded_file.name)
-        
-#         with st.spinner("Generating subtitles..."):
-#             subtitles = transcribe_audio(audio_path)
-#             st.success("Subtitles Generated!")
-
-#         # Display the video and subtitles
-#         st.video(uploaded_file)
-#         st.text_area("Generated Subtitles", subtitles, height=300)
-        
-#         # Load and search embeddings
-#         embeddings = load_embeddings()
-#         collection = save_to_chroma(embeddings)
-        
-#         if query:
-#             results = search_subtitles(query, collection)
-#             st.write("### Matching Subtitles:")
-#             for sub in results:
-#                 st.write(f"- {sub}")
-        
-#         if download_btn:
-#             with open("generated_subtitles.srt", "w") as f:
-#                 f.write(subtitles)
-#             st.download_button("Download SRT", "generated_subtitles.srt")
-
-# if __name__ == '__main__':
-#     main()
 import streamlit as st
 import whisper
 import ffmpeg
 import pandas as pd
 import pickle
 import os
-from chromadb.utils import embedding_functions
+import numpy as np
+from sentence_transformers import SentenceTransformer
 from chromadb import Client
 from chromadb.config import Settings
 
+embed_model = SentenceTransformer('all-MiniLM-L6-v2')
+
 def extract_audio(uploaded_file):
-    """Extracts audio from video or handles audio file directly."""
     audio_path = "temp_audio.wav"
-
-    # Save uploaded file temporarily
     temp_file = f"temp_{uploaded_file.name}"
     with open(temp_file, "wb") as f:
         f.write(uploaded_file.getvalue())
 
-    # Extract audio for video files, keep as-is for audio
-    if uploaded_file.name.endswith(('.mp4', '.mkv')):
-        ffmpeg.input(temp_file).output(audio_path).run(overwrite_output=True)
-    else:
-        audio_path = temp_file
-
-    return audio_path, temp_file
+    try:
+        if uploaded_file.name.endswith(('.mp4', '.mkv')):
+            ffmpeg.input(temp_file).output(audio_path).run(overwrite_output=True)
+        else:
+            audio_path = temp_file
+        return audio_path, temp_file
+    except Exception as e:
+        st.error(f"Error extracting audio: {str(e)}")
+        return None, None
 
 def transcribe_audio(audio_path):
-    """Transcribes audio to text using Whisper."""
     try:
         model = whisper.load_model("base")
-    except Exception:
-        st.warning("Downloading Whisper model. This may take a while...")
-        model = whisper.load_model("base")
+        result = model.transcribe(audio_path)
+
+        subtitles = []
+        for i, segment in enumerate(result['segments']):
+            start_time = format_timestamp(segment['start'])
+            end_time = format_timestamp(segment['end'])
+            text = segment['text']
+            subtitles.append(f"{i + 1}\n{start_time} --> {end_time}\n{text}\n")
+        
+        return subtitles
+    except Exception as e:
+        st.error(f"Error during transcription: {str(e)}")
+        return []
+
+def format_timestamp(seconds):
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int((seconds % 1) * 1000)
+    return f"{hours:02}:{minutes:02}:{secs:02},{millis:03}"
+
+def embed_subtitles(subtitles):
+    raw_texts = [line.split('\n')[2] for line in subtitles if line.strip()]
+    embeddings = embed_model.encode(raw_texts)
+
+    df = pd.DataFrame({
+        'subtitle': raw_texts,
+        'embedding': list(embeddings)
+    })
 
-    result = model.transcribe(audio_path)
-    return result['text']
+    with open('subtitle_embeddings.pkl', 'wb') as f:
+        pickle.dump(df, f)
 
-def load_embeddings():
-    """Loads subtitle embeddings from pkl file."""
-    if os.path.exists('subtitle_embeddings.pkl'):
-        with open('subtitle_embeddings.pkl', 'rb') as f:
-            embeddings = pickle.load(f)
-        return embeddings
-    else:
-        st.error("No embeddings file found.")
-        return pd.DataFrame()
+    return df
 
 def save_to_chroma(embeddings):
-    """Stores embeddings in Chroma DB."""
     client = Client(Settings())
-    
-    # Check if collection exists, else create
-    try:
-        collection = client.get_collection("subtitles")
-    except:
-        collection = client.create_collection(name="subtitles")
-    
+    collection = client.create_collection(name="subtitles")
+
     for idx, row in embeddings.iterrows():
         collection.add(
             documents=[row['subtitle']],
             ids=[str(idx)],
             embeddings=[row['embedding']]
         )
-    
     return collection
 
 def search_subtitles(query, collection):
-    """Searches for subtitles in Chroma DB."""
-    results = collection.query(query_texts=[query], n_results=5)
-    return results['documents']
+    try:
+        results = collection.query(query_texts=[query], n_results=5)
+        return results['documents']
+    except Exception as e:
+        st.error(f"Error searching subtitles: {str(e)}")
+        return []
 
 def main():
     st.set_page_config(page_title="Video/Audio Subtitle Generator", layout="wide")
@@ -166,43 +98,43 @@ def main():
         with st.spinner("Extracting audio..."):
             audio_path, temp_file = extract_audio(uploaded_file)
 
-        with st.spinner("Generating subtitles..."):
-            subtitles = transcribe_audio(audio_path)
-            st.success("Subtitles Generated!")
+        if audio_path:
+            with st.spinner("Generating subtitles..."):
+                subtitles = transcribe_audio(audio_path)
+                st.success("Subtitles Generated!")
 
-        # Display the media and subtitles
-        if uploaded_file.name.endswith(('.mp4', '.mkv')):
-            st.video(uploaded_file)
-        else:
-            st.audio(uploaded_file)
+            if uploaded_file.name.endswith(('.mp4', '.mkv')):
+                st.video(uploaded_file)
+            else:
+                st.audio(uploaded_file)
 
-        st.text_area("Generated Subtitles", subtitles, height=300)
+            st.write("### Generated Subtitles:")
+            for sub in subtitles:
+                st.text(sub)
 
-        # Load and search embeddings
-        embeddings = load_embeddings()
-        if not embeddings.empty:
-            collection = save_to_chroma(embeddings)
+            with st.spinner("Embedding and storing subtitles..."):
+                embeddings = embed_subtitles(subtitles)
+                
+                if embeddings.empty:
+                    st.warning("No subtitles generated.")
+                else:
+                    collection = save_to_chroma(embeddings)
 
             if query:
                 results = search_subtitles(query, collection)
                 st.write("### Matching Subtitles:")
-                for sub in results:
-                    st.write(f"- {sub}")
-
-        # Subtitle download option
-        if download_btn:
-            srt_content = f"1\n00:00:00,000 --> 00:00:10,000\n{subtitles}\n"
-            
-            st.download_button(
-                label="Download SRT",
-                data=srt_content.encode('utf-8'),
-                file_name="generated_subtitles.srt",
-                mime="text/plain"
-            )
-
-        # Cleanup temporary files
-        os.remove(audio_path)
-        os.remove(temp_file)
+                if results:
+                    for idx, sub in enumerate(results, start=1):
+                        st.write(f"{idx}. {sub}")
+                else:
+                    st.warning("No matching subtitles found.")
+
+            if download_btn:
+                with open("generated_subtitles.srt", "w") as f:
+                    f.writelines(subtitles)
+
+                with open("generated_subtitles.srt", "rb") as f:
+                    st.download_button("Download SRT", f, file_name="generated_subtitles.srt", mime="text/plain")
 
 if __name__ == '__main__':
     main()