RAG / app.py
Hanzo03's picture
inatial commit
ccdd4a4
# ======================== app.py ========================
import gradio as gr
import os
import json
import shutil # Used for file operations
# Import the core functions from your existing files
from modules.video_analyzer import analyze_video_for_ppe
from modules.rag_indexer import index_analysis_data
from modules.rag_query import run_query
# Configuration (Keep consistent with rag_indexer.py and rag_query.py)
VIDEO_FILENAME = "uploaded_video.mp4" # Temp name for the uploaded file
RAW_ANALYSIS_FILE = 'raw_analysis.json'
DB_PATH = "./chroma_db"
COLLECTION_NAME = 'video_analysis_data' # Use the general collection name
def pipeline_fn(video_file, user_query):
"""
The main function connecting the Gradio inputs to the RAG pipeline.
Args:
video_file: The temporary file object from Gradio (gr.File).
user_query: The text question from Gradio (gr.Textbox).
Returns:
The text response from the RAG query.
"""
if video_file is None:
return "Error: Please upload a video file first."
if not user_query:
return "Error: Please enter a question to query the video analysis."
# 1. Handle File Upload and Naming
# Gradio passes a temporary file path, we need to copy it
# and rename it so the analyzer can find it consistently.
try:
# We copy the file to the expected working directory
temp_video_path = os.path.join(os.getcwd(), VIDEO_FILENAME)
shutil.copy(video_file.name, temp_video_path)
print(f"Copied uploaded file to: {temp_video_path}")
except Exception as e:
return f"File handling error: {e}"
# 2. Analyze Video
print("\n--- STAGE 1: Analyzing Video ---")
# frames_per_sec=0.5 is a sensible default for a quick demo
analysis_results = analyze_video_for_ppe(
video_path=temp_video_path,
frames_per_sec= 2
)
# Save the raw analysis for the indexer to pick up
with open(RAW_ANALYSIS_FILE, 'w') as f:
json.dump(analysis_results, f, indent=4)
# 3. Index Data
print("\n--- STAGE 2: Indexing Analysis Data ---")
# This must be run to create/update the ChromaDB with the new analysis
index_analysis_data(json_file=RAW_ANALYSIS_FILE, collection_name=COLLECTION_NAME)
# 4. Execute RAG Query
print("\n--- STAGE 3: Executing RAG Query ---")
rag_answer = run_query(user_query)
# 5. Cleanup (Optional but Recommended for Demo)
os.remove(temp_video_path) # Remove the copied video file
os.remove(RAW_ANALYSIS_FILE) # Remove the temporary JSON file
return rag_answer
# --- Gradio Interface Definition ---
# Define the input components
video_input = gr.File(
label="Upload Video File (.mp4, .mov, etc.)",
file_types=["video"], # Restrict to video files
type="filepath"
)
query_input = gr.Textbox(
label="Ask a Question about the Video Content",
placeholder="e.g., What are people doing in the video?",
lines=2
)
# Define the output component
output_textbox = gr.Textbox(
label="RAG Analysis Result",
lines=10,
interactive=False
)
# Create the Gradio Interface
demo = gr.Interface(
fn=pipeline_fn,
inputs=[video_input, query_input],
outputs=output_textbox,
title="๐Ÿš€ Video Content RAG Pipeline",
description="Upload a video, and ask a question. The pipeline runs object detection, indexes the data, and uses Gemini to answer your question based on the analysis.",
)
if __name__ == "__main__":
print("Launching Gradio App...")
# This will open the app in your browser at http://127.0.0.1:7860/
demo.launch()