Update two_stage_processor.py

two_stage_processor.py

@@ -12,6 +12,8 @@
 import logging
 from pathlib import Path
 import tempfile
+import traceback
+from utilities import refine_mask_hq  # Import for MatAnyone refinement
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +40,7 @@ def _prog(pct: float, desc: str):
             
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
+                logger.error("Could not open video file")
                 return None, "Could not open video file"
             
             fps = cap.get(cv2.CAP_PROP_FPS)
@@ -49,9 +52,14 @@ def _prog(pct: float, desc: str):
             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))
+            # Setup output using app.py's create_video_writer
+            from app import create_video_writer  # Import here to avoid circular imports
+            out, actual_output_path = create_video_writer(output_path, fps, width, height)
+            if out is None:
+                cap.release()
+                logger.error("Could not create output video file")
+                return None, "Could not create output video file"
+            output_path = actual_output_path
             
             # Storage for masks (for potential reuse)
             masks = []
@@ -87,6 +95,7 @@ def _prog(pct: float, desc: str):
             try:
                 with open(mask_file, 'wb') as f:
                     pickle.dump(masks, f)
+                logger.info(f"Masks saved to {mask_file}")
             except Exception as e:
                 logger.warning(f"Failed to save masks: {e}")
             
@@ -94,7 +103,7 @@ def _prog(pct: float, desc: str):
             return output_path, f"Green screen created: {frame_count} frames"
             
         except Exception as e:
-            logger.error(f"Stage 1 error: {e}")
+            logger.error(f"Stage 1 error: {e}\n{traceback.format_exc()}")
             return None, f"Stage 1 failed: {str(e)}"
     
     def stage2_greenscreen_to_final(self, greenscreen_path, background, output_path, 
@@ -107,18 +116,14 @@ def _prog(pct: float, desc: str):
                 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
-            }
+            chroma_settings = CHROMA_PRESETS['standard']
         
         try:
             _prog(0.0, "Stage 2: Applying final background...")
             
             cap = cv2.VideoCapture(greenscreen_path)
             if not cap.isOpened():
+                logger.error("Could not open green screen video")
                 return None, "Could not open green screen video"
             
             fps = cap.get(cv2.CAP_PROP_FPS)
@@ -130,15 +135,21 @@ def _prog(pct: float, desc: str):
             if isinstance(background, str):
                 bg = cv2.imread(background)
                 if bg is None:
+                    logger.error("Could not load background image")
                     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))
+            # Setup output using app.py's create_video_writer
+            from app import create_video_writer
+            out, actual_output_path = create_video_writer(output_path, fps, width, height)
+            if out is None:
+                cap.release()
+                logger.error("Could not create output video file")
+                return None, "Could not create output video file"
+            output_path = actual_output_path
             
             frame_count = 0
             
@@ -163,13 +174,13 @@ def _prog(pct: float, desc: str):
             return output_path, f"Final video created: {frame_count} frames"
             
         except Exception as e:
-            logger.error(f"Stage 2 error: {e}")
+            logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
             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
+            logger.warning("SAM2 predictor not available, using 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
@@ -204,7 +215,7 @@ def _extract_person_mask(self, frame):
             return mask
             
         except Exception as e:
-            logger.error(f"Mask extraction error: {e}")
+            logger.error(f"Mask extraction error: {e}\n{traceback.format_exc()}")
             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
@@ -217,96 +228,111 @@ def _refine_mask(self, frame, mask):
             return mask
         
         try:
-            # MatAnyone refinement logic here
-            # This would depend on your MatAnyone implementation
-            return mask
-        except:
-            logger.warning("MatAnyone refinement failed, using original mask")
+            refined_mask = refine_mask_hq(frame, mask, self.matanyone_model)
+            logger.info("MatAnyone mask refinement successful")
+            return refined_mask
+        except Exception as e:
+            logger.warning(f"MatAnyone refinement failed: {e}\n{traceback.format_exc()}")
             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)
+        try:
+            # 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)
+        except Exception as e:
+            logger.error(f"Greenscreen application error: {e}\n{traceback.format_exc()}")
+            return frame
     
     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)
+        """Advanced chroma keying with spill suppression"""
+        try:
+            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)
+        except Exception as e:
+            logger.error(f"Chroma keying error: {e}\n{traceback.format_exc()}")
+            return frame
     
     def process_full_pipeline(self, video_path, background, final_output, 
                              chroma_settings=None, progress_callback=None):
         """
         Run the complete two-stage pipeline
         """
-        # 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
+            # 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:
+                logger.error(f"Stage 1 failed: {gs_msg}")
+                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 Exception as e:
+                logger.warning(f"Failed to clean up greenscreen file: {e}")
+            
+            if final_result is None:
+                logger.error(f"Stage 2 failed: {final_msg}")
+                return None, final_msg
+            
+            return final_result, final_msg
+        except Exception as e:
+            logger.error(f"Full pipeline error: {e}\n{traceback.format_exc()}")
+            return None, f"Full pipeline failed: {str(e)}"
 
-# Chroma key settings presets
+# Chroma key settings presets (aligned with app.py and ui_components.py)
 CHROMA_PRESETS = {
     'standard': {
         'key_color': [0, 255, 0],
@@ -314,13 +340,13 @@ def process_full_pipeline(self, video_path, background, final_output,
         'edge_softness': 2,
         'spill_suppression': 0.3
     },
-    'tight': {
+    'studio': {
         'key_color': [0, 255, 0],
         'tolerance': 30,
         'edge_softness': 1,
         'spill_suppression': 0.4
     },
-    'soft': {
+    'outdoor': {
         'key_color': [0, 255, 0],
         'tolerance': 50,
         'edge_softness': 3,