Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -5,10 +5,13 @@
 Stage 1: SAM2 creates reference mask from first 3 frames
 Stage 2: MatAnyone processes entire video → pha.mp4 
 Stage 3: Create green screen video using pha.mp4 as alpha matte
-Stage 4: Composite green screen video with background → final.mp4
-NEW: Auto-detection of optimal chroma key threshold
-NEW: Chunked processing for long videos with memory management
-FIXED: Memory test disabled to prevent predictor deletion
+Stage 4: Adaptive chroma key compositing with background → final.mp4
+FEATURES:
+- Auto-detection of optimal chroma key threshold
+- Adaptive iteration to find perfect threshold
+- Chunked processing for long videos with memory management
+- Dimension mismatch handling
+- Memory optimization with light/deep cleanup modes
 """
 
 # ==============================================================================
@@ -76,9 +79,14 @@ class ProcessingConfig:
     # Auto-detection settings
     AUTO_DETECTION_FRAMES = 5  # Number of frames to analyze for auto-detection
     AUTO_DETECTION_FALLBACK = 0.1  # Fallback threshold if auto-detection fails
-    MIN_THRESHOLD = 0.05  # Minimum allowed threshold
+    MIN_THRESHOLD = 0.02  # Minimum allowed threshold
     MAX_THRESHOLD = 0.3   # Maximum allowed threshold
     
+    # Adaptive optimization settings
+    ADAPTIVE_MAX_ITERATIONS = 10  # Maximum iterations for threshold optimization
+    ADAPTIVE_GREEN_TOLERANCE = 0.01  # Acceptable green residue level
+    ADAPTIVE_TRANSPARENCY_TOLERANCE = 0.1  # Acceptable transparency quality
+    
     # Quality settings
     VIDEO_CODEC = 'libx264'
     VIDEO_BITRATE = '8000k' 
@@ -514,10 +522,10 @@ class TwoStageProcessor:
     1. SAM2 creates reference mask from first few frames
     2. MatAnyone processes entire video using reference mask → pha.mp4
     3. Create green screen video using pha.mp4 as alpha matte
-    4. Composite green screen video with background → final.mp4
-    5. Auto-detects optimal chroma key threshold
+    4. Adaptive chroma key compositing with background → final.mp4
+    5. Auto-detects and iteratively optimizes chroma key threshold
     6. Chunked processing for long videos with memory management
-    7. FIXED: Memory test disabled to prevent predictor deletion
+    7. Dimension mismatch handling for robust processing
     """
     
     def __init__(self, sam2_handler, matanyone_handler, temp_dir: Optional[str] = None):
