Update app.py

app.py

@@ -1,150 +1,78 @@
 import streamlit as st
-import ffmpeg
-import numpy as np
 import pandas as pd
-from sentence_transformers import SentenceTransformer
+import chromadb
+from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-from transformers import pipeline
+import whisper
 import os
-import tempfile
-import shutil
-import chromadb
-
-# Initialize Chroma DB client
-client = chromadb.Client()
-
-# Sidebar for CSV Upload and Permanent Save
-st.sidebar.title("📂 Upload CSV File")
-csv_file = st.sidebar.file_uploader("Choose a CSV file", type=["csv"])
-
-# Save CSV permanently
-def save_csv_permanently(uploaded_file):
-    save_path = os.path.join(os.getcwd(), "permanent_subtitle_data.csv")
-    with open(save_path, "wb") as f:
-        f.write(uploaded_file.getbuffer())
-    return save_path
-
-# Load the CSV into Chroma DB
-def load_csv_to_chroma(csv_path):
-    df = pd.read_csv(csv_path)
 
-    # Ensure the embedding column is properly formatted
-    df['embedding'] = df['embedding'].apply(lambda x: np.array(eval(x)).tolist())
+# Initialize Chroma DB
+chroma_client = chromadb.Client()
+collection = chroma_client.create_collection(name="subtitles")
 
-    # Create Chroma collection
-    collection_name = "video_subtitles"
-    if collection_name in [col.name for col in client.list_collections()]:
-        client.delete_collection(name=collection_name)
+# Sidebar for subtitle upload
+st.sidebar.header("📂 Upload Subtitle CSV")
+subtitle_file = st.sidebar.file_uploader("Upload subtitle dataset (CSV)", type=["csv"])
 
-    collection = client.create_collection(name=collection_name)
+if subtitle_file:
+    # Read CSV and store in Chroma DB
+    subtitle_df = pd.read_csv(subtitle_file)
 
-    # Add data to Chroma DB
-    for i, row in df.iterrows():
+    for i, row in subtitle_df.iterrows():
         collection.add(
-            ids=[str(i)],
-            documents=[row['text']],
-            embeddings=[row['embedding']]
+            documents=[row['cleaned_subtitle']],
+            metadatas=[{"subtitle": row['subtitle']}],
+            ids=[f"subtitle_{i}"]
         )
+    st.sidebar.success("✅ Subtitles stored in Chroma DB permanently.")
 
-    return collection
-
-# Handle CSV upload and save permanently
-if csv_file:
-    st.sidebar.success("CSV uploaded successfully!")
-    
-    # Save CSV permanently
-    csv_path = save_csv_permanently(csv_file)
-    st.sidebar.success(f"CSV saved permanently at: {csv_path}")
-
-    # Load into Chroma DB
-    with st.spinner("Loading CSV into Chroma DB..."):
-        collection = load_csv_to_chroma(csv_path)
-    st.sidebar.success("CSV loaded into Chroma DB ✅")
-
-# Whisper model for transcription
-whisper_model = pipeline("automatic-speech-recognition", model="openai/whisper-base")
-
-# Function to extract audio from video
-def extract_audio(video_file, audio_file):
-    input_file = ffmpeg.input(video_file)
-    output_file = ffmpeg.output(input_file, audio_file, **{'vn': None, 'ar': 16000, 'ac': 1, 'f': 'wav'})
-    ffmpeg.run(output_file)
-
-# Function to transcribe audio
-def transcribe_audio(audio_file):
-    result = whisper_model(audio_file)
+# Function to transcribe video using Whisper
+def video_to_text(video_path):
+    model = whisper.load_model("small")
+    result = model.transcribe(video_path)
     return result['text']
 
-# Generate embeddings for the transcription
-def generate_embedding(text):
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-    return model.encode(text).tolist()
+# Function to match transcribed text with subtitles
+def match_with_chroma(transcription, collection):
+    # Retrieve all subtitle texts from Chroma DB
+    docs = collection.get()
+    subtitles = [doc['document'] for doc in docs['documents']]
 
-# Search subtitles in Chroma DB
-def search_in_chroma(transcribed_text, collection, top_k=10):
-    query_embedding = np.array(generate_embedding(transcribed_text))
+    # Vectorization with TF-IDF
+    vectorizer = TfidfVectorizer()
+    vectors = vectorizer.fit_transform([transcription] + subtitles)
 
-    # Query Chroma DB
-    results = collection.query(
-        query_embeddings=[query_embedding.tolist()],
-        n_results=top_k
-    )
+    # Cosine similarity
+    similarity_scores = cosine_similarity(vectors[0:1], vectors[1:])
 
-    # Prepare results with cosine similarity
-    subtitles = []
-    for i, doc in enumerate(results['documents'][0]):
-        embedding = np.array(results['embeddings'][0][i])
-        similarity = cosine_similarity([query_embedding], [embedding])[0][0]
+    # Best match
+    best_match_index = similarity_scores.argmax()
+    best_match = docs['metadatas'][best_match_index]['subtitle']
 
-        subtitles.append({
-            "text": doc,
-            "cosine_similarity": similarity
-        })
-
-    # Sort results by similarity
-    subtitles = sorted(subtitles, key=lambda x: x['cosine_similarity'], reverse=True)
-    return subtitles
+    return best_match
 
 # Streamlit UI
-st.title("🎥 Video Subtitle Search with Chroma DB")
+st.title("🎥 Video Subtitle Extractor")
 
 # Upload video
-uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"])
-
-if uploaded_file and csv_file:
-    # Create temporary directory
-    temp_dir = tempfile.mkdtemp()
-
-    # Save video temporarily
-    video_path = os.path.join(temp_dir, "temp_video.mp4")
-    with open(video_path, "wb") as f:
-        f.write(uploaded_file.getbuffer())
-
-    audio_path = os.path.join(temp_dir, "temp_audio.wav")
-
-    # Extract audio
-    st.info("Extracting audio...")
-    extract_audio(video_path, audio_path)
+video_file = st.file_uploader("Upload a video", type=["mp4", "mkv", "avi"])
 
-    # Transcribe audio
-    st.info("Transcribing audio...")
-    transcribed_text = transcribe_audio(audio_path)
-    st.text_area("Transcribed Text", transcribed_text, height=150)
+if video_file:
+    with st.spinner("Processing video..."):
+        # Save uploaded video temporarily
+        video_path = "uploaded_video.mp4"
+        with open(video_path, "wb") as f:
+            f.write(video_file.read())
 
-    # Search in Chroma DB
-    st.info("Searching subtitles in Chroma DB...")
-    matching_subtitles = search_in_chroma(transcribed_text, collection)
+        # Transcribe video
+        transcription = video_to_text(video_path)
 
-    # Display video
-    st.video(video_path)
+        # Find best-matching subtitle
+        best_match = match_with_chroma(transcription, collection)
 
-    # Display matching subtitles
-    st.subheader("📜 Matching Subtitles with Cosine Similarity")
-    for sub in matching_subtitles:
-        st.write(f"**Subtitle:** {sub['text']}")
-        st.write(f"**Cosine Similarity:** {sub['cosine_similarity']:.4f}")
-        st.write("---")
+        # Display results
+        st.subheader("🔍 Extracted Subtitle")
+        st.write(best_match)
 
-    # Cleanup temporary files and directory
-    shutil.rmtree(temp_dir)
-    st.success("Temporary files cleaned up successfully ✅")
+        # Clean up temporary video file
+        os.remove(video_path)