Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -3,353 +3,449 @@
 Two-Stage Green Screen Processing System
 Stage 1: Original → Green Screen
 Stage 2: Green Screen → Final Background
+
+This version is aligned with the current project structure:
+- Uses segment/refine helpers from utils.cv_processing
+- Has its own safe create_video_writer (no core.app dependency)
+- Supports cancel via stop_event
+- Robust SAM2 predictor handling
 """
 
+from __future__ import annotations
+
 import cv2
 import numpy as np
 import os
+import io
+import gc
 import pickle
 import logging
-from pathlib import Path
 import tempfile
 import traceback
-from utils.refinement.mask_refiner import refine_mask_hq  # Updated import path
+from pathlib import Path
+from typing import Optional, Dict, Any, Callable
+
+from utils.cv_processing import (
+    segment_person_hq,
+    refine_mask_hq,
+)
+
+try:
+    from utils.logger import get_logger
+    logger = get_logger(__name__)
+except Exception:
+    logger = logging.getLogger(__name__)
+
+
+# ---------------------------
+# Small local video I/O helper
+# ---------------------------
+def create_video_writer(
+    output_path: str,
+    fps: float,
+    width: int,
+    height: int,
+    prefer_mp4: bool = True,
+):
+    """
+    Create a cv2.VideoWriter with sane defaults.
+    Returns (writer, actual_output_path) or (None, output_path) on failure.
+    """
+    try:
+        ext = ".mp4" if prefer_mp4 else ".avi"
+        if not output_path:
+            output_path = tempfile.mktemp(suffix=ext)
+        else:
+            base, curr_ext = os.path.splitext(output_path)
+            if curr_ext.lower() not in [".mp4", ".avi", ".mov", ".mkv"]:
+                output_path = base + ext
+
+        # pick codec
+        # mp4v works widely on Spaces; if that fails, try XVID
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v") if prefer_mp4 else cv2.VideoWriter_fourcc(*"XVID")
+        writer = cv2.VideoWriter(output_path, fourcc, float(fps), (int(width), int(height)))
+        if not writer or not writer.isOpened():
+            # Fallback
+            alt_ext = ".avi" if prefer_mp4 else ".mp4"
+            alt_fourcc = cv2.VideoWriter_fourcc(*"XVID") if prefer_mp4 else cv2.VideoWriter_fourcc(*"mp4v")
+            alt_path = os.path.splitext(output_path)[0] + alt_ext
+            writer = cv2.VideoWriter(alt_path, alt_fourcc, float(fps), (int(width), int(height)))
+            if not writer or not writer.isOpened():
+                return None, output_path
+            return writer, alt_path
+
+        return writer, output_path
+    except Exception as e:
+        logger.error(f"create_video_writer failed: {e}")
+        return None, output_path
+
+
+# ---------------------------
+# Chroma key presets
+# ---------------------------
+CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {
+    'standard': {
+        'key_color': [0, 255, 0],     # pure green (BGR)
+        'tolerance': 38,              # color distance threshold
+        'edge_softness': 2,           # Gaussian kernel radius
+        'spill_suppression': 0.35,    # 0..1
+    },
+    'studio': {
+        'key_color': [0, 255, 0],
+        'tolerance': 30,
+        'edge_softness': 1,
+        'spill_suppression': 0.45,
+    },
+    'outdoor': {
+        'key_color': [0, 255, 0],
+        'tolerance': 50,
+        'edge_softness': 3,
+        'spill_suppression': 0.25,
+    },
+}
 
-logger = logging.getLogger(__name__)
 
 class TwoStageProcessor:
-    """Handle two-stage video processing with green screen intermediate"""
-    
+    """
+    Handle two-stage video processing with a green screen intermediate.
+    - Stage 1: generate clean green screen video (hard edges; great for chroma key)
+    - Stage 2: chroma-key that green to your final background
+    """
+
     def __init__(self, sam2_predictor=None, matanyone_model=None):
-        self.sam2_predictor = sam2_predictor
-        self.matanyone_model = matanyone_model
+        # We expect `sam2_predictor` to behave like SAM2ImagePredictor:
+        #   .set_image(np.ndarray)
+        #   .predict(point_coords=..., point_labels=..., multimask_output=True)
+        # If you passed a wrapper, we’ll try to unwrap it.
+        self.sam2 = self._unwrap_sam2(sam2_predictor)
+        self.matanyone = 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):
+
+        logger.info("TwoStageProcessor initialized. "
+                    f"SAM2 available: {self.sam2 is not None} | "
+                    f"MatAnyOne available: {self.matanyone is not None}")
+
+    # ---------------------------
+    # Stage 1: Original → Green
+    # ---------------------------
+    def stage1_extract_to_greenscreen(
+        self,
+        video_path: str,
+        output_path: str,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ):
         """
