app.py

import cv2
import numpy as np
from transformers import CLIPProcessor, CLIPModel, Blip2Processor, Blip2ForConditionalGeneration
import torch
from PIL import Image
import faiss
import logging
import gradio as gr
import tempfile
import os
import shutil
from tqdm.auto import tqdm
from pathlib import Path
from typing import List, Dict, Tuple, Optional
import gc
import warnings
warnings.filterwarnings("ignore")

# Configure model caching and environment
os.environ["TRANSFORMERS_CACHE"] = "./model_cache"
os.environ["HF_HOME"] = "./model_cache"
os.makedirs("./model_cache", exist_ok=True)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class VideoProcessor:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logging.info(f"Using device: {self.device}")
        
        # Load models with optimizations
        self._load_models()
        
        # Processing settings
        self.frame_interval = 30  # Process 1 frame every 30 frames
        self.max_frames = 50      # Maximum frames to process
        self.target_size = (224, 224)
        self.batch_size = 4 if torch.cuda.is_available() else 2

    def _load_models(self):
        """Load models with optimizations and proper configurations"""
        try:
            logging.info("Loading CLIP model...")
            self.clip_model = CLIPModel.from_pretrained(
                "openai/clip-vit-base-patch32",
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                cache_dir="./model_cache"
            ).to(self.device)
            self.clip_processor = CLIPProcessor.from_pretrained(
                "openai/clip-vit-base-patch32",
                cache_dir="./model_cache"
            )

            logging.info("Loading BLIP2 model...")
            model_name = "Salesforce/blip2-opt-2.7b"
            
            # Initialize BLIP2 with minimal configuration
            self.blip_processor = Blip2Processor.from_pretrained(
                model_name,
                cache_dir="./model_cache"
            )

            self.blip_model = Blip2ForConditionalGeneration.from_pretrained(
                model_name,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None,
                cache_dir="./model_cache",
                low_cpu_mem_usage=True
            ).to(self.device)

            # Set models to evaluation mode
            self.clip_model.eval()
            self.blip_model.eval()
            
            logging.info("Models loaded successfully!")
        except Exception as e:
            logging.error(f"Error loading models: {str(e)}")
            raise

    def _preprocess_frame(self, frame: np.ndarray) -> Image.Image:
        """Preprocess a single frame"""
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        return Image.fromarray(rgb_frame).resize(self.target_size, Image.LANCZOS)

    @torch.no_grad()
    def process_frame_batch(self, frames: List[np.ndarray]) -> Tuple[Optional[np.ndarray], Optional[List[str]]]:
        """Process a batch of frames efficiently"""
        try:
            # Convert frames to PIL Images
            pil_frames = [self._preprocess_frame(f) for f in frames]
            
            # Get CLIP features
            clip_inputs = self.clip_processor(
                images=pil_frames, 
                return_tensors="pt", 
                padding=True
            ).to(self.device)
            
            if self.device.type == "cuda":
                clip_inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in clip_inputs.items()}
            features = self.clip_model.get_image_features(**clip_inputs)
            
            # Get BLIP captions
            blip_inputs = self.blip_processor(
                images=pil_frames,
                return_tensors="pt",
                padding=True
            ).to(self.device)
            
            if self.device.type == "cuda":
                blip_inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in blip_inputs.items()}
            
            # Generate captions
            captions = self.blip_model.generate(
                **blip_inputs,
                max_length=30,
                min_length=10,
                num_beams=5,
                length_penalty=1.0,
                temperature=0.7,
                do_sample=False
            )
            
            captions = [self.blip_processor.decode(c, skip_special_tokens=True) for c in captions]
            
            # Clear GPU memory if needed
            if self.device.type == "cuda":
                torch.cuda.empty_cache()
            
            return features.cpu().numpy(), captions
            
        except Exception as e:
            logging.error(f"Error in batch processing: {str(e)}")
            return None, None

    def process_video(self, video_path: str, progress: gr.Progress) -> Tuple[Optional[faiss.Index], Optional[List[Dict]], str]:
        """Process video with batching and progress updates"""
        cap = None
        try:
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                raise ValueError(f"Could not open video file: {video_path}")
            
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = cap.get(cv2.CAP_PROP_FPS)
            
            # Calculate frames to process
            frames_to_process = min(self.max_frames, total_frames // self.frame_interval)
            progress(0, desc="Initializing video processing...")
            
            features_list = []
            frame_data = []
            current_batch = []
            batch_positions = []
            
            frame_count = 0
            processed_count = 0
            
            while processed_count < frames_to_process:
                ret, frame = cap.read()
                if not ret:
                    break
                
                if frame_count % self.frame_interval == 0:
                    current_batch.append(frame)
                    batch_positions.append(frame_count)
                    
                    # Process batch when full
                    if len(current_batch) == self.batch_size or frame_count == total_frames - 1:
                        progress(processed_count / frames_to_process, 
                               desc=f"Processing frames... {processed_count}/{frames_to_process}")
                        
                        features, captions = self.process_frame_batch(current_batch)
                        
                        if features is not None and captions is not None:
                            for i, (feat, cap_text) in enumerate(zip(features, captions)):
                                features_list.append(feat)
                                frame_data.append({
                                    'frame_number': batch_positions[i],
                                    'timestamp': batch_positions[i] / fps,
                                    'caption': cap_text
                                })
                        
                        processed_count += len(current_batch)
                        current_batch = []
                        batch_positions = []
                
                frame_count += 1
            
            # Create FAISS index
            if features_list:
                features_array = np.vstack(features_list)
                frame_index = faiss.IndexFlatL2(features_array.shape[1])
                frame_index.add(features_array)
                return frame_index, frame_data, "Video processed successfully!"
            else:
                return None, None, "No frames were processed successfully."

        except Exception as e:
            logging.error(f"Error processing video: {str(e)}")
            return None, None, f"Error processing video: {str(e)}"
            
        finally:
            if cap is not None:
                cap.release()
            gc.collect()
            if self.device.type == "cuda":
                torch.cuda.empty_cache()

class VideoQAInterface:
    def __init__(self):
        self.processor = VideoProcessor()
        self.frame_index = None
        self.frame_data = None
        self.processed = False
        self.current_video_path = None
        self.temp_dir = tempfile.mkdtemp()
        logging.info(f"Initialized temp directory: {self.temp_dir}")

    def __del__(self):
        """Cleanup temporary files"""
        try:
            if hasattr(self, 'temp_dir') and os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
                logging.info(f"Cleaned up temp directory: {self.temp_dir}")
        except Exception as e:
            logging.error(f"Error cleaning up temp directory: {str(e)}")

    def process_video(self, video_file, progress=gr.Progress()):
        """Process video with progress tracking"""
        if video_file is None:
            return "Please upload a video first."
            
        try:
            # Save uploaded video to temp directory
            temp_video_path = os.path.join(self.temp_dir, "input_video.mp4")
            shutil.copy2(video_file.name, temp_video_path)
            self.current_video_path = temp_video_path
            logging.info(f"Saved video to: {self.current_video_path}")
            
            progress(0, desc="Starting video processing...")
            self.frame_index, self.frame_data, message = self.processor.process_video(
                self.current_video_path, progress
            )
            
            if self.frame_index is not None:
                self.processed = True
                return "Video processed successfully! You can now ask questions."
            else:
                self.processed = False
                return message
                
        except Exception as e:
            self.processed = False
            logging.error(f"Error processing video: {str(e)}")
            return f"Error processing video: {str(e)}"

    @torch.no_grad()
    def answer_question(self, query):
        """Answer questions about the video"""
        if not self.processed or self.current_video_path is None:
            return None, "Please process a video first."
            
        try:
            # Get query features
            inputs = self.processor.clip_processor(text=[query], return_tensors="pt").to(self.processor.device)
            query_features = self.processor.clip_model.get_text_features(**inputs)
            
            # Search for relevant frames
            k = 4  # Number of frames to retrieve
            D, I = self.frame_index.search(query_features.cpu().numpy(), k)
            
            results = []
            for distance, idx in zip(D[0], I[0]):
                frame_info = self.frame_data[idx].copy()
                frame_info['relevance'] = float(1 / (1 + distance))
                results.append(frame_info)
            
            # Format output
            descriptions = []
            frames = []
            
            cap = cv2.VideoCapture(self.current_video_path)
            try:
                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))
                        
                        desc = f"Timestamp: {result['timestamp']:.2f}s\n"
                        desc += f"Scene Description: {result['caption']}\n"
                        desc += f"Relevance Score: {result['relevance']:.2f}"
                        descriptions.append(desc)
            finally:
                cap.release()
            
            if not frames:
                return None, "No relevant frames found."
            
            combined_desc = "\n\nFrame Analysis:\n\n"
            for i, desc in enumerate(descriptions, 1):
                combined_desc += f"Frame {i}:\n{desc}\n\n"
                
            return frames, combined_desc
            
        except Exception as e:
            logging.error(f"Error answering question: {str(e)}")
            return None, f"Error answering question: {str(e)}"

    def create_interface(self):
        """Create Gradio interface"""
        with gr.Blocks(title="Advanced Video Question Answering") as interface:
            gr.Markdown("# Advanced Video Question Answering")
            gr.Markdown("Upload a video and ask questions about any aspect of its content!")
            
            with gr.Row():
                with gr.Column():
                    video_input = gr.File(
                        label="Upload Video",
                        file_types=["video"]
                    )
                    status = gr.Textbox(label="Status", interactive=False)
                    process_btn = gr.Button("Process Video")
            
            with gr.Row():
                query_input = gr.Textbox(
                    label="Ask about the video",
                    placeholder="What's happening in the video?"
                )
                query_btn = gr.Button("Search")
            
            gallery = gr.Gallery(
                label="Retrieved Frames",
                show_label=True,
                columns=[2],
                rows=[2]
            )
            
            descriptions = gr.Textbox(
                label="Analysis",
                interactive=False,
                lines=10
            )
            
            # Set up event handlers
            process_btn.click(
                fn=self.process_video,
                inputs=[video_input],
                outputs=[status]
            )
            
            query_btn.click(
                fn=self.answer_question,
                inputs=[query_input],
                outputs=[gallery, descriptions]
            )
        
        return interface

# Create and launch the app
app = VideoQAInterface()
interface = app.create_interface()

if __name__ == "__main__":
    interface.launch(
        server_name="0.0.0.0",
        share=False,
        show_error=True
    )