Create app.py

app.py

@@ -0,0 +1,86 @@
+import streamlit as st
+import whisper
+import numpy as np
+import chromadb
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+import os
+
+# Set up the title and description
+st.title("🎥 AI-Powered Video Subtitle Extractor with Cosine Similarity & Chroma DB")
+st.write("Upload a video to extract speech, convert it to text, and find matching subtitles using cosine similarity.")
+
+# 🔥 Load Chroma DB
+chroma_path = "./chroma.sqlite3"  # Path to your local Chroma DB file
+chroma_client = chromadb.PersistentClient(path=chroma_path)
+collection = chroma_client.get_collection(name="subtitle_chunk1")
+
+# Load embedding model
+embedder = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Whisper model
+model = whisper.load_model("base")
+
+# Function to extract subtitles using cosine similarity
+def find_matching_subtitles(transcribed_text, top_k=5):
+    # Generate embedding for transcribed text
+    query_embedding = embedder.encode(transcribed_text).reshape(1, -1)
+
+    # Retrieve all stored subtitles from Chroma DB
+    results = collection.get()
+    all_embeddings = np.array(results['embeddings'])
+    all_documents = results['documents']
+    all_metadata = results['metadatas']
+
+    # Calculate cosine similarity
+    similarities = cosine_similarity(query_embedding, all_embeddings)[0]
+
+    # Get top K matches
+    top_indices = np.argsort(similarities)[-top_k:][::-1]
+
+    # Display matching subtitles
+    matching_subtitles = []
+    for idx in top_indices:
+        matching_subtitles.append({
+            "subtitle": all_documents[idx],
+            "similarity": similarities[idx],
+            "metadata": all_metadata[idx]
+        })
+
+    return matching_subtitles
+
+# Streamlit UI
+uploaded_file = st.file_uploader("Upload Video", type=["mp4", "mkv", "avi", "mov"])
+
+if uploaded_file:
+    # Save uploaded video temporarily
+    temp_video_path = os.path.join("temp_video", "uploaded_video.mp4")
+    os.makedirs("temp_video", exist_ok=True)
+    
+    with open(temp_video_path, "wb") as f:
+        f.write(uploaded_file.read())
+
+    # Transcribe video speech
+    st.info("⏳ Transcribing video speech...")
+    transcription = model.transcribe(temp_video_path)
+    transcribed_text = transcription['text']
+    
+    st.success("✅ Transcription complete!")
+    st.write("### Transcribed Speech:")
+    st.write(transcribed_text)
+
+    # Find matching subtitles
+    st.info("🔍 Finding matching subtitles...")
+    matching_subtitles = find_matching_subtitles(transcribed_text)
+
+    st.write("### 🎯 Matching Subtitles:")
+    for match in matching_subtitles:
+        st.write(f"**Subtitle:** {match['subtitle']}")
+        st.write(f"**Similarity:** {match['similarity']:.4f}")
+        st.write(f"**Metadata:** {match['metadata']}")
+        st.write("---")
+
+    # Clean up temporary video
+    os.remove(temp_video_path)
+
+st.sidebar.write("🔎 Upload a video to extract and match subtitles using Cosine Similarity & Chroma DB.")