-        Stage 1: Extract person and create green screen video
-        Also saves masks for potential reuse
+        Extract foreground to a pure green background.
+        Saves per-frame masks (pickle) next to the output for optional reuse.
         """
         def _prog(pct: float, desc: str):
             if progress_callback:
-                progress_callback(pct, desc)
-        
+                try:
+                    progress_callback(float(pct), str(desc))
+                except Exception:
+                    pass
+
         try:
-            _prog(0.0, "Stage 1: Extracting to green screen...")
-            
+            _prog(0.0, "Stage 1: Preparing…")
             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)
-            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+                return None, "Could not open input video"
+
+            fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
             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 using core app's create_video_writer
-            from core.app import create_video_writer  # Updated import path
-            out, actual_output_path = create_video_writer(output_path, fps, width, height)
-            if out is None:
+
+            writer, output_path = create_video_writer(output_path, fps, width, height)
+            if writer 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 = []
-            frame_count = 0
-            
+                return None, "Could not create output writer"
+
+            green_bg = np.zeros((height, width, 3), dtype=np.uint8)
+            green_bg[:, :] = [0, 255, 0]  # BGR Pure Green
+
+            masks: list[np.ndarray] = []
+            frame_idx = 0
+
             while True:
-                ret, frame = cap.read()
-                if not ret:
+                if stop_event is not None and stop_event.is_set():
+                    _prog(1.0, "Stage 1: Cancelled")
+                    break
+
+                ok, frame = cap.read()
+                if not ok:
                     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)
+
+                # 1) get a mask (SAM2 w/ smart points via segment_person_hq)
+                mask = self._get_mask(frame)
+
+                # 2) refine occasionally with MatAnyOne to keep it light
+                if (self.matanyone is not None) and (frame_idx % 3 == 0):
+                    try:
+                        mask = refine_mask_hq(frame, mask, self.matanyone, fallback_enabled=True)
+                    except Exception as e:
+                        logger.warning(f"MatAnyOne refine failed (frame {frame_idx}): {e}")
+
                 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
-            
+
+                # 3) HARD-edge composite to green (no feather here)
+                green = self._apply_greenscreen_hard(frame, mask, green_bg)
+                writer.write(green)
+
+                frame_idx += 1
+                if total > 0:
+                    pct = 0.05 + 0.9 * (frame_idx / total)
+                else:
+                    pct = min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 1: {frame_idx}/{total or '?'} frames")
+
             cap.release()
-            out.release()
-            
-            # Save masks for potential reuse
-            mask_file = self.mask_cache_dir / f"{Path(output_path).stem}_masks.pkl"
+            writer.release()
+
+            # Save masks (best-effort)
             try:
-                with open(mask_file, 'wb') as f:
+                mask_file = self.mask_cache_dir / (Path(output_path).stem + "_masks.pkl")
+                with open(mask_file, "wb") as f:
                     pickle.dump(masks, f)
-                logger.info(f"Masks saved to {mask_file}")
+                logger.info(f"Stage 1: saved masks → {mask_file}")
             except Exception as e:
-                logger.warning(f"Failed to save masks: {e}")
-            
-            _prog(1.0, "Stage 1 complete: Green screen created")
-            return output_path, f"Green screen created: {frame_count} frames"
-            
+                logger.warning(f"Stage 1: failed to save masks: {e}")
+
+            _prog(1.0, "Stage 1: Complete")
+            return output_path, f"Green screen video created ({frame_idx} frames)"
+
         except Exception as 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, 
-                                   chroma_settings=None, progress_callback=None):
+            return None, f"Stage 1 failed: {e}"
+
+    # ---------------------------
+    # Stage 2: Green → Final BG
+    # ---------------------------
+    def stage2_greenscreen_to_final(
+        self,
+        greenscreen_path: str,
+        background: np.ndarray | str,
+        output_path: str,
+        chroma_settings: Optional[Dict[str, Any]] = None,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ):
         """
