app.py

import cv2
import numpy as np
from transformers import CLIPProcessor, CLIPModel, BlipProcessor, BlipForConditionalGeneration
import torch
from PIL import Image
import faiss
import pickle
from typing import List, Dict, Tuple
import logging
import gradio as gr
import tempfile
import os
import shutil
from tqdm import tqdm
import torch.nn as nn
import math
from concurrent.futures import ThreadPoolExecutor
import numpy as np

class VideoRAGTool:
    def __init__(self, clip_model_name: str = "openai/clip-vit-base-patch32",
                 blip_model_name: str = "Salesforce/blip-image-captioning-base"):
        """Initialize with performance optimizations."""
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Initialize models with optimization flags
        self.clip_model = CLIPModel.from_pretrained(clip_model_name).to(self.device)
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
        self.blip_processor = BlipProcessor.from_pretrained(blip_model_name)
        self.blip_model = BlipForConditionalGeneration.from_pretrained(blip_model_name).to(self.device)
        
        # Enable eval mode for inference
        self.clip_model.eval()
        self.blip_model.eval()
        
        # Batch processing settings
        self.batch_size = 8  # Adjust based on your GPU memory
        
        self.frame_index = None
        self.frame_data = []
        self.logger = self._setup_logger()

    @torch.no_grad()  # Disable gradient computation for inference
    def generate_caption(self, image: Image.Image) -> str:
        """Optimized caption generation."""
        inputs = self.blip_processor(image, return_tensors="pt").to(self.device)
        out = self.blip_model.generate(**inputs, max_length=30, num_beams=2)
        return self.blip_processor.decode(out[0], skip_special_tokens=True)

    def get_video_info(self, video_path: str) -> Tuple[int, float]:
        """Get video frame count and FPS."""
        cap = cv2.VideoCapture(video_path)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        cap.release()
        return total_frames, fps

    def preprocess_frame(self, frame: np.ndarray, target_size: Tuple[int, int] = (224, 224)) -> Image.Image:
        """Preprocess frame with resizing for efficiency."""
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image = Image.fromarray(frame_rgb)
        return image.resize(target_size, Image.LANCZOS)

    @torch.no_grad()
    def process_batch(self, frames: List[Image.Image]) -> Tuple[np.ndarray, List[str]]:
        """Process a batch of frames efficiently."""
        # CLIP processing
        clip_inputs = self.clip_processor(images=frames, return_tensors="pt", padding=True).to(self.device)
        image_features = self.clip_model.get_image_features(**clip_inputs)
        
        # BLIP processing
        captions = []
        blip_inputs = self.blip_processor(images=frames, return_tensors="pt", padding=True).to(self.device)
        out = self.blip_model.generate(**blip_inputs, max_length=30, num_beams=2)
        
        for o in out:
            caption = self.blip_processor.decode(o, skip_special_tokens=True)
            captions.append(caption)
        
        return image_features.cpu().numpy(), captions

    def process_video(self, video_path: str, frame_interval: int = 30) -> None:
        """Optimized video processing with batching and progress tracking."""
        self.logger.info(f"Processing video: {video_path}")
        
        total_frames, fps = self.get_video_info(video_path)
        cap = cv2.VideoCapture(video_path)
        
        # Calculate total batches for progress bar
        frames_to_process = total_frames // frame_interval
        total_batches = math.ceil(frames_to_process / self.batch_size)
        
        current_batch = []
        features_list = []
        frame_count = 0
        
        with tqdm(total=frames_to_process, desc="Processing frames") as pbar:
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                
                if frame_count % frame_interval == 0:
                    # Preprocess frame
                    processed_frame = self.preprocess_frame(frame)
                    current_batch.append(processed_frame)
                    
                    # Process batch when it reaches batch_size
                    if len(current_batch) == self.batch_size:
                        batch_features, batch_captions = self.process_batch(current_batch)
                        
                        # Store results
                        for i, (features, caption) in enumerate(zip(batch_features, batch_captions)):
                            batch_frame_number = frame_count - (self.batch_size - i - 1) * frame_interval
                            self.frame_data.append({
                                'frame_number': batch_frame_number,
                                'timestamp': batch_frame_number / fps,
                                'caption': caption
                            })
                            features_list.append(features)
                        
                        current_batch = []
                        pbar.update(self.batch_size)
                
                frame_count += 1
            
            # Process remaining frames
            if current_batch:
                batch_features, batch_captions = self.process_batch(current_batch)
                for i, (features, caption) in enumerate(zip(batch_features, batch_captions)):
                    batch_frame_number = frame_count - (len(current_batch) - i - 1) * frame_interval
                    self.frame_data.append({
                        'frame_number': batch_frame_number,
                        'timestamp': batch_frame_number / fps,
                        'caption': caption
                    })
                    features_list.append(features)
        
        cap.release()
        
        if not features_list:
            raise ValueError("No frames were processed from the video")
        
        # Create FAISS index
        features_array = np.vstack(features_list)
        self.frame_index = faiss.IndexFlatL2(features_array.shape[1])
        self.frame_index.add(features_array)
        
        self.logger.info(f"Processed {len(self.frame_data)} frames from video")


    def query_video(self, query_text: str, k: int = 5) -> List[Dict]:
        """Query the video using natural language and return relevant frames."""
        self.logger.info(f"Processing query: {query_text}")
        
        inputs = self.clip_processor(text=[query_text], return_tensors="pt").to(self.device)
        text_features = self.clip_model.get_text_features(**inputs)
        
        distances, indices = self.frame_index.search(
            text_features.cpu().detach().numpy(), 
            k
        )
        
        results = []
        for i, (distance, idx) in enumerate(zip(distances[0], indices[0])):
            frame_info = self.frame_data[idx].copy()
            frame_info['relevance_score'] = float(1 / (1 + distance))
            results.append(frame_info)
            
        return results

