Update app.py

app.py

@@ -1,160 +1,84 @@
 import streamlit as st
 import whisper
-import tempfile
-import os
-import chromadb
-from pydub import AudioSegment
-from sentence_transformers import SentenceTransformer
-import numpy as np
+import moviepy.editor as mp
+import pandas as pd
 import pickle
-
-# 💡 Set Page Configuration
-st.set_page_config(page_title="Audio Subtitle Tool", layout="wide")
-
-# 🌐 Initialize ChromaDB
-db_client = chromadb.PersistentClient(path="./chroma_data")
-sub_db = db_client.get_or_create_collection(name="audio_subtitles")
-
-# ⚙️ Load Whisper Model
-@st.cache_resource()
-def load_transcriber():
-    return whisper.load_model("base")
-
-transcriber = load_transcriber()
-
-# ⚙️ Load Sentence Transformer
-@st.cache_resource()
-def load_encoder():
-    return SentenceTransformer("paraphrase-MiniLM-L6-v2")
-
-semantic_encoder = load_encoder()
-
-# 🔥 Load Subtitle Data from Pickle
-@st.cache_resource()
-def load_subtitle_store():
-    with open("subtitle_data.pkl", "rb") as file:
-        subtitle_store = pickle.load(file)
-
-# ⏱️ Format Time into SRT format
-def time_formatter(sec):
-    hr = int(sec // 3600)
-    mins = int((sec % 3600) // 60)
-    secs = int(sec % 60)
-    millis = int((sec % 1) * 1000)
-    return f"{hr:02}:{mins:02}:{secs:02},{millis:03}"
-
-# 🎙️ Transcribe Audio and Store in ChromaDB
-@st.cache_resource()
-def process_audio(file_path):
-    result = transcriber.transcribe(file_path, beam_size=1)
-    segments = result['segments']
-    srt_subtitles = []
-    plain_text_subs = []
-
-    for idx, seg in enumerate(segments):
-        start = time_formatter(seg['start'])
-        end = time_formatter(seg['end'])
-        content = seg['text']
-
-        # Create SRT format
-        srt_subtitles.append(f"{idx + 1}\n{start} --> {end}\n{content}\n")
-        plain_text_subs.append(content)
-
-        # Store in ChromaDB
-        sub_db.upsert(
-            documents=[content],
-            metadatas=[{"start_time": start, "end_time": end}],
-            ids=[f"subtitle_{idx}"]
+import os
+from chromadb.utils import embedding_functions
+from chromadb import Client
+from chromadb.config import Settings
+
+def extract_audio(video_file):
+    """Extracts audio from uploaded video file."""
+    video = mp.VideoFileClip(video_file)
+    audio_path = "temp_audio.wav"
+    video.audio.write_audiofile(audio_path)
+    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 "\n".join(srt_subtitles), " ".join(plain_text_subs)
-
-# 🔍 Subtitle Search Function
-def search_content(query):
-    query_vec = semantic_encoder.encode(query).tolist()
-    matches = sub_db.query(query_embeddings=[query_vec], n_results=5)
-
-    if 'documents' in matches and matches['documents']:
-        results = []
-        for idx, doc in enumerate(matches['documents'][0]):
-            meta = matches['metadatas'][0][idx]
-            results.append(f"📌 {doc} (From: {meta['start_time']} → To: {meta['end_time']})")
-        return results
-    return []
-
-# 🔥 Sidebar Navigation
-st.sidebar.title("🎯 Navigation")
-page = st.sidebar.radio("Choose Action", ["Generate Subtitles", "Search Subtitles"])
-
-# 🎛️ Sidebar Styling
-st.sidebar.markdown(
-    """
-    <style>
-    .sidebar .sidebar-content {
-        background-color: #f0f0f5;
-        color: #333;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-# 🎬 Header Section
-st.markdown(
-    """
-    <h1 style='text-align: center; color: #673AB7;'>🎤 Audio Subtitle Generator & Search</h1>
-    <p style='text-align: center; font-size: 18px; color: #666;'>Generate subtitles from audio files and search through them.</p>
-    """,
-    unsafe_allow_html=True
-)
-
-# 🚀 Load Subtitle Data
-load_subtitle_store()
-
-# 🎵 File Upload Section
-st.markdown("### 🎧 Upload Your Audio File")
-audio_file = st.file_uploader("Supported formats: MP3, WAV", type=["mp3", "wav"])
-
-# 🎯 Page Logic
-if page == "Generate Subtitles":
-    if audio_file:
-        st.audio(audio_file, format='audio/mp3')
-
-        with st.spinner("🔧 Processing audio... Please wait."):
-            # Temporary file handling
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio:
-                if audio_file.name.endswith(".mp3"):
-                    audio = AudioSegment.from_mp3(audio_file)
-                elif audio_file.name.endswith(".wav"):
-                    audio = AudioSegment.from_wav(audio_file)
-                audio = audio.set_frame_rate(16000)
-                audio.export(tmp_audio.name, format="wav")
-
-                # Transcribe and store in ChromaDB
-                subtitle_output, plain_output = process_audio(tmp_audio.name)
-                st.success("✅ Transcription Completed!")
-
-                # Display Subtitles
-                st.markdown("### 📜 Generated Subtitles")
-                st.text_area("Subtitles (SRT Format)", subtitle_output, height=300)
-
-                # Download Options
-                st.download_button("⬇️ Download SRT File", subtitle_output, file_name="subtitles.srt", mime="text/plain")
-                st.download_button("⬇️ Download Plain Text", plain_output, file_name="subtitles.txt", mime="text/plain")
-
-                os.remove(tmp_audio.name)
-
-elif page == "Search Subtitles":
-    st.subheader("🔎 Search Subtitles")
-    query_input = st.text_input("Enter text to search")
-    
-    if query_input:
-        with st.spinner("🔍 Searching..."):
-            search_matches = search_content(query_input)
-            
-            if search_matches:
-                st.success("✅ Results Found:")
-                for match in search_matches:
-                    st.write(match)
-            else:
-                st.warning("⚠️ No matching results. Try a different query.")
+    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()