-        Stage 2: Replace green screen with final background using chroma keying
+        Replace green screen with the given background using chroma keying.
+        `background` may be a path or an already-loaded image (BGR).
         """
         def _prog(pct: float, desc: str):
             if progress_callback:
-                progress_callback(pct, desc)
-        
-        if chroma_settings is None:
-            chroma_settings = CHROMA_PRESETS['standard']
-        
+                try:
+                    progress_callback(float(pct), str(desc))
+                except Exception:
+                    pass
+
         try:
-            _prog(0.0, "Stage 2: Applying final background...")
-            
+            _prog(0.0, "Stage 2: Preparing…")
             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)
-            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+            fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
             width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            
-            # Prepare background
+
+            writer, output_path = create_video_writer(output_path, fps, width, height)
+            if writer is None:
+                cap.release()
+                return None, "Could not create output writer"
+
+            # Load/resize background
             if isinstance(background, str):
-                bg = cv2.imread(background)
+                bg = cv2.imread(background, cv2.IMREAD_COLOR)
                 if bg is None:
-                    logger.error("Could not load background image")
+                    cap.release()
+                    writer.release()
                     return None, "Could not load background image"
             else:
                 bg = background
-            
-            bg = cv2.resize(bg, (width, height))
-            
-            # Setup output using core app's create_video_writer
-            from core.app import create_video_writer  # Updated import path
-            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
-            
+            bg = cv2.resize(bg, (width, height), interpolation=cv2.INTER_LANCZOS4)
+
+            settings = dict(CHROMA_PRESETS.get('standard', {}))
+            if chroma_settings:
+                settings.update(chroma_settings)
+
+            frame_idx = 0
+
             while True:
-                ret, frame = cap.read()
-                if not ret:
+                if stop_event is not None and stop_event.is_set():
+                    _prog(1.0, "Stage 2: Cancelled")
                     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
-            
+
+                ok, frame = cap.read()
+                if not ok:
+                    break
+
+                out = self._chroma_key_advanced(frame, bg, settings)
+                writer.write(out)
+
+                frame_idx += 1
+                if total > 0:
+                    pct = 0.05 + 0.9 * (frame_idx / total)
+                else:
+                    pct = min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 2: {frame_idx}/{total or '?'} frames")
+
             cap.release()
-            out.release()
-            
-            _prog(1.0, "Stage 2 complete: Final video created")
-            return output_path, f"Final video created: {frame_count} frames"
-            
+            writer.release()
+            _prog(1.0, "Stage 2: Complete")
+
+            return output_path, f"Final video created ({frame_idx} frames)"
+
         except Exception as 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:
-            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
-            return mask
-        
+            return None, f"Stage 2 failed: {e}"
+
+    # ---------------------------
+    # Full pipeline
+    # ---------------------------
+    def process_full_pipeline(
+        self,
+        video_path: str,
+        background: np.ndarray | str,
+        final_output: str,
+        chroma_settings: Optional[Dict[str, Any]] = None,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ):
+        """
+        Stage 1 (to temp greenscreen) → Stage 2 (final composite).
+        """
+        gs_temp = tempfile.mktemp(suffix="_greenscreen.mp4")
         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
+            gs_path, msg1 = self.stage1_extract_to_greenscreen(
+                video_path, gs_temp, progress_callback=progress_callback, stop_event=stop_event
             )
-            
-            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}\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
-            return mask
-    
-    def _refine_mask(self, frame, mask):
-        """Refine mask using MatAnyone if available"""
-        if self.matanyone_model is None:
-            logger.warning("MatAnyone not available, using original mask")
-            return mask
-        
+            if gs_path is None:
+                return None, msg1
+
+            result, msg2 = self.stage2_greenscreen_to_final(
+                gs_path, background, final_output,
+                chroma_settings=chroma_settings,
+                progress_callback=progress_callback,
+                stop_event=stop_event
+            )
+            if result is None:
+                return None, msg2
+
+            return result, msg2
+
+        finally:
+            # best-effort cleanup
+            try:
+                if os.path.exists(gs_temp):
+                    os.remove(gs_temp)
+            except Exception:
+                pass
+            gc.collect()
+
+    # ---------------------------
+    # Internals
+    # ---------------------------
+    def _unwrap_sam2(self, obj):
+        """
+        Try to get a callable SAM2-like predictor from whatever was passed.
+        Accepts:
+          - direct predictor (has set_image + predict)
+          - wrapper with .model that has set_image + predict
+          - wrapper with .predictor
+        """
         try:
-            refined_mask = refine_mask_hq(frame, mask, self.matanyone_model)
-            logger.info("MatAnyone mask refinement successful")
-            return refined_mask
+            if obj is None:
+                return None
+            # predictor directly?
+            if hasattr(obj, "set_image") and hasattr(obj, "predict"):
+                return obj
+            # wrapper.model?
+            model = getattr(obj, "model", None)
+            if model is not None and hasattr(model, "set_image") and hasattr(model, "predict"):
+                return model
+            # wrapper.predictor?
+            predictor = getattr(obj, "predictor", None)
+            if predictor is not None and hasattr(predictor, "set_image") and hasattr(predictor, "predict"):
+                return predictor
         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"""
