app.py

import streamlit as st
import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import base64
from PIL import Image
import io
import cv2
from insightface.app import FaceAnalysis
from moviepy.editor import VideoFileClip

# Load models
@st.cache_resource
def load_models():
    unified_model = SentenceTransformer("clip-ViT-B-32")
    face_app = FaceAnalysis(providers=['CPUExecutionProvider'])
    face_app.prepare(ctx_id=0, det_size=(640, 640))
    return unified_model, face_app

unified_model, face_app = load_models()

# Load data
@st.cache_data
def load_data(video_id):
    with open(f"{video_id}_summary.json", "r") as f:
        summary = json.load(f)
    with open(f"{video_id}_transcription.json", "r") as f:
        transcription = json.load(f)
    with open(f"{video_id}_unified_metadata.json", "r") as f:
        unified_metadata = json.load(f)
    with open(f"{video_id}_face_metadata.json", "r") as f:
        face_metadata = json.load(f)
    return summary, transcription, unified_metadata, face_metadata

video_id = "IMFUOexuEXw"
video_path = "avengers_interview.mp4"
summary, transcription, unified_metadata, face_metadata = load_data(video_id)

# Load FAISS indexes
@st.cache_resource
def load_indexes(video_id):
    unified_index = faiss.read_index(f"{video_id}_unified_index.faiss")
    face_index = faiss.read_index(f"{video_id}_face_index.faiss")
    return unified_index, face_index

unified_index, face_index = load_indexes(video_id)

# Search functions
def unified_search(query, index, metadata, model, n_results=5):
    if isinstance(query, str):
        query_vector = model.encode([query], convert_to_tensor=True).cpu().numpy()
    else:  # Assume it's an image
        query_vector = model.encode(query, convert_to_tensor=True).cpu().numpy()
    
    D, I = index.search(query_vector, n_results)
    results = [{'data': metadata[i], 'distance': d} for i, d in zip(I[0], D[0])]
    return results

def face_search(face_embedding, index, metadata, n_results=5):
    D, I = index.search(np.array(face_embedding).reshape(1, -1), n_results)
    results = [metadata[i] for i in I[0]]
    return results, D[0]

def detect_and_embed_face(image, face_app):
    img_array = np.array(image)
    faces = face_app.get(img_array)
    if len(faces) == 0:
        return None
    largest_face = max(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
    return largest_face.embedding

def create_video_clip(video_path, start_time, end_time, output_path):
    with VideoFileClip(video_path) as video:
        new_clip = video.subclip(start_time, end_time)
        new_clip.write_videofile(output_path, codec="libx264", audio_codec="aac")
    return output_path

# Streamlit UI
st.title("Video Analysis Dashboard")

# Sidebar with full video and scrollable transcript
st.sidebar.header("Full Video")
st.sidebar.video(video_path)

st.sidebar.header("Video Transcript")
transcript_text = transcription['transcription']
st.sidebar.text_area("Full Transcript", transcript_text, height=300)

# Main content
st.header("Video Summary")
col1, col2 = st.columns(2)

with col1:
    st.subheader("Prominent Faces")
    for face in summary['prominent_faces']:
        st.write(f"Face ID: {face['id']}, Appearances: {face['appearances']}")
        if 'thumbnail' in face:
            image = Image.open(io.BytesIO(base64.b64decode(face['thumbnail'])))
            st.image(image, caption=f"Face ID: {face['id']}", width=100)

with col2:
    st.subheader("Themes")
    for theme in summary['themes']:
        st.write(f"Theme ID: {theme['id']}, Keywords: {', '.join(theme['keywords'])}")

# Search functionality
st.header("Search")

search_type = st.selectbox("Select search type", ["Unified", "Face"])

if search_type == "Unified":
    search_method = st.radio("Choose search method", ["Text", "Image"])
    
    if search_method == "Text":
        query = st.text_input("Enter your search query")
        if st.button("Search"):
            results = unified_search(query, unified_index, unified_metadata, unified_model)
            st.subheader("Search Results")
            for result in results:
                st.write(f"Time: {result['data']['start']:.2f}s - {result['data']['end']:.2f}s, Distance: {result['distance']:.4f}")
                if 'text' in result['data']:
                    st.write(f"Text: {result['data']['text']}")
                clip_path = create_video_clip(video_path, result['data']['start'], result['data']['end'], f"temp_clip_{result['data']['start']}.mp4")
                st.video(clip_path)
                st.write("---")
    else:
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_column_width=True)
            if st.button("Search"):
                results = unified_search(image, unified_index, unified_metadata, unified_model)
                st.subheader("Image Search Results")
                for result in results:
                    st.write(f"Time: {result['data']['start']:.2f}s - {result['data']['end']:.2f}s, Distance: {result['distance']:.4f}")
                    clip_path = create_video_clip(video_path, result['data']['start'], result['data']['end'], f"temp_clip_{result['data']['start']}.mp4")
                    st.video(clip_path)
                    st.write("---")

elif search_type == "Face":
    face_search_type = st.radio("Choose face search method", ["Select from video", "Upload image"])
    
    if face_search_type == "Select from video":
        face_id = st.selectbox("Select a face", [face['id'] for face in summary['prominent_faces']])
        if st.button("Search"):
            selected_face = next(face for face in summary['prominent_faces'] if face['id'] == face_id)
            face_results, face_distances = face_search(selected_face['embedding'], face_index, face_metadata)
            st.subheader("Face Search Results")
            for result, distance in zip(face_results, face_distances):
                st.write(f"Time: {result['start']:.2f}s - {result['end']:.2f}s, Distance: {distance:.4f}")
                clip_path = create_video_clip(video_path, result['start'], result['end'], f"temp_face_clip_{result['start']}.mp4")
                st.video(clip_path)
                st.write("---")
    else:
        uploaded_file = st.file_uploader("Choose a face image...", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_column_width=True)
            if st.button("Search"):
                face_embedding = detect_and_embed_face(image, face_app)
                if face_embedding is not None:
                    face_results, face_distances = face_search(face_embedding, face_index, face_metadata)
                    st.subheader("Face Search Results")
                    for result, distance in zip(face_results, face_distances):
                        st.write(f"Time: {result['start']:.2f}s - {result['end']:.2f}s, Distance: {distance:.4f}")
                        clip_path = create_video_clip(video_path, result['start'], result['end'], f"temp_face_clip_{result['start']}.mp4")
                        st.video(clip_path)
                        st.write("---")
                else:
                    st.error("No face detected in the uploaded image. Please try another image.")