Update app.py

app.py

@@ -1,3 +1,86 @@
+# 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
@@ -8,10 +91,13 @@ from chromadb.utils import embedding_functions
 from chromadb import Client
 from chromadb.config import Settings
 
-def extract_audio(video_file):
-    """Extracts audio using ffmpeg."""
+def extract_audio(file_path):
+    """Extracts audio from video or handles audio file directly."""
     audio_path = "temp_audio.wav"
-    ffmpeg.input(video_file).output(audio_path).run(overwrite_output=True)
+    if file_path.endswith(('.mp4', '.mkv')):
+        ffmpeg.input(file_path).output(audio_path).run(overwrite_output=True)
+    else:
+        audio_path = file_path
     return audio_path
 
 def transcribe_audio(audio_path):
@@ -44,11 +130,11 @@ def search_subtitles(query, collection):
     return results['documents']
 
 def main():
-    st.set_page_config(page_title="Video Subtitle Generator", layout="wide")
-    st.title("🎥 Video Subtitle Generator")
+    st.set_page_config(page_title="Video/Audio Subtitle Generator", layout="wide")
+    st.title("🎥🎵 Video/Audio Subtitle Generator")
 
     with st.sidebar:
-        uploaded_file = st.file_uploader("Upload Video", type=["mp4", "mkv"])
+        uploaded_file = st.file_uploader("Upload Video/Audio", type=["mp4", "mkv", "mp3", "wav"])
         query = st.text_input("Search Subtitles")
         download_btn = st.button("Download Subtitles")
         
@@ -60,8 +146,12 @@ def main():
             subtitles = transcribe_audio(audio_path)
             st.success("Subtitles Generated!")
 
-        # Display the video and subtitles
-        st.video(uploaded_file)
+        # Display the media and subtitles
+        if uploaded_file.name.endswith(('.mp4', '.mkv')):
+            st.video(uploaded_file)
+        else:
+            st.audio(uploaded_file)
+        
         st.text_area("Generated Subtitles", subtitles, height=300)
         
         # Load and search embeddings