Create two_stage_processor.py

two_stage_processor.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python3
+"""
+Two-Stage Green Screen Processing System
+Stage 1: Original → Green Screen
+Stage 2: Green Screen → Final Background
+"""
+
+import cv2
+import numpy as np
+import os
+import pickle
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+class TwoStageProcessor:
+    """Handle two-stage video processing with green screen intermediate"""
+    
+    def __init__(self, sam2_predictor=None, matanyone_model=None):
+        self.sam2_predictor = sam2_predictor
+        self.matanyone_model = matanyone_model
+        self.mask_cache_dir = Path("/tmp/mask_cache")
+        self.mask_cache_dir.mkdir(exist_ok=True, parents=True)
+        
+    def stage1_extract_to_greenscreen(self, video_path, output_path, progress_callback=None):
+        """
+        Stage 1: Extract person and create green screen video
+        Also saves masks for potential reuse
+        """
+        def _prog(pct: float, desc: str):
+            if progress_callback:
+                progress_callback(pct, desc)
+        
+        try:
+            _prog(0.0, "Stage 1: Extracting to green screen...")
+            
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                return None, "Could not open video file"
+            
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            
+            # Pure green background for chroma keying
+            green_bg = np.zeros((height, width, 3), dtype=np.uint8)
+            green_bg[:, :] = [0, 255, 0]  # Pure green in BGR
+            
+            # Setup output
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+            
+            # Storage for masks (for potential reuse)
+            masks = []
+            frame_count = 0
+            
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                _prog(0.1 + (frame_count / max(1, total_frames)) * 0.8,
+                      f"Stage 1: Processing frame {frame_count + 1}/{total_frames}")
+                
+                # Get mask using SAM2
+                mask = self._extract_person_mask(frame)
+                masks.append(mask)
+                
+                # Refine mask every 3rd frame with MatAnyone
+                if frame_count % 3 == 0 and self.matanyone_model:
+                    mask = self._refine_mask(frame, mask)
+                
+                # Apply green screen with HARD edges for clean keying
+                result = self._apply_greenscreen_hard(frame, mask, green_bg)
+                out.write(result)
+                
+                frame_count += 1
+            
+            cap.release()
+            out.release()
+            
+            # Save masks for potential reuse
+            mask_file = self.mask_cache_dir / f"{Path(output_path).stem}_masks.pkl"
+            with open(mask_file, 'wb') as f:
+                pickle.dump(masks, f)
+            
+            _prog(1.0, "Stage 1 complete: Green screen created")
+            return output_path, f"Green screen created: {frame_count} frames"
+            
+        except Exception as e:
+            logger.error(f"Stage 1 error: {e}")
+            return None, f"Stage 1 failed: {str(e)}"
+    
+    def stage2_greenscreen_to_final(self, greenscreen_path, background, output_path, 
+                                   chroma_settings=None, progress_callback=None):
+        """
+        Stage 2: Replace green screen with final background using chroma keying
+        """
+        def _prog(pct: float, desc: str):
+            if progress_callback:
+                progress_callback(pct, desc)
+        
+        if chroma_settings is None:
+            chroma_settings = {
+                'key_color': [0, 255, 0],  # Green in BGR
+                'tolerance': 40,
+                'edge_softness': 2,
+                'spill_suppression': 0.3
+            }
+        
+        try:
+            _prog(0.0, "Stage 2: Applying final background...")
+            
+            cap = cv2.VideoCapture(greenscreen_path)
+            if not cap.isOpened():
+                return None, "Could not open green screen video"
+            
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            
+            # Prepare background
+            if isinstance(background, str):
+                bg = cv2.imread(background)
+                if bg is None:
+                    return None, "Could not load background image"
+            else:
+                bg = background
+            
+            bg = cv2.resize(bg, (width, height))
+            
+            # Setup output
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+            
+            frame_count = 0
+            
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                
+                _prog(0.1 + (frame_count / max(1, total_frames)) * 0.8,
+                      f"Stage 2: Compositing frame {frame_count + 1}/{total_frames}")
+                
+                # Apply chroma keying
+                result = self._chroma_key_advanced(frame, bg, chroma_settings)
+                out.write(result)
+                
+                frame_count += 1
+            
+            cap.release()
+            out.release()
+            
+            _prog(1.0, "Stage 2 complete: Final video created")
+            return output_path, f"Final video created: {frame_count} frames"
+            
+        except Exception as e:
+            logger.error(f"Stage 2 error: {e}")
+            return None, f"Stage 2 failed: {str(e)}"
+    
+    def _extract_person_mask(self, frame):
+        """Extract person mask using SAM2"""
+        if self.sam2_predictor is None:
+            # Fallback mask
+            h, w = frame.shape[:2]
+            mask = np.zeros((h, w), dtype=np.uint8)
+            mask[h//6:5*h//6, w//4:3*w//4] = 255
+            return mask
+        
+        try:
+            self.sam2_predictor.set_image(frame)
+            h, w = frame.shape[:2]
+            
+            # Strategic points for person
+            points = np.array([
+                [w//2, h//3],     # Head
+                [w//2, h//2],     # Torso
+                [w//2, 2*h//3],   # Lower body
+                [w//3, h//2],     # Left
+                [2*w//3, h//2],   # Right
+            ])
+            labels = np.ones(len(points))
+            
+            masks, scores, _ = self.sam2_predictor.predict(
+                point_coords=points,
+                point_labels=labels,
+                multimask_output=True
+            )
+            
+            best_idx = np.argmax(scores)
+            mask = masks[best_idx]
+            
+            if mask.dtype != np.uint8:
+                mask = (mask * 255).astype(np.uint8)
+            
+            return mask
+            
+        except Exception as e:
+            logger.error(f"Mask extraction error: {e}")
+            h, w = frame.shape[:2]
+            mask = np.zeros((h, w), dtype=np.uint8)
+            mask[h//6:5*h//6, w//4:3*w//4] = 255
+            return mask
+    
+    def _refine_mask(self, frame, mask):
+        """Refine mask using MatAnyone if available"""
+        if self.matanyone_model is None:
+            return mask
+        
+        try:
+            # MatAnyone refinement logic here
+            # This would depend on your MatAnyone implementation
+            return mask
+        except:
+            return mask
+    
+    def _apply_greenscreen_hard(self, frame, mask, green_bg):
+        """Apply green screen with hard edges for clean chroma keying"""
+        # Binary threshold for clean edges
+        _, mask_binary = cv2.threshold(mask, 140, 255, cv2.THRESH_BINARY)
+        
+        # No feathering - we want hard edges for chroma keying
+        mask_3ch = cv2.cvtColor(mask_binary, cv2.COLOR_GRAY2BGR)
+        mask_norm = mask_3ch.astype(float) / 255
+        
+        # Composite
+        result = frame * mask_norm + green_bg * (1 - mask_norm)
+        return result.astype(np.uint8)
+    
+    def _chroma_key_advanced(self, frame, background, settings):
+        """
+        Advanced chroma keying with spill suppression
+        """
+        key_color = np.array(settings['key_color'], dtype=np.uint8)
+        tolerance = settings['tolerance']
+        softness = settings['edge_softness']
+        spill_suppress = settings['spill_suppression']
+        
+        # Convert to float for processing
+        frame_float = frame.astype(np.float32)
+        bg_float = background.astype(np.float32)
+        
+        # Calculate color distance from key color
+        diff = np.abs(frame_float - key_color)
+        distance = np.sqrt(np.sum(diff ** 2, axis=2))
+        
+        # Create mask based on distance
+        mask = np.where(distance < tolerance, 0, 1)
+        
+        # Edge softening
+        if softness > 0:
+            mask = cv2.GaussianBlur(mask.astype(np.float32), 
+                                   (softness*2+1, softness*2+1), 
+                                   softness)
+        
+        # Spill suppression - reduce green in edges
+        if spill_suppress > 0:
+            green_channel = frame_float[:, :, 1]
+            spill_mask = np.where(mask < 1, 1 - mask, 0)
+            green_suppression = green_channel * spill_mask * spill_suppress
+            frame_float[:, :, 1] -= green_suppression
+            frame_float = np.clip(frame_float, 0, 255)
+        
+        # Expand mask to 3 channels
+        mask_3ch = np.stack([mask] * 3, axis=2)
+        
+        # Composite
+        result = frame_float * mask_3ch + bg_float * (1 - mask_3ch)
+        return np.clip(result, 0, 255).astype(np.uint8)
+    
+    def process_full_pipeline(self, video_path, background, final_output, 
+                             chroma_settings=None, progress_callback=None):
+        """
+        Run the complete two-stage pipeline
+        """
+        import tempfile
+        
+        # Stage 1: Create green screen
+        greenscreen_path = tempfile.mktemp(suffix='_greenscreen.mp4')
+        gs_result, gs_msg = self.stage1_extract_to_greenscreen(
+            video_path, greenscreen_path, progress_callback
+        )
+        
+        if gs_result is None:
+            return None, gs_msg
+        
+        # Stage 2: Apply final background
+        final_result, final_msg = self.stage2_greenscreen_to_final(
+            greenscreen_path, background, final_output, 
+            chroma_settings, progress_callback
+        )
+        
+        # Cleanup
+        try:
+            os.remove(greenscreen_path)
+        except:
+            pass
+        
+        return final_result, final_msg
+
+# Chroma key settings presets
+CHROMA_PRESETS = {
+    'standard': {
+        'key_color': [0, 255, 0],
+        'tolerance': 40,
+        'edge_softness': 2,
+        'spill_suppression': 0.3
+    },
+    'tight': {
+        'key_color': [0, 255, 0],
+        'tolerance': 30,
+        'edge_softness': 1,
+        'spill_suppression': 0.4
+    },
+    'soft': {
+        'key_color': [0, 255, 0],
+        'tolerance': 50,
+        'edge_softness': 3,
+        'spill_suppression': 0.2
+    }
+}