class VideoRAGApp:
    def __init__(self):
        self.rag_tool = VideoRAGTool()
        self.current_video_path = None
        self.processed = False
        self.temp_dir = tempfile.mkdtemp()

    def __del__(self):
        """Cleanup temporary files on deletion"""
        if hasattr(self, 'temp_dir') and os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir, ignore_errors=True)

    def process_video(self, video_file):
        """Process uploaded video and return status message"""
        try:
            if video_file is None:
                return "Please upload a video first."

            video_path = video_file.name
            temp_video_path = os.path.join(self.temp_dir, "current_video.mp4")
            shutil.copy2(video_path, temp_video_path)
            
            self.current_video_path = temp_video_path
            
            self.rag_tool.process_video(self.current_video_path)
            self.processed = True
            return "Video processed successfully! You can now ask questions about the video."
            
        except Exception as e:
            self.processed = False
            return f"Error processing video: {str(e)}"

    def query_video(self, query_text):
        """Query the video and return relevant frames with descriptions"""
        if not self.processed:
            return None, "Please process a video first."
        
        try:
            results = self.rag_tool.query_video(query_text, k=4)
            
            frames = []
            descriptions = []
            
            cap = cv2.VideoCapture(self.current_video_path)
            
            for result in results:
                frame_number = result['frame_number']
                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
                ret, frame = cap.read()
                
                if ret:
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frames.append(Image.fromarray(frame_rgb))
                    
                    description = f"Timestamp: {result['timestamp']:.2f}s\n"
                    description += f"Scene Description: {result['caption']}\n"
                    description += f"Relevance Score: {result['relevance_score']:.2f}"
                    descriptions.append(description)
            
            cap.release()
            
            # Combine all descriptions with frame numbers
            combined_description = "\n\nFrame Analysis:\n\n"
            for i, desc in enumerate(descriptions, 1):
                combined_description += f"Frame {i}:\n{desc}\n\n"
            
            return frames, combined_description
            
        except Exception as e:
            return None, f"Error querying video: {str(e)}"

    def create_interface(self):
        """Create and return Gradio interface"""
        with gr.Blocks(title="Video Chat RAG") as interface:
            gr.Markdown("# Video Chat RAG")
            gr.Markdown("Upload a video and ask questions about its content!")
            
            with gr.Row():
                video_input = gr.File(
                    label="Upload Video",
                    file_types=["video"],
                )
                process_button = gr.Button("Process Video")
            
            status_output = gr.Textbox(
                label="Status",
                interactive=False
            )
            
            with gr.Row():
                query_input = gr.Textbox(
                    label="Ask about the video",
                    placeholder="What's happening in the video?"
                )
                query_button = gr.Button("Search")
            
            with gr.Row():
                gallery = gr.Gallery(
                    label="Retrieved Frames",
                    show_label=True,
                    elem_id="gallery",
                    columns=[2],
                    rows=[2],
                    height="auto"
                )
                
            descriptions = gr.Textbox(
                label="Scene Descriptions",
                interactive=False,
                lines=10
            )
            
            process_button.click(
                fn=self.process_video,
                inputs=[video_input],
                outputs=[status_output]
            )
            
            query_button.click(
                fn=self.query_video,
                inputs=[query_input],
                outputs=[gallery, descriptions]
            )
        
        return interface

# Initialize and create the interface
app = VideoRAGApp()
interface = app.create_interface()

# Launch the app
if __name__ == "__main__":
    interface.launch()