Update app.py

app.py

@@ -4,243 +4,219 @@ from transformers import CLIPProcessor, CLIPModel, Blip2Processor, Blip2ForCondi
 import torch
 from PIL import Image
 import faiss
-from typing import List, Dict, Tuple
 import logging
 import gradio as gr
 import tempfile
 import os
-import shutil
-from tqdm import tqdm
+from tqdm.auto import tqdm
 from pathlib import Path
-from moviepy.video.io.VideoFileClip import VideoFileClip
+import time
 
-class VideoRAGSystem:
+class VideoProcessor:
     def __init__(self):
-        self.logger = self.setup_logger()
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.logger.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
 
-        # Initialize models
-        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(self.device)
+    def load_models(self):
+        """Load models with optimizations"""
+        # Load CLIP
+        self.clip_model = CLIPModel.from_pretrained(
+            "openai/clip-vit-base-patch32",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        ).to(self.device)
         self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-        
-        self.blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+
+        # Load BLIP2 with reduced size
         self.blip_model = Blip2ForConditionalGeneration.from_pretrained(
             "Salesforce/blip2-opt-2.7b",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None
         ).to(self.device)
+        self.blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
 
-        # Vector store setup
-        self.frame_index = None
-        self.frame_data = []
-        self.target_size = (224, 224)
-        
-        # Create directories for storing processed data
-        self.temp_dir = tempfile.mkdtemp()
-        self.frames_dir = os.path.join(self.temp_dir, "frames")
-        os.makedirs(self.frames_dir, exist_ok=True)
-
-    def setup_logger(self) -> logging.Logger:
-        logger = logging.getLogger('VideoRAGSystem')
-        if logger.handlers:
-            logger.handlers.clear()
-        logger.setLevel(logging.INFO)
-        handler = logging.StreamHandler()
-        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-        handler.setFormatter(formatter)
-        logger.addHandler(handler)
-        return logger
-
-    def split_video(self, video_path: str, timestamp_ms: int, context_seconds: int = 3) -> str:
-        """Extract a clip around the specified timestamp"""
-        timestamp_sec = timestamp_ms / 1000
-        output_path = os.path.join(self.temp_dir, "clip.mp4")
-        
-        with VideoFileClip(video_path) as video:
-            duration = video.duration
-            start_time = max(timestamp_sec - context_seconds, 0)
-            end_time = min(timestamp_sec + context_seconds, duration)
-            clip = video.subclip(start_time, end_time)
-            clip.write_videofile(output_path, audio_codec='aac')
-        
-        return output_path
+        # Set models to evaluation mode
+        self.clip_model.eval()
+        self.blip_model.eval()
 
     @torch.no_grad()
-    def analyze_frame(self, image: Image.Image) -> Dict:
-        """Comprehensive frame analysis"""
+    def process_frame_batch(self, frames):
+        """Process a batch of frames efficiently"""
         try:
-            # Generate caption
-            inputs = self.blip_processor(image, return_tensors="pt").to(self.device)
-            if self.device.type == "cuda":
-                inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in inputs.items()}
-            caption = self.blip_model.generate(**inputs, max_length=50)
-            caption_text = self.blip_processor.decode(caption[0], skip_special_tokens=True)
-
-            # Get visual features
-            clip_inputs = self.clip_processor(images=image, return_tensors="pt").to(self.device)
+            # Convert frames to PIL Images
+            pil_frames = [Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)).resize(self.target_size) 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)
-
-            return {
-                "caption": caption_text,
-                "features": features.cpu().numpy()
-            }
+            
+            # 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()}
+            captions = self.blip_model.generate(**blip_inputs, max_length=30)
+            captions = [self.blip_processor.decode(c, skip_special_tokens=True) for c in captions]
+            
+            return features.cpu().numpy(), captions
         except Exception as e:
-            self.logger.error(f"Frame analysis error: {str(e)}")
-            return None
+            print(f"Error in batch processing: {str(e)}")
+            return None, None
 
