wehub-resource-sync
296 lines
10 KiB
Python
296 lines
10 KiB
Python
import streamlit as st
|
|
import google.generativeai as genai
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import base64
|
|
from pathlib import Path
|
|
import mimetypes
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Page configuration
|
|
st.set_page_config(
|
|
page_title="Video RAG with Gemini",
|
|
page_icon="🎬",
|
|
layout="wide",
|
|
initial_sidebar_state="expanded"
|
|
)
|
|
|
|
# ===========================
|
|
# Video Processing Class
|
|
# ===========================
|
|
class VideoProcessor:
|
|
def __init__(self, api_key):
|
|
genai.configure(api_key=api_key)
|
|
self.model = genai.GenerativeModel('gemini-1.5-pro')
|
|
|
|
def upload_video(self, video_path, display_name=None):
|
|
"""Upload video to Gemini File API"""
|
|
try:
|
|
video_file = genai.upload_file(
|
|
path=video_path,
|
|
display_name=display_name or "uploaded_video"
|
|
)
|
|
return video_file
|
|
except Exception as e:
|
|
st.error(f"Error uploading video: {str(e)}")
|
|
return None
|
|
|
|
def wait_for_file_processing(self, video_file):
|
|
"""Wait for video file to be processed by Gemini"""
|
|
try:
|
|
while video_file.state.name == "PROCESSING":
|
|
time.sleep(2)
|
|
video_file = genai.get_file(video_file.name)
|
|
|
|
if video_file.state.name == "FAILED":
|
|
raise ValueError("Video processing failed")
|
|
|
|
return video_file
|
|
except Exception as e:
|
|
st.error(f"Error processing video: {str(e)}")
|
|
return None
|
|
|
|
def chat_with_video(self, video_file, prompt):
|
|
"""Generate response based on video content and user prompt"""
|
|
try:
|
|
response = self.model.generate_content([
|
|
video_file,
|
|
prompt
|
|
])
|
|
return response.text
|
|
except Exception as e:
|
|
st.error(f"Error generating response: {str(e)}")
|
|
return None
|
|
|
|
# ===========================
|
|
# Helper Functions
|
|
# ===========================
|
|
def is_video_file(file):
|
|
"""Check if uploaded file is a video"""
|
|
if file is None:
|
|
return False
|
|
|
|
mime_type, _ = mimetypes.guess_type(file.name)
|
|
return mime_type and mime_type.startswith('video/')
|
|
|
|
def get_file_size_mb(file):
|
|
"""Get file size in MB"""
|
|
return len(file.getvalue()) / (1024 * 1024)
|
|
|
|
def reset_chat():
|
|
"""Reset chat history and video state"""
|
|
st.session_state.messages = []
|
|
if 'video_file' in st.session_state:
|
|
try:
|
|
# Clean up uploaded file from Gemini
|
|
genai.delete_file(st.session_state.video_file.name)
|
|
except:
|
|
pass
|
|
del st.session_state.video_file
|
|
if 'video_processor' in st.session_state:
|
|
del st.session_state.video_processor
|
|
if 'video_name' in st.session_state:
|
|
del st.session_state.video_name
|
|
|
|
def display_video(video_bytes, video_name):
|
|
"""Display uploaded video"""
|
|
st.markdown(f"### 🎬 {video_name}")
|
|
st.video(video_bytes)
|
|
|
|
# ===========================
|
|
# Session State Initialization
|
|
# ===========================
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = []
|
|
|
|
if "video_file" not in st.session_state:
|
|
st.session_state.video_file = None
|
|
|
|
if "video_processor" not in st.session_state:
|
|
st.session_state.video_processor = None
|
|
|
|
if "video_name" not in st.session_state:
|
|
st.session_state.video_name = None
|
|
|
|
# ===========================
|
|
# Sidebar Configuration
|
|
# ===========================
|
|
with st.sidebar:
|
|
st.header("🔑 API Configuration")
|
|
|
|
# Try to get API key from environment first
|
|
default_api_key = os.getenv("GEMINI_API_KEY", "")
|
|
|
|
api_key = st.text_input(
|
|
"Gemini API Key",
|
|
value=default_api_key,
|
|
type="password",
|
|
help="Get your API key from https://aistudio.google.com/app/apikey"
|
|
)
|
|
|
|
if api_key:
|
|
if st.session_state.video_processor is None:
|
|
st.session_state.video_processor = VideoProcessor(api_key)
|
|
|
|
st.markdown("---")
|
|
|
|
st.header("📹 Upload Video")
|
|
uploaded_file = st.file_uploader(
|
|
"Choose a video file",
|
|
type=['mp4', 'avi', 'mov', 'mkv', 'webm'],
|
|
help="Supported formats: MP4, AVI, MOV, MKV, WEBM"
|
|
)
|
|
|
|
if uploaded_file is not None:
|
|
if not is_video_file(uploaded_file):
|
|
st.error("Please upload a valid video file.")
|
|
else:
|
|
file_size = get_file_size_mb(uploaded_file)
|
|
st.info(f"File size: {file_size:.2f} MB")
|
|
|
|
if file_size > 100: # Gemini has file size limits
|
|
st.warning("Large files may take longer to process or may fail. Consider compressing your video.")
|
|
|
|
# Process video if not already processed
|
|
if (st.session_state.video_file is None or
|
|
st.session_state.video_name != uploaded_file.name):
|
|
|
|
if st.session_state.video_processor is None:
|
|
st.error("Please enter your Gemini API key first.")
|
|
else:
|
|
with st.spinner("Uploading and processing video... This may take a few minutes."):
|
|
# Save uploaded file temporarily
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp_file:
|
|
tmp_file.write(uploaded_file.getvalue())
|
|
tmp_file_path = tmp_file.name
|
|
|
|
try:
|
|
# Upload to Gemini
|
|
video_file = st.session_state.video_processor.upload_video(
|
|
tmp_file_path,
|
|
uploaded_file.name
|
|
)
|
|
|
|
if video_file:
|
|
# Wait for processing
|
|
processed_file = st.session_state.video_processor.wait_for_file_processing(video_file)
|
|
|
|
if processed_file:
|
|
st.session_state.video_file = processed_file
|
|
st.session_state.video_name = uploaded_file.name
|
|
st.success("✅ Video processed successfully!")
|
|
|
|
# Reset chat when new video is uploaded
|
|
st.session_state.messages = []
|
|
|
|
finally:
|
|
# Clean up temporary file
|
|
try:
|
|
os.unlink(tmp_file_path)
|
|
except:
|
|
pass
|
|
|
|
# Display video preview
|
|
if st.session_state.video_file:
|
|
display_video(uploaded_file.getvalue(), uploaded_file.name)
|
|
|
|
st.markdown("---")
|
|
|
|
# Control buttons
|
|
col1, col2 = st.columns(2)
|
|
with col1:
|
|
if st.button("🗑️ Clear Chat"):
|
|
st.session_state.messages = []
|
|
st.rerun()
|
|
|
|
with col2:
|
|
if st.button("🔄 Reset All"):
|
|
reset_chat()
|
|
st.rerun()
|
|
|
|
# ===========================
|
|
# Main Chat Interface
|
|
# ===========================
|
|
st.title("🎬 Video RAG with Gemini")
|
|
st.markdown("Upload a video and chat with it using Google's Gemini AI!")
|
|
|
|
# Check if ready to chat
|
|
if not api_key:
|
|
st.info("👈 Please enter your Gemini API key in the sidebar to get started.")
|
|
elif st.session_state.video_file is None:
|
|
st.info("👈 Please upload a video file in the sidebar to start chatting.")
|
|
else:
|
|
st.success(f"✅ Ready to chat about: **{st.session_state.video_name}**")
|
|
|
|
# Display chat history
|
|
for message in st.session_state.messages:
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
# Example prompts
|
|
if len(st.session_state.messages) == 0:
|
|
st.markdown("### 💡 Try these example questions:")
|
|
example_prompts = [
|
|
"What is happening in this video?",
|
|
"Summarize the main events",
|
|
"Describe the people and objects you see",
|
|
"What is the setting or environment?",
|
|
"What actions are taking place?"
|
|
]
|
|
|
|
cols = st.columns(2)
|
|
for i, example in enumerate(example_prompts):
|
|
with cols[i % 2]:
|
|
if st.button(f"💬 {example}", key=f"example_{i}"):
|
|
# Simulate user input
|
|
st.session_state.messages.append({"role": "user", "content": example})
|
|
st.rerun()
|
|
|
|
# Chat input
|
|
if prompt := st.chat_input("Ask a question about your video..."):
|
|
# Add user message to chat history
|
|
st.session_state.messages.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.markdown(prompt)
|
|
|
|
# Generate and display assistant response
|
|
with st.chat_message("assistant"):
|
|
message_placeholder = st.empty()
|
|
|
|
with st.spinner("Analyzing video and generating response..."):
|
|
response = st.session_state.video_processor.chat_with_video(
|
|
st.session_state.video_file,
|
|
prompt
|
|
)
|
|
|
|
if response:
|
|
# Simulate streaming effect
|
|
full_response = ""
|
|
for chunk in response.split():
|
|
full_response += chunk + " "
|
|
message_placeholder.markdown(full_response + "▌")
|
|
time.sleep(0.05)
|
|
|
|
# Final response without cursor
|
|
message_placeholder.markdown(response)
|
|
|
|
# Add assistant response to chat history
|
|
st.session_state.messages.append({"role": "assistant", "content": response})
|
|
else:
|
|
st.error("Failed to generate response. Please try again.")
|
|
|
|
# ===========================
|
|
# Footer
|
|
# ===========================
|
|
st.markdown("---")
|
|
st.markdown("""
|
|
<div style='text-align: center; color: #666;'>
|
|
<p>Built with ❤️ using Gemini API and Streamlit |
|
|
<a href='https://ai.google.dev/gemini-api/docs/video-understanding' target='_blank'>Learn more about Gemini Video API</a></p>
|
|
</div>
|
|
""", unsafe_allow_html=True)
|