Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python3
 """
-Two-Stage Green-Screen Processing System with Comprehensive Debug Outputs ✅ 2025-09-09
-FIXED: Now uses enhanced SAM2Handler with full-body detection strategies
-Stage 1: Original → keyed background (auto-selected colour)
-Stage 2: Keyed video → final composite (hybrid chroma + segmentation rescue)
-DEBUG: Added extensive debug outputs to isolate quality degradation at every step
+EFFICIENT Two-Stage Green-Screen Processing System ✅ 2025-09-09
+CORRECTED APPROACH: SAM2 creates reference mask from first 3 frames, MatAnyone handles entire video
+Stage 1: SAM2 processes first 3 frames → Creates reference mask
+Stage 2: MatAnyone processes entire video using reference mask → Final composite
+This is the intended efficient approach, not frame-by-frame SAM2 processing!
 """
 from __future__ import annotations
 import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
@@ -55,17 +55,17 @@ def create_video_writer(output_path: str, fps: float, width: int, height: int, p
     raise RuntimeError(f"Could not create video writer for {output_path}")
 
 # ---------------------------------------------------------------------------
-# TwoStageProcessor: Main processing pipeline with debug outputs
+# TwoStageProcessor: EFFICIENT processing pipeline
 # ---------------------------------------------------------------------------
 class TwoStageProcessor:
     """
-    Two-stage video background replacement with comprehensive debugging:
-    Stage 1: Person segmentation → Green screen video  
-    Stage 2: Green screen video → Final composite
+    EFFICIENT Two-stage video background replacement:
+    Stage 1: SAM2 processes first 3 frames → Creates reference mask
+    Stage 2: MatAnyone processes entire video using reference mask → Final output
     """
     
     def __init__(self, sam2_handler, matanyone_processor=None):
-        """Initialize with enhanced SAM2Handler and optional MatAnyone."""
+        """Initialize with enhanced SAM2Handler and MatAnyone."""
         # CRITICAL: Ensure we receive the full SAM2Handler, not just predictor
         if hasattr(sam2_handler, 'create_mask'):
             self.sam2_handler = sam2_handler
@@ -100,90 +100,6 @@ def __init__(self, sam2_handler, matanyone_processor=None):
             logger.warning(f"Could not access quality_manager profile: {e}, using default")
             self.quality_profile = {"name": "default", "settings": {}}
 
-    def _get_mask(self, frame_rgb, frame_index):
-        """Get person mask with extensive debug outputs at every step."""
-        logger.info(f"=== DEBUG: TwoStageProcessor._get_mask called for frame {frame_index} ===")
-        logger.info(f"DEBUG: Input frame shape: {frame_rgb.shape}, dtype: {frame_rgb.dtype}")
-        
-        try:
-            # Check what type of SAM2 handler we have
-            if hasattr(self.sam2_handler, 'create_mask'):
-                logger.info("DEBUG: Using ENHANCED SAM2Handler.create_mask() with full-body detection")
-                
-                # Call SAM2Handler's create_mask method (which has its own debug outputs)
-                mask = self.sam2_handler.create_mask(frame_rgb)
-                
-            else:
-                logger.info("DEBUG: Using basic SAM2 segmentation fallback")
-                mask = self._basic_sam2_fallback(frame_rgb)
-            
-            # DEBUG: Check if TwoStageProcessor is getting the mask correctly
-            if mask is not None:
-                logger.info(f"DEBUG: TwoStageProcessor received mask shape: {mask.shape}, dtype: {mask.dtype}")
-                logger.info(f"DEBUG: TwoStageProcessor received mask range: {mask.min():.3f} to {mask.max():.3f}")
-                logger.info(f"DEBUG: TwoStageProcessor received mask non-zero pixels: {np.count_nonzero(mask)}")
-                
-                # Convert to uint8 for saving if needed
-                if mask.dtype != np.uint8:
-                    debug_mask = (mask * 255).astype(np.uint8)
-                else:
-                    debug_mask = mask.copy()
-                
-                cv2.imwrite(f"/tmp/debug_twostage_received_mask_frame_{frame_index}.png", debug_mask)
-                logger.info(f"DEBUG: Saved TwoStageProcessor received mask to /tmp/debug_twostage_received_mask_frame_{frame_index}.png")
-                
-                # DISABLED: No additional cleaning - let SAM2Handler's cleaning stand
-                # mask = self._maximum_quality_mask_cleaning(mask)  # DISABLED FOR TESTING
-                logger.info("DEBUG: Double cleaning DISABLED - using mask as received from SAM2Handler")
-                
-                return mask
-            else:
-                logger.warning(f"DEBUG: No mask received from SAM2Handler for frame {frame_index}")
-                return None
-                
-        except Exception as e:
-            logger.error(f"DEBUG: Exception in TwoStageProcessor._get_mask for frame {frame_index}: {e}")
-            logger.error(f"DEBUG: Traceback: {traceback.format_exc()}")
-            return None
-
-    def _basic_sam2_fallback(self, frame_rgb):
-        """Fallback for basic SAM2 predictor without enhanced features."""
-        try:
-            logger.info("DEBUG: Running basic SAM2 fallback")
-            self.sam2_handler.set_image(frame_rgb)
-            # Simple center point prompt
-            h, w = frame_rgb.shape[:2]
-            center_point = np.array([[w//2, h//2]])
-            center_label = np.array([1])
-            
-            masks, scores, _ = self.sam2_handler.predict(
-                point_coords=center_point,
-                point_labels=center_label,
-                multimask_output=True
-            )
-            
-            if len(masks) > 0:
-                # Select best mask by score
-                best_idx = np.argmax(scores)
-                mask = masks[best_idx].astype(np.float32)
-                logger.info(f"DEBUG: Basic SAM2 fallback - selected mask {best_idx} with score {scores[best_idx]:.3f}")
-                return mask
-            else:
-                logger.warning("DEBUG: Basic SAM2 fallback produced no masks")
-                return None
-        except Exception as e:
-            logger.error(f"DEBUG: Basic SAM2 fallback failed: {e}")
-            return None
-
-    def _maximum_quality_mask_cleaning(self, mask):
-        """
-        DISABLED FOR TESTING - Maximum quality mask post-processing.
-        Returns mask unchanged to prevent double cleaning artifacts.
-        """
-        logger.info("DEBUG: _maximum_quality_mask_cleaning called but DISABLED")
-        logger.info("DEBUG: Returning mask unchanged to prevent double cleaning artifacts")
-        return mask
-
     def process_video(self, 
                      input_path: str,
                      background_image_path: Optional[str] = None,
@@ -191,46 +107,48 @@ def process_video(self,
                      callback: Optional[Callable] = None,
                      trim_seconds: float = 5.0) -> Tuple[Optional[str], Optional[str], str]:
         """
-        Process video with two-stage background replacement.
-        Returns: (final_output_path, green_screen_path, status_message)
+        EFFICIENT video processing with SAM2 reference mask + MatAnyone.
+        Returns: (final_output_path, reference_mask_path, status_message)
         """
         try:
-            self.temp_dir = tempfile.mkdtemp(prefix="twostage_debug_")
-            logger.info(f"DEBUG: Processing video: {input_path}")
-            logger.info(f"DEBUG: Temp directory: {self.temp_dir}")
+            self.temp_dir = tempfile.mkdtemp(prefix="twostage_efficient_")
+            logger.info(f"🎬 EFFICIENT Processing: {input_path}")
+            logger.info(f"📁 Temp directory: {self.temp_dir}")
             
-            # Stage 1: Create green screen video
-            green_screen_path = self._stage1_create_greenscreen(
-                input_path, trim_seconds, callback
+            # Stage 1: Create reference mask from first 3 frames using SAM2
+            reference_mask_path = self._stage1_create_reference_mask(
+                input_path, callback
             )
             
-            if not green_screen_path:
-                return None, None, "❌ Stage 1 failed: Could not create green screen video"
+            if not reference_mask_path:
+                return None, None, "❌ Stage 1 failed: Could not create reference mask"
             
             if callback:
-                callback(50, "Stage 1 complete, starting Stage 2...")
+                callback(25, "Reference mask created, starting MatAnyone processing...")
             
-            # Stage 2: Apply background to green screen video  
-            final_output_path = self._stage2_apply_background(
-                green_screen_path, background_image_path, background_preset, callback
+            # Stage 2: Process entire video with MatAnyone using reference mask
+            final_output_path = self._stage2_matanyone_processing(
+                input_path, reference_mask_path, background_image_path, 
+                background_preset, callback, trim_seconds
             )
             
             if not final_output_path:
-                return None, green_screen_path, "❌ Stage 2 failed: Could not apply background"
+                return None, reference_mask_path, "❌ Stage 2 failed: MatAnyone processing failed"
             
-            logger.info(f"✅ Two-stage processing complete: {final_output_path}")
-            return final_output_path, green_screen_path, "✅ Video processed successfully!"
+            logger.info(f"✅ EFFICIENT processing complete: {final_output_path}")
+            return final_output_path, reference_mask_path, "✅ Video processed successfully with efficient SAM2 + MatAnyone pipeline!"
             
         except Exception as e:
-            logger.error(f"Two-stage processing failed: {e}")
+            logger.error(f"EFFICIENT processing failed: {e}")
             logger.error(traceback.format_exc())
             return None, None, f"❌ Processing failed: {str(e)}"
         
         finally:
-            self._cleanup()
+            # Don't auto-cleanup - let caller decide
+            pass
 
-    def _stage1_create_greenscreen(self, input_path: str, trim_seconds: float, callback: Optional[Callable]) -> Optional[str]:
-        """Stage 1: Create green screen video using person segmentation with debug."""
+    def _stage1_create_reference_mask(self, input_path: str, callback: Optional[Callable]) -> Optional[str]:
+        """Stage 1: Create reference mask from first 3 frames using SAM2."""
         try:
             cap = cv2.VideoCapture(input_path)
             if not cap.isOpened():
@@ -243,230 +161,362 @@ def _stage1_create_greenscreen(self, input_path: str, trim_seconds: float, callb
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
             total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             
-            if trim_seconds > 0:
-                max_frames = int(fps * trim_seconds)
-                total_frames = min(total_frames, max_frames)
-            
-            logger.info(f"DEBUG: Stage 1 - {total_frames} frames, {fps:.2f} FPS, {width}x{height}")
+            logger.info(f"📊 Video: {total_frames} frames, {fps:.2f} FPS, {width}x{height}")
             
-            # Output path
-            green_screen_path = os.path.join(self.temp_dir, "green_screen_debug.mp4")
-            writer = create_video_writer(green_screen_path, fps, width, height)
-            
-            # Define green screen color (force consistent green)
-            green_color = (0, 255, 0)  # Pure green in BGR
-            
-            frame_count = 0
-            successful_masks = 0
-            failed_masks = 0
+            # Process first 3 frames to create robust reference mask
+            reference_frames = []
+            masks = []
             
-            while frame_count < total_frames:
+            for frame_idx in range(min(3, total_frames)):
                 ret, frame = cap.read()
                 if not ret:
                     break
                 
+                reference_frames.append(frame)
+                
                 # Convert to RGB for SAM2
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 
-                # Get person mask with debug outputs
-                mask = self._get_mask(frame_rgb, frame_count)
+                logger.info(f"🎯 Processing reference frame {frame_idx + 1}/3 with SAM2...")
+                mask = self.sam2_handler.create_mask(frame_rgb)
                 
                 if mask is not None:
-                    successful_masks += 1
-                    
                     # Ensure mask is proper format
                     if mask.dtype != np.uint8:
                         mask = (mask * 255).astype(np.uint8)
-                    
-                    # Create green screen frame
-                    green_frame = np.full_like(frame, green_color, dtype=np.uint8)
-                    
-                    # Apply mask to original frame
-                    if len(mask.shape) == 2:
-                        mask_3ch = cv2.merge([mask, mask, mask])
-                    else:
-                        mask_3ch = mask
-                    
-                    # Normalize mask to 0-1 range
-                    mask_norm = mask_3ch.astype(np.float32) / 255.0
-                    
-                    # Composite: person on green background
-                    result_frame = (frame.astype(np.float32) * mask_norm + 
-                                   green_frame.astype(np.float32) * (1 - mask_norm)).astype(np.uint8)
-                    
-                    writer.write(result_frame)
-                    
-                    # Save debug frame occasionally
-                    if frame_count % 30 == 0:  # Every 30th frame
-                        debug_frame_path = f"/tmp/debug_greenscreen_frame_{frame_count}.png"
-                        cv2.imwrite(debug_frame_path, result_frame)
-                        logger.info(f"DEBUG: Saved green screen frame to {debug_frame_path}")
-                        
+                    masks.append(mask)
+                    logger.info(f"✅ Reference frame {frame_idx + 1} mask created successfully")
                 else:
-                    failed_masks += 1
-                    # No mask detected - use pure green frame
-                    logger.warning(f"DEBUG: No mask detected for frame {frame_count}, using pure green")
-                    green_frame = np.full_like(frame, green_color, dtype=np.uint8)
-                    writer.write(green_frame)
+                    logger.warning(f"⚠️ Reference frame {frame_idx + 1} failed to generate mask")
                 
-                frame_count += 1
-                
-                # Progress callback
-                if callback and frame_count % 10 == 0:
-                    progress = int(25 * frame_count / total_frames)  # Stage 1 is 0-25%
-                    callback(progress, f"Creating green screen... {frame_count}/{total_frames}")
+                if callback and frame_idx < 2:
+                    callback(5 + frame_idx * 7, f"Processing reference frame {frame_idx + 1}/3...")
             
             cap.release()
-            writer.release()
             
-            logger.info(f"DEBUG: Stage 1 complete - Successful masks: {successful_masks}, Failed masks: {failed_masks}")
+            if not masks:
+                logger.error("❌ No valid masks generated from reference frames")
+                return None
+            
+            # Create composite reference mask from all successful masks
+            logger.info(f"🔄 Creating composite reference mask from {len(masks)} frames...")
+            reference_mask = self._create_composite_mask(masks, height, width)
+            
+            # Save reference mask
+            reference_mask_path = os.path.join(self.temp_dir, "reference_mask.png")
+            cv2.imwrite(reference_mask_path, reference_mask)
             
-            if os.path.exists(green_screen_path) and os.path.getsize(green_screen_path) > 0:
-                logger.info(f"✅ Stage 1 complete: {green_screen_path}")
-                self._cleanup_list.append(green_screen_path)
-                return green_screen_path
+            # Verify mask was saved
+            if os.path.exists(reference_mask_path) and os.path.getsize(reference_mask_path) > 0:
+                logger.info(f"✅ Reference mask saved: {reference_mask_path}")
+                logger.info(f"📐 Mask coverage: {np.mean(reference_mask / 255.0):.3f}")
+                self._cleanup_list.append(reference_mask_path)
+                return reference_mask_path
             else:
-                logger.error("Stage 1 failed: Empty output file")
+                logger.error("❌ Failed to save reference mask")
                 return None
                 
         except Exception as e:
-            logger.error(f"Stage 1 failed: {e}")
+            logger.error(f"❌ Stage 1 failed: {e}")
             logger.error(traceback.format_exc())
             return None
 
-    def _stage2_apply_background(self, green_screen_path: str, background_image_path: Optional[str], 
-                                background_preset: str, callback: Optional[Callable]) -> Optional[str]:
-        """Stage 2: Apply background to green screen video using professional chroma key."""
+    def _create_composite_mask(self, masks: List[np.ndarray], height: int, width: int) -> np.ndarray:
+        """Create a robust composite mask from multiple reference frames."""
         try:
-            # Load background
-            background = self._prepare_background(background_image_path, background_preset)
-            if background is None:
-                logger.error("Failed to prepare background")
+            if len(masks) == 1:
+                return masks[0]
+            
+            # Method 1: Union of all masks (conservative approach)
+            composite = np.zeros((height, width), dtype=np.float32)
+            
+            for mask in masks:
+                # Normalize mask to 0-1
+                if mask.dtype != np.float32:
+                    mask_norm = mask.astype(np.float32) / 255.0
+                else:
+                    mask_norm = mask
+                
+                # Resize if needed
+                if mask_norm.shape != (height, width):
+                    mask_norm = cv2.resize(mask_norm, (width, height))
+                
+                # Union operation
+                composite = np.maximum(composite, mask_norm)
+            
+            # Method 2: Voting-based approach for stability
+            vote_threshold = len(masks) * 0.6  # 60% of frames must agree
+            vote_composite = np.zeros((height, width), dtype=np.float32)
+            
+            for mask in masks:
+                mask_norm = mask.astype(np.float32) / 255.0
+                if mask_norm.shape != (height, width):
+                    mask_norm = cv2.resize(mask_norm, (width, height))
+                vote_composite += (mask_norm > 0.5).astype(np.float32)
+            
+            voting_mask = (vote_composite >= vote_threshold).astype(np.float32)
+            
+            # Combine both approaches - use union but refine with voting
+            final_mask = composite * 0.7 + voting_mask * 0.3
+            
+            # Post-process for smoothness
+            final_mask = cv2.GaussianBlur(final_mask, (5, 5), 1.0)
+            final_mask = (final_mask > 0.4).astype(np.float32)
+            
+            # Convert back to uint8
+            return (final_mask * 255).astype(np.uint8)
+            
+        except Exception as e:
+            logger.error(f"Composite mask creation failed: {e}")
+            # Fallback to first mask
+            return masks[0] if masks else np.zeros((height, width), dtype=np.uint8)
+
+    def _stage2_matanyone_processing(self, input_path: str, reference_mask_path: str,
+                                   background_image_path: Optional[str], background_preset: str,
+                                   callback: Optional[Callable], trim_seconds: float) -> Optional[str]:
+        """Stage 2: Process entire video with MatAnyone using reference mask."""
+        try:
+            if not self.matanyone:
+                logger.error("❌ MatAnyone processor not available")
                 return None
             
-            cap = cv2.VideoCapture(green_screen_path)
-            if not cap.isOpened():
-                logger.error(f"Cannot open green screen video: {green_screen_path}")
+            # Trim video if needed
+            if trim_seconds > 0:
+                trimmed_video_path = self._trim_video(input_path, trim_seconds, callback)
+                if not trimmed_video_path:
+                    return None
+                video_to_process = trimmed_video_path
+            else:
+                video_to_process = input_path
+            
+            if callback:
+                callback(50, "Starting MatAnyone processing...")
+            
+            # Create output directory for MatAnyone
+            matanyone_output_dir = os.path.join(self.temp_dir, "matanyone_output")
+            os.makedirs(matanyone_output_dir, exist_ok=True)
+            
+            logger.info(f"🚀 Processing entire video with MatAnyone...")
+            logger.info(f"📹 Video: {video_to_process}")
+            logger.info(f"🎭 Reference mask: {reference_mask_path}")
+            logger.info(f"📁 Output dir: {matanyone_output_dir}")
+            
+            # Process with MatAnyone
+            matanyone_result = self.matanyone.process_video(
+                input_path=video_to_process,
+                mask_path=reference_mask_path,
+                output_path=matanyone_output_dir,
+                max_size=1024  # Reasonable resolution for processing
+            )
+            
+            if callback:
+                callback(80, "MatAnyone processing complete, applying background...")
+            
+            # Handle MatAnyone result (it might return tuple or string)
+            alpha_video_path = None
+            if matanyone_result:
+                if isinstance(matanyone_result, (list, tuple)) and len(matanyone_result) > 1:
+                    alpha_video_path = matanyone_result[1]  # Second element is usually alpha video
+                elif isinstance(matanyone_result, str):
+                    alpha_video_path = matanyone_result
+                else:
+                    # Search for alpha files in output directory
+                    for pattern in ["alpha.mp4", "alpha.mkv", "alpha.mov", "alpha.webm"]:
+                        candidate = os.path.join(matanyone_output_dir, pattern)
+                        if os.path.exists(candidate):
+                            alpha_video_path = candidate
+                            break
+            
+            if not alpha_video_path or not os.path.exists(alpha_video_path):
+                logger.error("❌ MatAnyone failed to produce alpha video")
                 return None
             
-            # Video properties
-            fps = cap.get(cv2.CAP_PROP_FPS)
-            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            logger.info(f"✅ MatAnyone alpha video created: {alpha_video_path}")
             
-            logger.info(f"DEBUG: Stage 2 - {total_frames} frames, {fps:.2f} FPS, {width}x{height}")
+            # Stage 3: Composite with background
+            final_output_path = self._stage3_composite_background(
+                video_to_process, alpha_video_path, background_image_path, 
+                background_preset, callback
+            )
             
-            # Resize background to match video
-            background = cv2.resize(background, (width, height))
-            logger.info(f"DEBUG: Background resized to {width}x{height}")
+            return final_output_path
             
-            # Output path
-            final_output_path = os.path.join(self.temp_dir, "final_output_debug.mp4")
-            writer = create_video_writer(final_output_path, fps, width, height)
+        except Exception as e:
+            logger.error(f"❌ Stage 2 failed: {e}")
+            logger.error(traceback.format_exc())
+            return None
+
+    def _trim_video(self, input_path: str, trim_seconds: float, callback: Optional[Callable]) -> Optional[str]:
+        """Trim video to specified duration."""
+        try:
+            from moviepy.editor import VideoFileClip
             
-            frame_count = 0
-            while frame_count < total_frames:
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                
-                # Professional chroma key for green screen
-                result_frame = self._professional_chroma_key(frame, background)
-                writer.write(result_frame)
-                
-                # Save debug frame occasionally
-                if frame_count % 30 == 0:  # Every 30th frame
-                    debug_final_path = f"/tmp/debug_final_frame_{frame_count}.png"
-                    cv2.imwrite(debug_final_path, result_frame)
-                    logger.info(f"DEBUG: Saved final frame to {debug_final_path}")
-                
-                frame_count += 1
+            trimmed_path = os.path.join(self.temp_dir, "trimmed_video.mp4")
+            
+            logger.info(f"✂️ Trimming video to {trim_seconds}s...")
+            
+            clip = VideoFileClip(input_path)
+            duration = min(trim_seconds, float(clip.duration or trim_seconds))
+            trimmed_clip = clip.subclip(0, duration)
+            
+            # Write with optimized settings
+            trimmed_clip.write_videofile(
+                trimmed_path,
+                audio=False,
+                fps=clip.fps,
+                codec="libx264",
+                logger=None,
+                ffmpeg_params=["-preset", "fast", "-pix_fmt", "yuv420p"]
+            )
+            
+            clip.close()
+            trimmed_clip.close()
+            
+            if os.path.exists(trimmed_path):
+                logger.info(f"✅ Video trimmed: {trimmed_path}")
+                self._cleanup_list.append(trimmed_path)
+                return trimmed_path
+            else:
+                logger.error("❌ Video trimming failed")
+                return None
                 
-                # Progress callback
-                if callback and frame_count % 10 == 0:
-                    progress = int(25 + 75 * frame_count / total_frames)  # Stage 2 is 25-100%
-                    callback(progress, f"Applying background... {frame_count}/{total_frames}")
+        except Exception as e:
+            logger.error(f"❌ Video trimming failed: {e}")
+            return None
+
+    def _stage3_composite_background(self, original_video_path: str, alpha_video_path: str,
+                                   background_image_path: Optional[str], background_preset: str,
+                                   callback: Optional[Callable]) -> Optional[str]:
+        """Stage 3: Composite alpha video with background."""
+        try:
+            from moviepy.editor import VideoFileClip
             
-            cap.release()
-            writer.release()
+            # Load videos
+            original_clip = VideoFileClip(original_video_path)
+            alpha_clip = VideoFileClip(alpha_video_path)
+            
+            # Get video properties
+            width = int(original_clip.w)
+            height = int(original_clip.h)
+            
+            # Prepare background
+            background = self._prepare_background(background_image_path, background_preset, width, height)
+            if background is None:
+                logger.error("❌ Failed to prepare background")
+                return None
+            
+            logger.info(f"🎨 Compositing final video...")
+            
+            def composite_frame(get_frame, t):
+                """Composite frame with alpha matte."""
+                try:
+                    original_frame = get_frame(t)
+                    
+                    # Get alpha frame (handle duration mismatch)
+                    alpha_duration = float(alpha_clip.duration or 0)
+                    alpha_t = min(t, max(0, alpha_duration - 0.01)) if alpha_duration > 0 else 0
+                    alpha_frame = alpha_clip.get_frame(alpha_t)
+                    
+                    # Process alpha
+                    if alpha_frame.ndim == 3:
+                        alpha_frame = alpha_frame[:, :, 0]  # Take first channel
+                    
+                    # Normalize alpha to 0-1 range
+                    if alpha_frame.max() > 1.1:
+                        alpha_frame = alpha_frame / 255.0
+                    
+                    alpha_frame = np.clip(alpha_frame, 0, 1)
+                    
+                    # Resize alpha if needed
+                    if alpha_frame.shape != (height, width):
+                        alpha_frame = cv2.resize(alpha_frame, (width, height))
+                    
+                    # Add channel dimension
+                    alpha_frame = alpha_frame[:, :, np.newaxis]
+                    
+                    # Prepare background
+                    bg_normalized = background.astype(np.float32) / 255.0
+                    
+                    # Composite
+                    result = alpha_frame * original_frame + (1 - alpha_frame) * bg_normalized
+                    
+                    return np.clip(result, 0, 1).astype(np.float32)
+                    
+                except Exception as e:
+                    logger.error(f"Frame compositing error: {e}")
+                    return original_frame if 'original_frame' in locals() else np.zeros((height, width, 3), dtype=np.float32)
+            
+            # Create final composite
+            final_clip = original_clip.fl(composite_frame)
+            
+            # Output path
+            final_output_path = os.path.join(self.temp_dir, "final_output.mp4")
+            
+            # Write final video
+            final_clip.write_videofile(
+                final_output_path,
+                audio=False,
+                fps=original_clip.fps,
+                codec="libx264",
+                logger=None,
+                ffmpeg_params=["-preset", "fast", "-pix_fmt", "yuv420p"]
+            )
+            
+            original_clip.close()
+            alpha_clip.close()
+            final_clip.close()
+            
+            if callback:
+                callback(100, "Final compositing completed!")
             
             if os.path.exists(final_output_path) and os.path.getsize(final_output_path) > 0:
-                logger.info(f"✅ Stage 2 complete: {final_output_path}")
+                logger.info(f"✅ Final video created: {final_output_path}")
                 return final_output_path
             else:
-                logger.error("Stage 2 failed: Empty output file")
+                logger.error("❌ Final video creation failed")
                 return None
                 
         except Exception as e:
-            logger.error(f"Stage 2 failed: {e}")
+            logger.error(f"❌ Stage 3 failed: {e}")
             logger.error(traceback.format_exc())
             return None
 
-    def _professional_chroma_key(self, frame, background):
-        """Professional chroma key for green screen removal."""
-        # Convert to HSV for better green detection
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-        
-        # Define green color range in HSV
-        # Green hue is around 60 in OpenCV HSV (0-179 range)
-        lower_green = np.array([40, 40, 40])    # Lower bound
-        upper_green = np.array([80, 255, 255])  # Upper bound
-        
-        # Create mask for green pixels
-        green_mask = cv2.inRange(hsv, lower_green, upper_green)
-        
-        # Smooth the mask
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
-        green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel)
-        green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel)
-        
-        # Blur for soft edges
-        green_mask = cv2.GaussianBlur(green_mask, (5, 5), 0)
-        
-        # Convert mask to 3-channel and normalize
-        green_mask_3ch = cv2.merge([green_mask, green_mask, green_mask])
-        alpha = green_mask_3ch.astype(np.float32) / 255.0
-        
-        # Composite: background where green, original where not green
-        result = (background.astype(np.float32) * alpha + 
-                 frame.astype(np.float32) * (1 - alpha)).astype(np.uint8)
-        
-        return result
-
-    def _prepare_background(self, background_image_path: Optional[str], background_preset: str) -> Optional[np.ndarray]:
+    def _prepare_background(self, background_image_path: Optional[str], background_preset: str, 
+                          width: int, height: int) -> Optional[np.ndarray]:
         """Prepare background image from file or preset."""
         try:
             if background_image_path and os.path.exists(background_image_path):
-                logger.info(f"Loading custom background: {background_image_path}")
+                logger.info(f"📷 Loading custom background: {background_image_path}")
                 background = cv2.imread(background_image_path)
                 if background is not None:
+                    # Convert BGR to RGB and resize
+                    background = cv2.cvtColor(background, cv2.COLOR_BGR2RGB)
+                    background = cv2.resize(background, (width, height))
                     return background
                 logger.warning(f"Failed to load custom background, falling back to preset")
             
             # Generate preset background
-            logger.info(f"Using preset background: {background_preset}")
-            return self._generate_preset_background(background_preset)
+            logger.info(f"🎨 Using preset background: {background_preset}")
+            return self._generate_preset_background(background_preset, width, height)
             
         except Exception as e:
             logger.error(f"Background preparation failed: {e}")
             return None
 
-    def _generate_preset_background(self, preset: str) -> np.ndarray:
-        """Generate preset background (fallback implementation)."""
-        # Simple solid color backgrounds
+    def _generate_preset_background(self, preset: str, width: int, height: int) -> np.ndarray:
+        """Generate preset background."""
         color_map = {
-            "office": (240, 240, 240),     # Light gray
-            "blue": (255, 200, 100),       # Light blue
+            "office": (245, 246, 248),     # Light gray
+            "blue": (100, 200, 255),       # Light blue
             "studio": (200, 200, 200),     # Medium gray  
             "white": (255, 255, 255),      # White
             "black": (0, 0, 0)             # Black
         }
         
-        color = color_map.get(preset, (240, 240, 240))  # Default to office
-        background = np.full((720, 1280, 3), color, dtype=np.uint8)  # Default size
+        color = color_map.get(preset, (245, 246, 248))  # Default to office
+        background = np.full((height, width, 3), color, dtype=np.uint8)
         
-        logger.info(f"Generated {preset} background with color {color}")
+        logger.info(f"Generated {preset} background ({width}x{height}) with color {color}")
         return background
 
     def _cleanup(self):
@@ -486,18 +536,15 @@ def _cleanup(self):
             logger.warning(f"Cleanup failed: {e}")
 
 class BasicSAM2Wrapper:
-    """Wrapper for basic SAM2 predictor to provide create_mask interface with debug."""
+    """Wrapper for basic SAM2 predictor to provide create_mask interface."""
     
     def __init__(self, predictor):
         self.predictor = predictor
-        logger.info("DEBUG: Created BasicSAM2Wrapper for compatibility")
+        logger.info("Created BasicSAM2Wrapper for compatibility")
     
     def create_mask(self, frame_rgb):
-        """Basic mask creation using center point prompt with debug outputs."""
+        """Basic mask creation using center point prompt."""
         try:
-            logger.info("DEBUG: BasicSAM2Wrapper.create_mask() called")
-            logger.info(f"DEBUG: BasicSAM2Wrapper input frame shape: {frame_rgb.shape}")
-            
             self.predictor.set_image(frame_rgb)
             
             # Simple center point prompt
@@ -515,22 +562,12 @@ def create_mask(self, frame_rgb):
                 # Select best mask by score
                 best_idx = np.argmax(scores)
                 mask = masks[best_idx].astype(np.float32)
-                
-                logger.info(f"DEBUG: BasicSAM2Wrapper generated {len(masks)} masks")
-                logger.info(f"DEBUG: BasicSAM2Wrapper selected mask {best_idx} with score {scores[best_idx]:.3f}")
-                logger.info(f"DEBUG: BasicSAM2Wrapper mask has {np.count_nonzero(mask)} non-zero pixels")
-                
-                # Save debug mask
-                debug_mask = (mask * 255).astype(np.uint8)
-                cv2.imwrite("/tmp/debug_basic_sam2_wrapper_mask.png", debug_mask)
-                logger.info("DEBUG: Saved BasicSAM2Wrapper mask to /tmp/debug_basic_sam2_wrapper_mask.png")
-                
+                logger.info(f"BasicSAM2Wrapper: Generated mask with {np.count_nonzero(mask)} pixels")
                 return mask
             
-            logger.warning("DEBUG: BasicSAM2Wrapper - No masks generated")
+            logger.warning("BasicSAM2Wrapper: No masks generated")
             return None
             
         except Exception as e:
-            logger.error(f"DEBUG: BasicSAM2Wrapper failed: {e}")
-            logger.error(f"DEBUG: BasicSAM2Wrapper traceback: {traceback.format_exc()}")
+            logger.error(f"BasicSAM2Wrapper failed: {e}")
             return None