-    def extract_keyframes(self, video_path: str, max_frames: int = 15) -> List[Dict]:
-        """Extract and analyze key frames"""
+    def process_video(self, video_path: str, progress=gr.Progress()):
+        """Process video with batching and progress updates"""
         cap = cv2.VideoCapture(video_path)
-        frames_info = []
-        frame_count = 0
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        interval = max(1, total_frames // max_frames)
-
-        with tqdm(total=max_frames, desc="Analyzing frames") as pbar:
-            while len(frames_info) < max_frames and cap.isOpened():
+        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...")
+        
+        features_list = []
+        frame_data = []
+        current_batch = []
+        batch_positions = []
+        
+        try:
+            frame_count = 0
+            processed_count = 0
+            
+            while cap.isOpened() and processed_count < frames_to_process:
                 ret, frame = cap.read()
                 if not ret:
                     break
-
-                if frame_count % interval == 0:
-                    # Save frame
-                    frame_path = os.path.join(self.frames_dir, f"frame_{frame_count}.jpg")
-                    cv2.imwrite(frame_path, frame)
                     
-                    # Analyze frame
-                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                    image = Image.fromarray(frame_rgb).resize(self.target_size, Image.LANCZOS)
-                    analysis = self.analyze_frame(image)
+                if frame_count % self.frame_interval == 0:
+                    current_batch.append(frame)
+                    batch_positions.append(frame_count)
                     
-                    if analysis is not None:
-                        frames_info.append({
-                            "frame_number": frame_count,
-                            "timestamp": frame_count / cap.get(cv2.CAP_PROP_FPS),
-                            "path": frame_path,
-                            "caption": analysis["caption"],
-                            "features": analysis["features"]
-                        })
-                        pbar.update(1)
-
+                    # Process batch when full
+                    if len(current_batch) == self.batch_size or frame_count == total_frames - 1:
+                        features, captions = self.process_frame_batch(current_batch)
+                        
+                        if features is not None and captions is not None:
+                            for i, (feat, cap) in enumerate(zip(features, captions)):
+                                features_list.append(feat)
+                                frame_data.append({
+                                    'frame_number': batch_positions[i],
+                                    'timestamp': batch_positions[i] / fps,
+                                    'caption': cap
+                                })
+                        
+                        processed_count += len(current_batch)
+                        current_batch = []
+                        batch_positions = []
+                        
+                        # Update progress
+                        progress(processed_count / frames_to_process, 
+                               desc=f"Processing frames... {processed_count}/{frames_to_process}")
+                
                 frame_count += 1
-
-        cap.release()
-        return frames_info
-
-    def process_video(self, video_path: str):
-        """Process video and build search index"""
-        self.logger.info(f"Processing video: {video_path}")
-        
-        try:
-            # Extract and analyze frames
-            frames_info = self.extract_keyframes(video_path)
-            self.frame_data = frames_info
-
-            # Build FAISS index
-            if frames_info:
-                features = np.vstack([frame["features"] for frame in frames_info])
-                self.frame_index = faiss.IndexFlatL2(features.shape[1])
-                self.frame_index.add(features)
                 
-            self.logger.info(f"Processed {len(frames_info)} frames successfully")
-            return True
+            cap.release()
             
+            # 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:
-            self.logger.error(f"Video processing error: {str(e)}")
-            return False
-
-    @torch.no_grad()
-    def search_frames(self, query: str, k: int = 4) -> List[Dict]:
-        """Search for relevant frames based on the query"""
-        try:
-            # Process query
-            inputs = self.clip_processor(text=[query], return_tensors="pt").to(self.device)
-            if self.device.type == "cuda":
-                inputs = {k: v.half() if v.dtype == torch.float32 else v for k, v in inputs.items()}
-            query_features = self.clip_model.get_text_features(**inputs)
-
-            # Search
-            distances, indices = self.frame_index.search(
-                query_features.cpu().numpy(),
-                k
-            )
-
-            # Prepare results
-            results = []
-            for distance, idx in zip(distances[0], indices[0]):
-                frame_info = self.frame_data[idx].copy()
-                frame_info["relevance"] = float(1 / (1 + distance))
-                results.append(frame_info)
-
-            return results
-
-        except Exception as e:
-            self.logger.error(f"Search error: {str(e)}")
-            return []
+            cap.release()
+            return None, None, f"Error processing video: {str(e)}"
 
 class VideoQAInterface:
     def __init__(self):
-        self.rag_system = VideoRAGSystem()
-        self.current_video = None
+        self.processor = VideoProcessor()
+        self.frame_index = None
+        self.frame_data = None
         self.processed = False
 
-    def process_video(self, video_file):
-        """Handle video upload and processing"""
+    def process_video(self, video_file, progress=gr.Progress()):
+        """Process video with progress tracking"""
         try:
             if video_file is None:
-                return "Please upload a video first.", gr.Progress(0)
-
-            self.current_video = video_file.name
-            success = self.rag_system.process_video(self.current_video)
+                return "Please upload a video first."
+                
+            progress(0, desc="Starting video processing...")
+            self.frame_index, self.frame_data, message = self.processor.process_video(
+                video_file.name, progress
+            )
             
-            if success:
+            if self.frame_index is not None:
                 self.processed = True
-                return "Video processed successfully! You can now ask questions.", gr.Progress(100)
+                return "Video processed successfully! You can now ask questions."
             else:
-                return "Error processing video. Please try again.", gr.Progress(0)
-
+                self.processed = False
+                return message
+                
         except Exception as e:
             self.processed = False
-            return f"Error: {str(e)}", gr.Progress(0)
+            return f"Error: {str(e)}"
 
+    @torch.no_grad()
     def answer_question(self, query):
-        """Handle question answering"""
+        """Answer questions about the video"""
         if not self.processed:
             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
-            results = self.rag_system.search_frames(query)
+            k = 4  # Number of frames to retrieve
+            D, I = self.frame_index.search(query_features.cpu().numpy(), k)
             
-            if not results:
-                return None, "No relevant frames found."
-
-            # Prepare output
-            frames = []
+            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(video_file.name)
             for result in results:
-                # Load frame
-                frame = Image.open(result["path"])
-                frames.append(frame)
+                frame_number = result['frame_number']
+                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+                ret, frame = cap.read()
                 
-                # Prepare description
-                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)
-
-            # Combine descriptions
+                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)
+            
+            cap.release()
+            
             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:
