Update core/app.py

core/app.py

@@ -1,587 +1,457 @@
 #!/usr/bin/env python3
 """
-BackgroundFX Pro – Main Application Entry Point
-Refactored modular architecture – orchestrates specialised components
-
-2025-08-27 update:
-- Robust Two-Stage importer: tries package import, then direct file import (bypasses
-  any side-effect errors in processing/__init__.py). Clear logging of where it loaded.
-- Config hardening: adds safe defaults for fields like max_model_size/use_nvenc/etc.
-- Defensive error logs with tracebacks for easier diagnosis.
+Two-Stage Green-Screen Processing System  ✅ 2025-08-26
+Stage 1: Original → keyed background (auto-selected colour)
+Stage 2: Keyed video → final composite  (hybrid chroma + segmentation rescue)
+
+Aligned with current project layout:
+  * uses helpers from utils.cv_processing (segment_person_hq, refine_mask_hq)
+  * safe local create_video_writer (no core.app dependency)
+  * cancel support via stop_event
+  * progress_callback(pct, desc)
+  * fully self-contained – just drop in and import TwoStageProcessor
 """
 
 from __future__ import annotations
 
-# ─────────────────────────────────────────────────────────────────────────────
-# 0) Early env/threading hygiene (must run first)
-# ─────────────────────────────────────────────────────────────────────────────
-import early_env  # sets OMP/MKL/OPENBLAS + torch threads safely
-
-import logging
-import threading
-import traceback
-import sys
-import os
-import importlib
-import importlib.util
+import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
 from pathlib import Path
-from typing import Optional, Tuple, Dict, Any, Callable
-
-# ─────────────────────────────────────────────────────────────────────────────
-# 1) Logging
-# ─────────────────────────────────────────────────────────────────────────────
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger("core.app")
-
-# Ensure project root is importable (helps HF Spaces and local runs)
-PROJECT_FILE = Path(__file__).resolve()
-CORE_DIR = PROJECT_FILE.parent
-ROOT = CORE_DIR.parent
-if str(ROOT) not in sys.path:
-    sys.path.insert(0, str(ROOT))
-
-# ─────────────────────────────────────────────────────────────────────────────
-# 2) Patch Gradio schema early (HF Spaces quirk)
-# ─────────────────────────────────────────────────────────────────────────────
-try:
-    import gradio_client.utils as gc_utils
-
-    _orig_get_type = gc_utils.get_type
-
-    def _patched_get_type(schema):
-        if not isinstance(schema, dict):
-            if isinstance(schema, bool):
-                return "boolean"
-            if isinstance(schema, str):
-                return "string"
-            if isinstance(schema, (int, float)):
-                return "number"
-            return "string"
-        return _orig_get_type(schema)
-
-    gc_utils.get_type = _patched_get_type
-    logger.info("Gradio schema patch applied")
-except Exception as e:
-    logger.warning(f"Gradio patch failed: {e}")
-
-# ─────────────────────────────────────────────────────────────────────────────
-# 3) Core config + components
-# ─────────────────────────────────────────────────────────────────────────────
-from config.app_config import get_config
-from core.exceptions import ModelLoadingError, VideoProcessingError
-from utils.hardware.device_manager import DeviceManager
-from utils.system.memory_manager import MemoryManager
-from models.loaders.model_loader import ModelLoader
-from processing.video.video_processor import CoreVideoProcessor
-from processing.audio.audio_processor import AudioProcessor
-from utils.monitoring.progress_tracker import ProgressTracker
-from utils.cv_processing import validate_video_file
-
-# ─────────────────────────────────────────────────────────────────────────────
-# 3.1) Optional: Two-stage import with package and file fallbacks
-# ─────────────────────────────────────────────────────────────────────────────
-def _import_two_stage() -> tuple[bool, Any, Dict[str, Any], str, str]:
-    """
-    Returns (available, TwoStageProcessor|None, CHROMA_PRESETS|{}, import_origin, error_text)
-
-    Tries:
-      1) package imports (preferred)
-      2) direct file imports (bypass processing/__init__.py side-effects)
-    """
-    pkg_paths = [
-        "processing.two_stage.two_stage_processor",
-        "two_stage_processor",
-        "processing.two_stage_processor",
-    ]
-    fs_paths = [
-        ROOT / "processing" / "two_stage" / "two_stage_processor.py",
-        ROOT / "processing" / "two_stage_processor.py",
-        ROOT / "two_stage_processor.py",
-    ]
-
-    # Package imports
-    last_err_text = ""
-    for mod_path in pkg_paths:
-        try:
-            mod = importlib.import_module(mod_path)
-            TSP = getattr(mod, "TwoStageProcessor")
-            PRESETS = getattr(mod, "CHROMA_PRESETS", {"standard": {}})
-            logger.info(f"Two-stage import OK from package '{mod_path}'")
-            return True, TSP, PRESETS, f"pkg:{mod_path}", ""
-        except Exception as e:
-            tb = traceback.format_exc()
-            last_err_text += f"[pkg:{mod_path}] {repr(e)}\n{tb}\n"
-
-    # Direct file imports (bypass __init__.py)
-    for path in fs_paths:
-        try:
-            if not path.exists():
-                continue
-            spec = importlib.util.spec_from_file_location("bx_two_stage", str(path))
-            if not spec or not spec.loader:
-                continue
-            mod = importlib.util.module_from_spec(spec)
-            sys.modules["bx_two_stage"] = mod
-            spec.loader.exec_module(mod)  # type: ignore[attr-defined]
-            TSP = getattr(mod, "TwoStageProcessor")
-            PRESETS = getattr(mod, "CHROMA_PRESETS", {"standard": {}})
-            logger.info(f"Two-stage import OK from file '{path}'")
-            return True, TSP, PRESETS, f"file:{path}", ""
-        except Exception as e:
-            tb = traceback.format_exc()
-            last_err_text += f"[file:{path}] {repr(e)}\n{tb}\n"
+from typing import Optional, Dict, Any, Callable, Tuple, List
 