@@ -709,11 +717,11 @@ def _process_single_video(self,
             if self.memory_manager.should_clear_memory():
                 self.memory_manager.cleanup_stage("GreenScreen")
             
-            # Stage 4: Final compositing
+            # Stage 4: Adaptive final compositing
             if callback:
-                callback("Final compositing with background...", 90)
-            logger.info("STAGE 4: Final compositing with background...")
-            final_path = self._stage4_final_compositing(green_screen_path, background_path, output_path, quality)
+                callback("Adaptive compositing with background...", 90)
+            logger.info("STAGE 4: Adaptive compositing with background...")
+            final_path = self._stage4_adaptive_compositing(green_screen_path, background_path, output_path, quality)
             
             # Final memory cleanup
             self.memory_manager.cleanup_stage("Final")
@@ -982,7 +990,7 @@ def _trim_video(self, input_path: str, output_path: str, seconds: int):
             raise
 
 # ==============================================================================
-# CHAPTER 10: STAGE 3 - GREEN SCREEN VIDEO CREATION
+# CHAPTER 10: STAGE 3 - GREEN SCREEN VIDEO CREATION WITH DIMENSION FIX
 # ==============================================================================
 
     @MemoryTester.monitor_memory_during_processing
@@ -1096,54 +1104,55 @@ def _stage3_create_green_screen_video(self, original_video_path: str, alpha_vide
             raise
 
 # ==============================================================================
-# CHAPTER 11: STAGE 4 - FINAL COMPOSITING WITH AUTO-DETECTION
+# CHAPTER 11: STAGE 4 - ADAPTIVE CHROMA KEY COMPOSITING
 # ==============================================================================
 
-    @MemoryTester.monitor_memory_during_processing
-    def _stage4_final_compositing(self, green_screen_path: str, background_path: str, output_path: str, quality: str) -> str:
+    @MemoryTester.monitor_memory_during_processing  
+    def _stage4_adaptive_compositing(self, green_screen_path: str, background_path: str, 
+                                    output_path: str, quality: str) -> str:
         """
-        Stage 4: Final compositing of green screen video with background.
+        Stage 4: Final compositing with adaptive threshold optimization.
         
-        Uses MoviePy for high-quality chroma key compositing with auto-detected threshold.
+        Iteratively adjusts chroma key threshold until green is properly removed.
         """
         try:
             # Get quality profile
             profile = QualityManager.get_profile(quality)
             
-            # Load green screen video
+            # Load videos/images
             green_clip = VideoFileClip(green_screen_path)
             
-            # Load and prepare background
             if background_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
-                # Video background
                 bg_clip = VideoFileClip(background_path)
-                # Loop background if shorter than green screen video
                 if bg_clip.duration < green_clip.duration:
                     bg_clip = bg_clip.loop(duration=green_clip.duration)
-                # Trim background if longer
                 bg_clip = bg_clip.subclip(0, green_clip.duration)
             else:
-                # Image background
                 bg_clip = ImageClip(background_path, duration=green_clip.duration)
             
-            # Resize background to match green screen video
             bg_clip = bg_clip.resize((green_clip.w, green_clip.h))
             
-            # Auto-detect optimal chroma key threshold
-            logger.info("Auto-detecting optimal chroma key threshold...")
-            optimal_threshold = self._auto_detect_chroma_threshold(green_screen_path)
-            logger.info(f"Auto-detected threshold: {optimal_threshold}")
+            # Start with auto-detected threshold
+            initial_threshold = self._auto_detect_chroma_threshold(green_screen_path)
+            logger.info(f"Initial auto-detected threshold: {initial_threshold:.4f}")
+            
+            # Adaptive optimization
+            best_threshold = self._optimize_chroma_threshold(
+                green_clip, bg_clip, initial_threshold, 
+                max_iterations=self.config.ADAPTIVE_MAX_ITERATIONS
+            )
+            
+            logger.info(f"✅ Optimized threshold: {best_threshold:.4f} (started from {initial_threshold:.4f})")
             
-            # Apply chroma key with auto-detected threshold
+            # Apply final chroma key with optimized threshold
             green_screen_keyed = green_clip.fx(vfx.mask_color, 
                                              color=self.config.GREEN_COLOR_NORMALIZED,
-                                             thr=optimal_threshold,  # Use auto-detected value
-                                             s=0.2)    # Increased smoothing
+                                             thr=best_threshold,
+                                             s=0.1)  # Reduced smoothing for sharper edges
             
-            # Composite layers
+            # Composite and write
             final_clip = CompositeVideoClip([bg_clip, green_screen_keyed])
             
-            # Write final video
             write_params = {
                 'codec': self.config.VIDEO_CODEC,
                 'bitrate': profile['bitrate'],
@@ -1152,7 +1161,6 @@ def _stage4_final_compositing(self, green_screen_path: str, background_path: str
                 'logger': None
             }
             
-            # Add CRF if specified
             if 'crf' in profile:
                 write_params['ffmpeg_params'] = ['-crf', str(profile['crf'])]
             
@@ -1160,21 +1168,182 @@ def _stage4_final_compositing(self, green_screen_path: str, background_path: str
             
             # Cleanup
             green_clip.close()
-            bg_clip.close() 
+            bg_clip.close()
             final_clip.close()
             
-            # Verify output
             if not os.path.exists(output_path):
                 raise RuntimeError("Final output file was not created")
-                
-            file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
-            logger.info(f"✅ Final compositing completed: {output_path} ({file_size:.1f} MB)")
+            
+            file_size = os.path.getsize(output_path) / (1024 * 1024)
+            logger.info(f"✅ Adaptive compositing completed: {output_path} ({file_size:.1f} MB)")
             
             return output_path
             
         except Exception as e:
-            logger.error(f"Stage 4 failed: {e}")
+            logger.error(f"Stage 4 adaptive failed: {e}")
             raise
+    
+    def _optimize_chroma_threshold(self, green_clip, bg_clip, initial_threshold: float, 
+                                  max_iterations: int = 10) -> float:
+        """
+        Iteratively optimize chroma key threshold by analyzing output quality.
+        """
+        threshold = initial_threshold
+        best_threshold = threshold
+        best_score = float('inf')
+        
+        # Binary search bounds
+        low = max(self.config.MIN_THRESHOLD, initial_threshold * 0.5)
+        high = min(self.config.MAX_THRESHOLD, initial_threshold * 2.0)
+        
+        # Track tested thresholds to avoid repetition
+        tested_thresholds = set()
+        
+        for iteration in range(max_iterations):
+            # Round threshold to avoid tiny differences
+            threshold = round(threshold, 4)
+            
+            # Skip if already tested
+            if threshold in tested_thresholds:
+                logger.info(f"  Threshold {threshold:.4f} already tested, adjusting...")
+                threshold = (low + high) / 2
+                continue
+                
+            tested_thresholds.add(threshold)
+            logger.info(f"🔄 Optimization iteration {iteration + 1}/{max_iterations}, testing threshold: {threshold:.4f}")
+            
+            # Apply chroma key with current threshold
+            keyed = green_clip.fx(vfx.mask_color,
+                                color=self.config.GREEN_COLOR_NORMALIZED,
+                                thr=threshold,
+                                s=0.1)
+            
+            # Composite
+            test_composite = CompositeVideoClip([bg_clip, keyed])
+            
+            # Extract test frames (beginning, middle, end)
+            test_times = [
+                green_clip.duration * 0.1,
+                green_clip.duration * 0.5,
+                green_clip.duration * 0.9
+            ]
+            
+            total_green_score = 0
+            total_transparency_score = 0
+            
+            for test_time in test_times:
+                test_frame = test_composite.get_frame(test_time)
+                bg_frame = bg_clip.get_frame(test_time)
+                
+                # Analyze the frame for green residue
+                green_score = self._analyze_green_residue(test_frame)
+                transparency_score = self._analyze_transparency_quality(test_frame, bg_frame)
+                
+                total_green_score += green_score
+                total_transparency_score += transparency_score
+            
+            # Average scores
+            avg_green_score = total_green_score / len(test_times)
+            avg_transparency_score = total_transparency_score / len(test_times)
+            
+            # Combined score (lower is better)
+            total_score = avg_green_score + avg_transparency_score * 0.5
+            
+            logger.info(f"  📊 Green residue: {avg_green_score:.4f}, Transparency: {avg_transparency_score:.4f}, Total: {total_score:.4f}")
+            
+            # Update best if improved
+            if total_score < best_score:
+                best_score = total_score
+                best_threshold = threshold
+                logger.info(f"  ✅ New best threshold: {best_threshold:.4f} (score: {best_score:.4f})")
+            
+            # Check if we're good enough
+            if avg_green_score < self.config.ADAPTIVE_GREEN_TOLERANCE and \
+               avg_transparency_score < self.config.ADAPTIVE_TRANSPARENCY_TOLERANCE:
+                logger.info(f"  🎯 Acceptable quality reached! Stopping optimization.")
+                break
+            
+            # Adjust threshold using binary search
+            if avg_green_score > 0.05:  # Too much green remains
+                logger.info(f"  🟢 Too much green, decreasing threshold")
+                high = threshold
+                threshold = (low + threshold) / 2
+            elif avg_transparency_score > 0.3:  # Too much was removed
+                logger.info(f"  👤 Subject too transparent, increasing threshold")
+                low = threshold
+                threshold = (threshold + high) / 2
+            else:
+                # Fine-tune around current value
+                if avg_green_score > avg_transparency_score:
+                    threshold *= 0.95  # Slightly more aggressive
+                else:
+                    threshold *= 1.05  # Slightly less aggressive
+            
+            # Ensure we stay in bounds
+            threshold = np.clip(threshold, self.config.MIN_THRESHOLD, self.config.MAX_THRESHOLD)
+            
+            # Clean up test composite
+            test_composite.close()
+            keyed.close()
+            
+            # Stop if converged
+            if abs(threshold - best_threshold) < 0.001 and iteration > 3:
+                logger.info("  📍 Converged, stopping optimization")
+                break
+        
+        return best_threshold
+    
+    def _analyze_green_residue(self, frame: np.ndarray) -> float:
+        """
+        Analyze how much green remains in the frame.
+        Returns score from 0 (no green) to 1 (lots of green).
+        """
+        # Convert to float
+        img = frame.astype(np.float32) / 255.0
+        
+        # Detect pure green pixels
+        green_pixels = (
+            (img[:,:,1] > 0.7) &   # High green
+            (img[:,:,0] < 0.3) &   # Low red
+            (img[:,:,2] < 0.3)     # Low blue
+        )
+        
+        # Calculate percentage of green pixels
+        green_ratio = np.sum(green_pixels) / (frame.shape[0] * frame.shape[1])
+        
+        # Also check for greenish tint in other pixels
+        greenish_pixels = (
+            (img[:,:,1] > img[:,:,0] * 1.5) &  # Green > Red * 1.5
+            (img[:,:,1] > img[:,:,2] * 1.5) &  # Green > Blue * 1.5
+            (img[:,:,1] > 0.4)                 # Significant green
+        )
+        
+        greenish_ratio = np.sum(greenish_pixels) / (frame.shape[0] * frame.shape[1])
+        
+        # Combined score
+        score = green_ratio + greenish_ratio * 0.3
+        
+        return min(1.0, score)
+    
+    def _analyze_transparency_quality(self, composite_frame: np.ndarray, bg_frame: np.ndarray) -> float:
+        """
+        Analyze if too much of the subject was removed.
+        Returns score from 0 (good) to 1 (too much removed).
+        """
+        # Calculate difference between composite and background
+        diff = np.abs(composite_frame.astype(np.float32) - bg_frame.astype(np.float32))
+        
+        # Sum of differences (more difference = more of subject preserved)
+        total_diff = np.sum(diff) / (255.0 * 3 * composite_frame.shape[0] * composite_frame.shape[1])
+        
+        # If difference is too small, too much was removed
+        if total_diff < 0.05:  # Less than 5% different from background
+            return 1.0  # Bad - subject was removed
+        elif total_diff > 0.3:  # More than 30% different
+            return 0.0  # Good - subject well preserved
+        else:
+            # Linear interpolation
+            return 1.0 - (total_diff - 0.05) / 0.25
 
 # ==============================================================================
 # CHAPTER 12: AUTO CHROMA KEY THRESHOLD DETECTION
@@ -1182,16 +1351,16 @@ def _stage4_final_compositing(self, green_screen_path: str, background_path: str
 
     def _auto_detect_chroma_threshold(self, green_screen_path: str) -> float:
         """
-        Auto-detect optimal chroma key threshold by analyzing green screen video.
+        Auto-detect initial chroma key threshold by analyzing green screen video.
         
         This method:
         1. Samples frames from the green screen video
         2. Identifies green background pixels vs person pixels
         3. Calculates color distance between greenest background and person
-        4. Sets threshold as 80% of that distance for safety margin
+        4. Sets threshold as percentage of that distance for initial guess
         """
         try:
-            logger.info("Analyzing green screen video for auto-threshold detection...")
+            logger.info("Analyzing green screen video for initial threshold detection...")
             
             # Open green screen video
             cap = cv2.VideoCapture(green_screen_path)
@@ -1201,7 +1370,6 @@ def _auto_detect_chroma_threshold(self, green_screen_path: str) -> float:
             
             # Get video properties
             total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            fps = cap.get(cv2.CAP_PROP_FPS)
             
             # Sample frames evenly across the video
             frame_indices = np.linspace(0, total_frames - 1, 
@@ -1239,13 +1407,13 @@ def _auto_detect_chroma_threshold(self, green_screen_path: str) -> float:
             
             logger.info(f"Analyzed {len(green_pixels)} green pixels and {len(person_pixels)} person pixels")
             
-            # Calculate optimal threshold
-            threshold = self._calculate_optimal_threshold(green_pixels, person_pixels)
+            # Calculate initial threshold
+            threshold = self._calculate_initial_threshold(green_pixels, person_pixels)
             
             # Clamp to safe range
             threshold = np.clip(threshold, self.config.MIN_THRESHOLD, self.config.MAX_THRESHOLD)
             
-            logger.info(f"Calculated threshold: {threshold:.3f}")
+            logger.info(f"Initial threshold calculated: {threshold:.4f}")
             return threshold
             
         except Exception as e:
@@ -1277,12 +1445,9 @@ def _analyze_frame_colors(self, frame_rgb: np.ndarray) -> Tuple[list, list]:
             )
             
             # Sample pixels (subsample for performance)
-            h, w = frame_rgb.shape[:2]
-            
             # Sample green pixels
             green_coords = np.where(green_mask)
             if len(green_coords[0]) > 1000:
-                # Subsample if too many
                 indices = np.random.choice(len(green_coords[0]), 1000, replace=False)
                 green_coords = (green_coords[0][indices], green_coords[1][indices])
             
@@ -1291,7 +1456,6 @@ def _analyze_frame_colors(self, frame_rgb: np.ndarray) -> Tuple[list, list]:
             # Sample person pixels
             person_coords = np.where(person_mask)
             if len(person_coords[0]) > 1000:
-                # Subsample if too many
                 indices = np.random.choice(len(person_coords[0]), 1000, replace=False)
                 person_coords = (person_coords[0][indices], person_coords[1][indices])
             
@@ -1303,9 +1467,9 @@ def _analyze_frame_colors(self, frame_rgb: np.ndarray) -> Tuple[list, list]:
             logger.warning(f"Frame color analysis failed: {e}")
             return [], []
     
-    def _calculate_optimal_threshold(self, green_pixels: np.ndarray, person_pixels: np.ndarray) -> float:
+    def _calculate_initial_threshold(self, green_pixels: np.ndarray, person_pixels: np.ndarray) -> float:
         """
-        Calculate optimal threshold based on color analysis.
+        Calculate initial threshold based on color analysis.
         
         Strategy:
         1. Find the "least green" green pixels (edge of green screen)
@@ -1320,7 +1484,6 @@ def _calculate_optimal_threshold(self, green_pixels: np.ndarray, person_pixels:
                 person_lab = self._rgb_to_lab_batch(person_pixels)
                 
                 # Use clustering to find representative colors
-                # Find the least green of the green pixels
                 kmeans_green = KMeans(n_clusters=min(5, len(green_lab)//10), random_state=42, n_init=10)
                 green_clusters = kmeans_green.fit_predict(green_lab)
                 green_centers = kmeans_green.cluster_centers_
@@ -1343,11 +1506,10 @@ def _calculate_optimal_threshold(self, green_pixels: np.ndarray, person_pixels:
                 color_distance = np.linalg.norm(main_green_cluster - closest_person_cluster)
                 
                 # Convert LAB distance to threshold
-                # LAB distances typically range 0-100, we want threshold 0-1
-                # Use 60% of the distance as threshold for safety margin
-                threshold = (color_distance / 100.0) * 0.6
+                # Start conservative - use 40% of distance for initial guess
+                threshold = (color_distance / 100.0) * 0.4
                 
-                logger.info(f"Color distance analysis: LAB distance={color_distance:.2f}, threshold={threshold:.3f}")
+                logger.info(f"Color distance analysis: LAB distance={color_distance:.2f}, initial threshold={threshold:.3f}")
                 
                 return threshold
             else:
@@ -1356,7 +1518,6 @@ def _calculate_optimal_threshold(self, green_pixels: np.ndarray, person_pixels:
                 
         except Exception as e:
             logger.warning(f"Threshold calculation failed: {e}")
-            # Fallback: analyze in RGB space
             return self._simple_rgb_threshold(green_pixels, person_pixels)
     
     def _simple_rgb_threshold(self, green_pixels: np.ndarray, person_pixels: np.ndarray) -> float:
@@ -1369,8 +1530,8 @@ def _simple_rgb_threshold(self, green_pixels: np.ndarray, person_pixels: np.ndar
             green_distances = np.linalg.norm(person_pixels - avg_green, axis=1)
             min_distance = np.min(green_distances)
             
-            # Use 70% of minimum distance as threshold
-            threshold = min_distance * 0.7
+            # Use 50% of minimum distance as initial threshold
+            threshold = min_distance * 0.5
             
             logger.info(f"RGB fallback: min_distance={min_distance:.3f}, threshold={threshold:.3f}")
             return threshold
@@ -1466,8 +1627,7 @@ def get_processing_stats(self) -> Dict[str, Any]:
 
 def create_two_stage_processor(sam2_handler, matanyone_handler, **kwargs):
     """
-
-Factory function to create TwoStageProcessor with compatibility layer.
+    Factory function to create TwoStageProcessor with compatibility layer.
     
     This provides a clean interface for integration with existing systems.
     """
@@ -1483,6 +1643,7 @@ def create_two_stage_processor(sam2_handler, matanyone_handler, **kwargs):
     config = ProcessingConfig()
     logger.info(f"Default config: {config.REFERENCE_FRAMES} reference frames, green={config.GREEN_COLOR}")
     logger.info(f"Auto-detection: {config.AUTO_DETECTION_FRAMES} analysis frames, fallback={config.AUTO_DETECTION_FALLBACK}")
+    logger.info(f"Adaptive optimization: {config.ADAPTIVE_MAX_ITERATIONS} max iterations")
     logger.info(f"Chunked processing: {config.MAX_CHUNK_DURATION}s chunks with {config.CHUNK_OVERLAP_FRAMES} frame overlap")
     logger.info(f"Memory management: Cache clearing={'ON' if config.CLEAR_CACHE_AFTER_STAGE else 'OFF'}")
     logger.info(f"Memory testing: {'DISABLED' if not config.ENABLE_MEMORY_TESTING else 'ENABLED'}")