+            logger.warning(f"SAM2 unwrap failed: {e}")
+        return None
+
+    def _get_mask(self, frame: np.ndarray) -> np.ndarray:
+        """
+        Use our project’s enhanced segmentation helper so validation/fallbacks are consistent.
+        """
+        predictor = self.sam2
         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)
+            mask = segment_person_hq(frame, predictor, fallback_enabled=True)
+            return mask
         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"""
+            logger.warning(f"Segmentation failed, using geometric fallback: {e}")
+            h, w = frame.shape[:2]
+            m = np.zeros((h, w), dtype=np.uint8)
+            m[h//6:5*h//6, w//4:3*w//4] = 255
+            return m
+
+    def _apply_greenscreen_hard(self, frame: np.ndarray, mask: np.ndarray, green_bg: np.ndarray) -> np.ndarray:
+        """
+        Hard-edge composite to pure green for very clean keying later.
+        """
         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)
+            if mask.ndim == 3:
+                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+            if mask.dtype != np.uint8:
+                mask = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
+
+            _, mask_bin = cv2.threshold(mask, 140, 255, cv2.THRESH_BINARY)
+            mask3 = cv2.cvtColor(mask_bin, cv2.COLOR_GRAY2BGR).astype(np.float32) / 255.0
+
+            out = frame.astype(np.float32) * mask3 + green_bg.astype(np.float32) * (1.0 - mask3)
+            return np.clip(out, 0, 255).astype(np.uint8)
         except Exception as e:
-            logger.error(f"Chroma keying error: {e}\n{traceback.format_exc()}")
+            logger.error(f"Greenscreen composite failed: {e}")
             return frame
-    
-    def process_full_pipeline(self, video_path, background, final_output, 
-                             chroma_settings=None, progress_callback=None):
+
+    def _chroma_key_advanced(self, frame_bgr: np.ndarray, bg_bgr: np.ndarray, settings: Dict[str, Any]) -> np.ndarray:
         """
-        Run the complete two-stage pipeline
+        Distance-to-key color mask + soft edge + spill suppression (green reduction).
         """
         try:
-            # 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)}"
+            key = np.array(settings.get("key_color", [0, 255, 0]), dtype=np.float32)
+            tol = float(settings.get("tolerance", 40))
+            soft = int(settings.get("edge_softness", 2))
+            spill = float(settings.get("spill_suppression", 0.3))
 
-# Chroma key settings presets (aligned with app.py and ui_components.py)
-CHROMA_PRESETS = {
-    'standard': {
-        'key_color': [0, 255, 0],
-        'tolerance': 40,
-        'edge_softness': 2,
-        'spill_suppression': 0.3
-    },
-    'studio': {
-        'key_color': [0, 255, 0],
-        'tolerance': 30,
-        'edge_softness': 1,
-        'spill_suppression': 0.4
-    },
-    'outdoor': {
-        'key_color': [0, 255, 0],
-        'tolerance': 50,
-        'edge_softness': 3,
-        'spill_suppression': 0.2
-    }
-}
+            f = frame_bgr.astype(np.float32)
+            b = bg_bgr.astype(np.float32)
+
+            # distance (BGR space)
+            diff = f - key
+            dist = np.sqrt((diff ** 2).sum(axis=2))
+
+            # inside green → 0, far from green → 1
+            mask = np.clip((dist - tol) / max(tol, 1.0), 0.0, 1.0)
+
+            if soft > 0:
+                ksize = max(1, soft * 2 + 1)
+                mask = cv2.GaussianBlur(mask.astype(np.float32), (ksize, ksize), soft)
+
+            # spill suppression
+            if spill > 0:
+                # where mask < 1.0 (near edges), reduce green channel proportionally
+                spill_zone = 1.0 - mask
+                g = f[:, :, 1]
+                f[:, :, 1] = np.clip(g - g * spill_zone * spill, 0, 255)
+
+            mask3 = np.stack([mask] * 3, axis=2)
+            out = f * mask3 + b * (1.0 - mask3)
+            return np.clip(out, 0, 255).astype(np.uint8)
+        except Exception as e:
+            logger.error(f"Chroma keying failed: {e}")
+            return frame_bgr