-    logger.error("Two-stage import failed from all paths:\n%s", last_err_text or "(no traceback)")
-    return False, None, {"standard": {}}, "", last_err_text or "(no traceback)"
+from utils.cv_processing import segment_person_hq, refine_mask_hq
 
-TWO_STAGE_AVAILABLE, _TwoStageProcessor, CHROMA_PRESETS, TWO_STAGE_IMPORT_ORIGIN, TWO_STAGE_IMPORT_ERROR = _import_two_stage()
-
-# ╔══════════════════════════════════════════════════════════════════════════╗
-# ║                           VideoProcessor class                           ║
-# ╚══════════════════════════════════════════════════════════════════════════╝
-class VideoProcessor:
-    """
-    Main orchestrator – coordinates all specialised components.
-    """
+# Project logger if available
+try:
+    from utils.logger import get_logger
+    logger = get_logger(__name__)
+except Exception:
+    logger = logging.getLogger(__name__)
 
-    def __init__(self):
-        self.config = get_config()
-        self._patch_config_defaults(self.config)  # ensure missing attrs exist
 
-        self.device_manager = DeviceManager()
-        self.memory_manager = MemoryManager(self.device_manager.get_optimal_device())
-        self.model_loader = ModelLoader(self.device_manager, self.memory_manager)
+# ---------------------------------------------------------------------------
+# Local video-writer helper
+# ---------------------------------------------------------------------------
+def create_video_writer(output_path: str, fps: float, width: int, height: int, prefer_mp4: bool = True):
+    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
+
+        fourcc = cv2.VideoWriter_fourcc(*("mp4v" if prefer_mp4 else "XVID"))
+        writer = cv2.VideoWriter(output_path, fourcc, float(fps), (int(width), int(height)))
+        if writer is None or not writer.isOpened():
+            alt_ext = ".avi" if prefer_mp4 else ".mp4"
+            alt_fourcc = cv2.VideoWriter_fourcc(*("XVID" if prefer_mp4 else "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 writer is None 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
+
+
+# ---------------------------------------------------------------------------
+# Key-colour helpers (fast, no external deps)
+# ---------------------------------------------------------------------------
+def _bgr_to_hsv_hue_deg(bgr: np.ndarray) -> np.ndarray:
+    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
+    # OpenCV H is 0-180; scale to degrees 0-360
+    return hsv[..., 0].astype(np.float32) * 2.0
+
+
+def _hue_distance(a_deg: float, b_deg: float) -> float:
+    """Circular distance on the hue wheel (degrees)."""
+    d = abs(a_deg - b_deg) % 360.0
+    return min(d, 360.0 - d)
+
+
+def _key_candidates_bgr() -> dict:
+    return {
+        "green":   {"bgr": np.array([  0,255,  0], dtype=np.uint8), "hue": 120.0},
+        "blue":    {"bgr": np.array([255,  0,  0], dtype=np.uint8), "hue": 240.0},
+        "cyan":    {"bgr": np.array([255,255,  0], dtype=np.uint8), "hue": 180.0},
+        "magenta": {"bgr": np.array([255,  0,255], dtype=np.uint8), "hue": 300.0},
+    }
+
+
+def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dict:
+    """Pick the candidate colour farthest from the actor’s dominant hues."""
+    try:
+        fg = frame_bgr[mask_uint8 > 127]
+        if fg.size < 1_000:
+            return _key_candidates_bgr()["green"]
+
+        fg_hue = _bgr_to_hsv_hue_deg(fg.reshape(-1, 1, 3)).reshape(-1)
+        hist, edges = np.histogram(fg_hue, bins=36, range=(0.0, 360.0))
+        top_idx = np.argsort(hist)[-3:]
+        top_hues = [(edges[i] + edges[i+1]) * 0.5 for i in top_idx]
+
+        best_name, best_score = None, -1.0
+        for name, info in _key_candidates_bgr().items():
+            cand_hue = info["hue"]
+            score = min(abs((cand_hue - th + 180) % 360 - 180) for th in top_hues)
+            if score > best_score:
+                best_name, best_score = name, score
+        return _key_candidates_bgr().get(best_name, _key_candidates_bgr()["green"])
+    except Exception:
+        return _key_candidates_bgr()["green"]
+
+
+# ---------------------------------------------------------------------------
+# Chroma presets
+# ---------------------------------------------------------------------------
+CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {
+    'standard': {'key_color': [0,255,0], 'tolerance': 38, 'edge_softness': 2, 'spill_suppression': 0.35},
+    '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},
+}
+
+
+# ---------------------------------------------------------------------------
+# Two-Stage Processor
+# ---------------------------------------------------------------------------
+class TwoStageProcessor:
+    def __init__(self, sam2_predictor=None, matanyone_model=None):
+        self.sam2  = self._unwrap_sam2(sam2_predictor)
+        self.matanyone = matanyone_model
+        self.mask_cache_dir = Path("/tmp/mask_cache")
+        self.mask_cache_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
+
+    # ---------------------------------------------------------------------
+    # Stage 1 – Original → keyed (green/blue/…)  -- chooses colour on 1st frame
+    # ---------------------------------------------------------------------
+    def stage1_extract_to_greenscreen(
+        self,
+        video_path: str,
+        output_path: str,
+        *,
+        key_color_mode: str = "auto",           # "auto" | "green" | "blue" | "cyan" | "magenta"
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
+    ) -> Tuple[Optional[dict], str]:
+
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
 
-        self.audio_processor = AudioProcessor()
-        self.core_processor: CoreVideoProcessor | None = None
-        self.two_stage_processor: Any | None = None  # instance of TwoStageProcessor if available
+        try:
+            _prog(0.0, "Stage 1: opening video…")
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                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
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            writer, out_path = create_video_writer(output_path, fps, w, h)
+            if writer is None:
+                cap.release()
+                return None, "Could not create output writer"
+
+            key_info: dict | None = None
+            chosen_bgr = np.array([0, 255, 0], np.uint8)  # default
+            probe_done = False
+            masks: List[np.ndarray] = []
+            frame_idx = 0
+
+            green_bg_template = np.zeros((h, w, 3), np.uint8)  # overwritten per-frame
+
+            while True:
+                if stop_event and stop_event.is_set():
+                    _prog(1.0, "Stage 1: cancelled")
+                    break
+
+                ok, frame = cap.read()
+                if not ok:
+                    break
+
+                mask = self._get_mask(frame)
+
+                # decide key colour once
+                if not probe_done:
+                    if key_color_mode.lower() == "auto":
+                        key_info = _choose_best_key_color(frame, mask)
+                        chosen_bgr = key_info["bgr"]
+                    else:
+                        cand = _key_candidates_bgr().get(key_color_mode.lower())
+                        if cand is not None:
+                            chosen_bgr = cand["bgr"]
+                    probe_done = True
+                    logger.info(f"[TwoStage] Using key colour: {key_color_mode} → {chosen_bgr.tolist()}")
+
+                # optional refine
+                if self.matanyone 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 fail f={frame_idx}: {e}")
 
-        self.models_loaded = False
-        self.loading_lock = threading.Lock()
-        self.cancel_event = threading.Event()
-        self.progress_tracker: ProgressTracker | None = None
+                # composite
+                green_bg_template[:] = chosen_bgr
+                gs = self._apply_greenscreen_hard(frame, mask, green_bg_template)
+                writer.write(gs)
+                masks.append(self._to_binary_mask(mask))
 
-        logger.info(f"VideoProcessor on device: {self.device_manager.get_optimal_device()}")
+                frame_idx += 1
+                pct = 0.05 + 0.9 * (frame_idx / total) if total else min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 1: {frame_idx}/{total or '?'}")
 
-    # ─────────────────────────────────────────────────────────────────────
-    # Config hardening – add missing attributes to avoid AttributeError
-    # ─────────────────────────────────────────────────────────────────────
-    @staticmethod
-    def _patch_config_defaults(cfg: Any) -> None:
-        """
-        Some downstream modules may read fields that older configs don't define.
-        Add safe defaults here so AttributeErrors never occur.
-        """
-        defaults = {
-            # video i/o & writer
-            "use_nvenc": False,
-            "prefer_mp4": True,
-            "video_codec": "mp4v",
-            "audio_copy": True,
-            "ffmpeg_path": "ffmpeg",
-            # model/resource guards
-            "max_model_size": 0,            # bytes or MB; downstream should treat 0/None as "no limit"
-            "max_model_size_bytes": 0,
-            # housekeeping
-            "output_dir": str(ROOT / "outputs"),
-        }
-        for k, v in defaults.items():
-            if not hasattr(cfg, k):
-                setattr(cfg, k, v)
-
-        # Ensure output dir exists
-        try:
-            Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
-        except Exception:
-            out = ROOT / "outputs"
-            out.mkdir(parents=True, exist_ok=True)
-            cfg.output_dir = str(out)
-
-    # ─────────────────────────────────────────────────────────────────────
-    # Progress helper
-    # ─────────────────────────────────────────────────────────────────────
-    def _init_progress(self, video_path: str, cb: Optional[Callable] = None):
-        try:
-            import cv2
-            cap = cv2.VideoCapture(video_path)
-            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             cap.release()
-            if total <= 0:
-                total = 100
-            self.progress_tracker = ProgressTracker(total, cb)
-        except Exception as e:
-            logger.warning(f"Progress init failed: {e}")
-            self.progress_tracker = ProgressTracker(100, cb)
-
-    # ─────────────────────────────────────────────────────────────────────
-    # Model loading
-    # ─────────────────────────────────────────────────────────────────────
-    def load_models(self, progress_callback: Optional[Callable] = None) -> str:
-        with self.loading_lock:
-            if self.models_loaded:
-                return "Models already loaded and validated"
+            writer.release()
 
+            # save mask cache
             try:
-                self.cancel_event.clear()
-                if progress_callback:
-                    progress_callback(
-                        0.0, f"Loading on {self.device_manager.get_optimal_device()}"
-                    )
-
-                sam2_loaded, mat_loaded = self.model_loader.load_all_models(
-                    progress_callback=progress_callback, cancel_event=self.cancel_event
-                )
-
-                if self.cancel_event.is_set():
-                    return "Model loading cancelled"
-
-                # Unwrap actual predictor / model objects
-                sam2_predictor = getattr(sam2_loaded, "model", None) if sam2_loaded else None
-                mat_model = getattr(mat_loaded, "model", None) if mat_loaded else None
-
-                # Core single-stage processor
-                self.core_processor = CoreVideoProcessor(
-                    config=self.config, models=self.model_loader
-                )
-
-                # Two-stage processor (optional)
-                self.two_stage_processor = None
-                if TWO_STAGE_AVAILABLE and (_TwoStageProcessor is not None) and (sam2_predictor or mat_model):
-                    try:
-                        self.two_stage_processor = _TwoStageProcessor(
-                            sam2_predictor=sam2_predictor, matanyone_model=mat_model
-                        )
-                        logger.info("Two-stage processor initialised (%s)", TWO_STAGE_IMPORT_ORIGIN or "unknown")
-                    except Exception as e:
-                        logger.warning("Two-stage init failed: %r\n%s", e, traceback.format_exc())
-                        self.two_stage_processor = None
-
-                self.models_loaded = True
-                msg = self.model_loader.get_load_summary()
-                msg += (
-                    "\n✅ Two-stage processor ready"
-                    if self.two_stage_processor
-                    else "\n⚠️ Two-stage processor not available"
-                )
-                logger.info(msg)
-                return msg
-
-            except (AttributeError, ModelLoadingError) as e:
-                self.models_loaded = False
-                err = f"Model loading failed: {e}"
-                logger.error(err)
-                return err
+                cache_file = self.mask_cache_dir / (Path(out_path).stem + "_masks.pkl")
+                with open(cache_file, "wb") as f:
+                    pickle.dump(masks, f)
             except Exception as e:
-                self.models_loaded = False
-                err = f"Unexpected error during model loading: {e}"
-                logger.error(err)
-                logger.debug("Traceback:\n%s", traceback.format_exc())
-                return err
-
-    # ─────────────────────────────────────────────────────────────────────
-    # Public entry – process video
-    # ─────────────────────────────────────────────────────────────────────
-    def process_video(
+                logger.warning(f"mask cache save fail: {e}")
+
+            _prog(1.0, "Stage 1: complete")
+            return (
+                {"path": out_path, "frames": frame_idx, "key_bgr": chosen_bgr.tolist()},
+                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: {e}"
+
+    # ---------------------------------------------------------------------
+    # Stage 2 – keyed video → final composite  (hybrid matte)
+    # ---------------------------------------------------------------------
+    def stage2_greenscreen_to_final(
         self,
-        video_path: str,
-        background_choice: str,
-        custom_background_path: Optional[str] = None,
-        progress_callback: Optional[Callable] = None,
-        use_two_stage: bool = False,
-        chroma_preset: str = "standard",
-        key_color_mode: str = "auto",
-        preview_mask: bool = False,
-        preview_greenscreen: bool = False,
+        gs_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,
     ) -> Tuple[Optional[str], str]:
-        if not self.models_loaded or not self.core_processor:
-            return None, "Models not loaded. Please click “Load Models” first."
-        if self.cancel_event.is_set():
-            return None, "Processing cancelled"
-
-        self._init_progress(video_path, progress_callback)
 
-        ok, why = validate_video_file(video_path)
-        if not ok:
-            return None, f"Invalid video: {why}"
+        def _prog(p, d):
+            if progress_callback:
+                try:
+                    progress_callback(float(p), str(d))
+                except Exception:
+                    pass
 
         try:
-            if use_two_stage:
-                if not TWO_STAGE_AVAILABLE or self.two_stage_processor is None:
-                    return None, "Two-stage processing not available on this build"
-                return self._process_two_stage(
-                    video_path,
-                    background_choice,
-                    custom_background_path,
-                    progress_callback,
-                    chroma_preset,
-                    key_color_mode,
-                )
+            _prog(0.0, "Stage 2: opening keyed video…")
+            cap = cv2.VideoCapture(gs_path)
+            if not cap.isOpened():
+                return None, "Could not open keyed video"
+
+            fps   = cap.get(cv2.CAP_PROP_FPS) or 25.0
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
+            w     = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            h     = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            writer, out_path = create_video_writer(output_path, fps, w, h)
+            if writer is None:
+                cap.release()
+                return None, "Could not create output writer"
+
+            # background
+            if isinstance(background, str):
+                bg = cv2.imread(background, cv2.IMREAD_COLOR)
+                if bg is None:
+                    cap.release()
+                    writer.release()
+                    return None, "Could not load background"
             else:
-                return self._process_single_stage(
-                    video_path,
-                    background_choice,
-                    custom_background_path,
-                    progress_callback,
-                    preview_mask,
-                    preview_greenscreen,
-                )
-
-        except VideoProcessingError as e:
-            logger.error(f"Processing failed: {e}")
-            return None, f"Processing failed: {e}"
+                bg = background
+            bg = cv2.resize(bg, (w, h), interpolation=cv2.INTER_LANCZOS4).astype(np.uint8)
+
+            # settings
+            settings = dict(CHROMA_PRESETS['standard'])
+            if chroma_settings:
+                settings.update(chroma_settings)
+
+            # load cached masks if any
+            cache_file = self.mask_cache_dir / (Path(gs_path).stem + "_masks.pkl")
+            cached_masks = None
+            if cache_file.exists():
+                try:
+                    with open(cache_file, 'rb') as f:
+                        cached_masks = pickle.load(f)
+                except Exception as e:
+                    logger.warning(f"mask cache load fail: {e}")
+
+            frame_idx = 0
+            while True:
+                if stop_event and stop_event.is_set():
+                    _prog(1.0, "Stage 2: cancelled")
+                    break
+                ok, frame = cap.read()
+                if not ok:
+                    break
+
+                if cached_masks and frame_idx < len(cached_masks):
+                    seg_mask = cached_masks[frame_idx]
+                else:
+                    seg_mask = self._segmentation_mask_on_stage2(frame)
+
+                composite = self._chroma_key_advanced(frame, bg, settings, seg_mask)
+
+                writer.write(composite)
+                frame_idx += 1
+                pct = 0.05 + 0.9 * (frame_idx / total) if total else min(0.95, 0.05 + frame_idx * 0.002)
+                _prog(pct, f"Stage 2: {frame_idx}/{total or '?'}")
+
+            cap.release()
+            writer.release()
+            _prog(1.0, "Stage 2: complete")
+            return out_path, f"Final video created ({frame_idx} frames)"
         except Exception as e:
-            logger.error(f"Unexpected processing error: {e}")
-            logger.debug("Traceback:\n%s", traceback.format_exc())
-            return None, f"Unexpected error: {e}"
-
-    # ─────────────────────────────────────────────────────��───────────────
-    # Private – single-stage
-    # ─────────────────────────────────────────────────────────────────────
-    def _process_single_stage(
-        self,
-        video_path: str,
-        background_choice: str,
-        custom_background_path: Optional[str],
-        progress_callback: Optional[Callable],
-        preview_mask: bool,
-        preview_greenscreen: bool,
-    ) -> Tuple[Optional[str], str]:
-        import time
-
-        ts = int(time.time())
-        out_dir = Path(self.config.output_dir) / "single_stage"
-        out_dir.mkdir(parents=True, exist_ok=True)
-        out_path = str(out_dir / f"processed_{ts}.mp4")
-
-        result = self.core_processor.process_video(
-            input_path=video_path,
-            output_path=out_path,
-            bg_config={
-                "background_choice": background_choice,
-                "custom_path": custom_background_path,
-            },
-        )
-        if not result:
-            return None, "Video processing failed"
-
-        # Mux original audio back unless preview flags request otherwise
-        if not (preview_mask or preview_greenscreen):
-            try:
-                final_path = self.audio_processor.add_audio_to_video(
-                    original_video=video_path, processed_video=out_path
-                )
-            except Exception as e:
-                logger.warning("Audio mux failed, returning video without audio: %r", e)
-                final_path = out_path
-        else:
-            final_path = out_path
-
-        msg = (
-            "Processing completed.\n"
-            f"Frames: {result.get('frames', 'unknown')}\n"
-            f"Background: {background_choice}\n"
-            f"Mode: Single-stage\n"
-            f"Device: {self.device_manager.get_optimal_device()}"
-        )
-        return final_path, msg
-
-    # ─────────────────────────────────────────────────────────────────────
-    # Private – two-stage
-    # ─────────────────────────────────────────────────────────────────────
-    def _process_two_stage(
+            logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
+            return None, f"Stage 2 failed: {e}"
+
+    # ---------------------------------------------------------------------
+    # Full pipeline – now passes chosen key into Stage 2
+    # ---------------------------------------------------------------------
+    def process_full_pipeline(
         self,
         video_path: str,
-        background_choice: str,
-        custom_background_path: Optional[str],
-        progress_callback: Optional[Callable],
-        chroma_preset: str,
-        key_color_mode: str,
+        background: np.ndarray | str,
+        final_output: str,
+        *,
+        key_color_mode: str = "auto",
+        chroma_settings: Optional[Dict[str, Any]] = None,
+        progress_callback: Optional[Callable[[float, str], None]] = None,
+        stop_event: Optional["threading.Event"] = None,
     ) -> Tuple[Optional[str], str]:
-        if self.two_stage_processor is None:
-            return None, "Two-stage processor not available"
-
-        import cv2, time
+        gs_tmp = tempfile.mktemp(suffix="_gs.mp4")
+        try:
+            gs_info, msg1 = self.stage1_extract_to_greenscreen(
+                video_path, gs_tmp,
+                key_color_mode=key_color_mode,
+                progress_callback=progress_callback, stop_event=stop_event
+            )
+            if gs_info is None:
+                return None, msg1
 
-        # Determine output geometry from source
-        cap = cv2.VideoCapture(video_path)
-        if not cap.isOpened():
-            return None, "Could not open input video"
-        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) or 1280
-        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) or 720
-        cap.release()
+            # inject key colour into chroma settings for Stage 2
+            chosen_key = gs_info.get("key_bgr", [0, 255, 0])
+            cs = dict(chroma_settings or CHROMA_PRESETS['standard'])
+            cs['key_color'] = chosen_key
 
-        # Prepare background
-        try:
-            background = self.core_processor.prepare_background(
-                background_choice, custom_background_path, w, h
+            result, msg2 = self.stage2_greenscreen_to_final(
+                gs_info["path"], background, final_output,
+                chroma_settings=cs, progress_callback=progress_callback, stop_event=stop_event
             )
-        except Exception as e:
-            logger.error("Background preparation failed: %r", e)
-            return None, f"Failed to prepare background: {e}"
-
-        if background is None:
-            return None, "Failed to prepare background"
-
-        ts = int(time.time())
-        out_dir = Path(self.config.output_dir) / "two_stage"
-        out_dir.mkdir(parents=True, exist_ok=True)
-        final_out = str(out_dir / f"final_{ts}.mp4")
-
-        chroma_cfg = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS.get("standard", {}))
-        logger.info(
-            "Two-stage with preset: %s | key_color_mode=%s | origin=%s",
-            chroma_preset, key_color_mode, TWO_STAGE_IMPORT_ORIGIN or "unknown"
-        )
-
-        result, message = self.two_stage_processor.process_full_pipeline(
-            video_path,
-            background,
-            final_out,
-            key_color_mode=key_color_mode,
-            chroma_settings=chroma_cfg,
-            progress_callback=progress_callback,
-        )
-        if result is None:
-            return None, message
-
-        # Mux audio from original (same logic as single-stage)
+            return result, msg2
+        finally:
+            try:
+                os.remove(gs_tmp)
+            except Exception:
+                pass
+            gc.collect()
+
+    # ---------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------
+    def _unwrap_sam2(self, obj):
         try:
-            final_path = self.audio_processor.add_audio_to_video(
-                original_video=video_path, processed_video=result
-            )
+            if obj is None:
+                return None
+            if all(hasattr(obj, attr) for attr in ("set_image", "predict")):
+                return obj
+            for attr in ("model", "predictor"):
+                inner = getattr(obj, attr, None)
+                if inner and all(hasattr(inner, a) for a in ("set_image", "predict")):
+                    return inner
         except Exception as e:
-            logger.warning("Audio mux failed for two-stage; returning video without audio: %r", e)
-            final_path = result
-
-        msg = (
-            "Two-stage processing completed.\n"
-            f"Background: {background_choice}\n"
-            f"Chroma Preset: {chroma_preset}\n"
-            f"Device: {self.device_manager.get_optimal_device()}"
-        )
-        return final_path, msg
-
-    # ─────────────────────────────────────────────────────────────────────
-    # Status helpers
-    # ─────────────────────────────────────────────────────────────────────
-    def get_status(self) -> Dict[str, Any]:
-        status = {
-            "models_loaded": self.models_loaded,
-            "two_stage_available": bool(TWO_STAGE_AVAILABLE and (self.two_stage_processor is not None)),
-            "two_stage_origin": TWO_STAGE_IMPORT_ORIGIN or "",
-            "two_stage_error": TWO_STAGE_IMPORT_ERROR[:5000] if TWO_STAGE_IMPORT_ERROR else "",
-            "device": str(self.device_manager.get_optimal_device()),
-            "core_processor_loaded": self.core_processor is not None,
-            "config": self._safe_config_dict(),
-            "memory_usage": self._safe_memory_usage(),
-        }
-        try:
-            status["sam2_loaded"] = self.model_loader.get_sam2() is not None
-            status["matanyone_loaded"] = self.model_loader.get_matanyone() is not None
-        except Exception:
-            status["sam2_loaded"] = False
-            status["matanyone_loaded"] = False
-
-        if self.progress_tracker:
-            status["progress"] = self.progress_tracker.get_all_progress()
-        return status
+            logger.warning(f"SAM2 unwrap fail: {e}")
+        return None
 
-    def _safe_config_dict(self) -> Dict[str, Any]:
+    def _get_mask(self, frame: np.ndarray) -> np.ndarray:
         try:
-            return self.config.to_dict()
-        except Exception:
-            keys = ["use_nvenc", "prefer_mp4", "video_codec", "audio_copy",
-                    "ffmpeg_path", "max_model_size", "max_model_size_bytes", "output_dir"]
-            return {k: getattr(self.config, k, None) for k in keys}
+            return segment_person_hq(frame, self.sam2, fallback_enabled=True)
+        except Exception as e:
+            logger.warning(f"Segmentation fallback: {e}")
+            h, w = frame.shape[:2]
+            m = np.zeros((h, w), np.uint8)
+            m[h//6:5*h//6, w//4:3*w//4] = 255
+            return m
+
+    def _apply_greenscreen_hard(self, frame, mask, green_bg):
+        mask_u8 = self._to_binary_mask(mask)
+        mk = cv2.cvtColor(mask_u8, cv2.COLOR_GRAY2BGR).astype(np.float32) / 255.0
+        out = frame.astype(np.float32) * mk + green_bg.astype(np.float32) * (1.0 - mk)
+        return np.clip(out, 0, 255).astype(np.uint8)
 
-    def _safe_memory_usage(self) -> Dict[str, Any]:
+    @staticmethod
+    def _to_binary_mask(mask: np.ndarray) -> np.ndarray:
+        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) if mask.max() <= 1.0 else np.clip(mask, 0, 255).astype(np.uint8)
+        _, binm = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+        return binm
+
+    def _segmentation_mask_on_stage2(self, frame_bgr: np.ndarray) -> Optional[np.ndarray]:
         try:
-            return self.memory_manager.get_memory_usage()
+            if self.sam2 is None:
+                return None
+            return self._get_mask(frame_bgr)
         except Exception:
-            return {}
-
-    def cancel_processing(self):
-        self.cancel_event.set()
-        logger.info("Cancellation requested")
+            return None
 
-    def cleanup_resources(self):
-        try:
-            self.memory_manager.cleanup_aggressive()
-        except Exception:
-            pass
+    def _chroma_key_advanced(
+        self,
+        frame_bgr: np.ndarray,
+        bg_bgr:    np.ndarray,
+        settings:  Dict[str, Any],
+        seg_mask: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
         try:
-            self.model_loader.cleanup()
-        except Exception:
-            pass
-        logger.info("Resources cleaned up")
-
-
-# ╔══════════════════════════════════════════════════════════════════════════╗
-# ║                         Singleton instance + wrappers                    ║
-# ╚══════════════════════════════════════════════════════════════════════════╝
-processor = VideoProcessor()
-
-def load_models_with_validation(progress_callback: Optional[Callable] = None) -> str:
-    return processor.load_models(progress_callback)
-
-def process_video_fixed(
-    video_path: str,
-    background_choice: str,
-    custom_background_path: Optional[str],
-    progress_callback: Optional[Callable] = None,
-    use_two_stage: bool = False,
-    chroma_preset: str = "standard",
-    key_color_mode: str = "auto",
-    preview_mask: bool = False,
-    preview_greenscreen: bool = False,
-) -> Tuple[Optional[str], str]:
-    return processor.process_video(
-        video_path,
-        background_choice,
-        custom_background_path,
-        progress_callback,
-        use_two_stage,
-        chroma_preset,
-        key_color_mode,
-        preview_mask,
-        preview_greenscreen,
-    )
-
-def get_model_status() -> Dict[str, Any]:
-    return processor.get_status()
-
-def get_cache_status() -> Dict[str, Any]:
-    return processor.get_status()
-
-PROCESS_CANCELLED = processor.cancel_event
-
-
-# ╔══════════════════════════════════════════════════════════════════════════╗
-# ║                                    CLI                                  ║
-# ╚══════════════════════════════════════════════════════════════════════════╝
-def main():
-    try:
-        logger.info("Starting BackgroundFX Pro")
-        logger.info(f"Device: {processor.device_manager.get_optimal_device()}")
-        logger.info(
-            "Two-stage available at import-time: %s (origin='%s')",
-            TWO_STAGE_AVAILABLE, TWO_STAGE_IMPORT_ORIGIN or "unknown"
-        )
-
-        from ui.ui_components import create_interface
-        demo = create_interface()
-        demo.queue().launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True,
-            debug=False,
-        )
-    finally:
-        processor.cleanup_resources()
-
-
-if __name__ == "__main__":
-    main()
+            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))
+
+            f = frame_bgr.astype(np.float32)
+            b = bg_bgr.astype(np.float32)
+
+            diff = np.linalg.norm(f - key, axis=2)
+            alpha = np.clip((diff - tol * 0.6) / max(1e-6, tol * 0.4), 0.0, 1.0)
+            if soft > 0:
+                k = soft * 2 + 1
+                alpha = cv2.GaussianBlur(alpha, (k, k), soft)
+
+            # segmentation rescue
+            if seg_mask is not None:
+                if seg_mask.ndim == 3:
+                    seg_mask = cv2.cvtColor(seg_mask, cv2.COLOR_BGR2GRAY)
+                seg = seg_mask.astype(np.float32) / 255.0
+                seg = cv2.GaussianBlur(seg, (5, 5), 1.0)
+                alpha = np.clip(np.maximum(alpha, seg * 0.85), 0.0, 1.0)
+
+            # spill suppression
+            if spill > 0:
+                zone = 1.0 - alpha
+                g = f[:, :, 1]
+                f[:, :, 1] = np.clip(g - g * zone * spill, 0, 255)
+
+            mask3 = np.stack([alpha] * 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 key error: {e}")
+            return frame_bgr