-            return None, f"Error: {str(e)}"
+            return None, f"Error answering question: {str(e)}"
 
     def create_interface(self):
         """Create Gradio interface"""
@@ -249,45 +225,42 @@ class VideoQAInterface:
             gr.Markdown("Upload a video and ask questions about any aspect of 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.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_button = gr.Button("Search")
+                query_btn = gr.Button("Search")
             
             gallery = gr.Gallery(
                 label="Retrieved Frames",
                 show_label=True,
                 columns=[2],
-                rows=[2],
-                height="auto"
+                rows=[2]
             )
-                
+            
             descriptions = gr.Textbox(
-                label="Scene Analysis",
+                label="Analysis",
                 interactive=False,
                 lines=10
             )
             
-            process_button.click(
+            # Set up event handlers
+            process_btn.click(
                 fn=self.process_video,
                 inputs=[video_input],
-                outputs=[status_output]
+                outputs=[status]
             )
             
-            query_button.click(
+            query_btn.click(
                 fn=self.answer_question,
                 inputs=[query_input],
                 outputs=[gallery, descriptions]
@@ -295,10 +268,9 @@ class VideoQAInterface:
         
         return interface
 
-# Initialize and create the interface
+# Create and launch the app
 app = VideoQAInterface()
 interface = app.create_interface()
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()