Update core/app.py

core/app.py

@@ -1,457 +1,482 @@
 #!/usr/bin/env python3
 """
-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
+BackgroundFX Pro – Main Application Entry Point
+Refactored modular architecture – orchestrates specialised components
 """
 
 from __future__ import annotations
 
-import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
-from pathlib import Path
-from typing import Optional, Dict, Any, Callable, Tuple, List
-
-from utils.cv_processing import segment_person_hq, refine_mask_hq
+# ── Early env/threading hygiene (safe default to silence libgomp) ────────────
+import os
+if not os.environ.get("OMP_NUM_THREADS", "").isdigit():
+    os.environ["OMP_NUM_THREADS"] = "2"
 
-# Project logger if available
+# If you use early_env in your project, keep this import (harmless if absent)
 try:
-    from utils.logger import get_logger
-    logger = get_logger(__name__)
+    import early_env  # sets OMP/MKL/OPENBLAS + torch threads safely
 except Exception:
-    logger = logging.getLogger(__name__)
+    pass
 
+import logging
+import threading
+import traceback
+import sys
+from pathlib import Path
+from typing import Optional, Tuple, Dict, Any, Callable
+
+# ── Logging ──────────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("core.app")
+
+# ── Ensure project root importable ───────────────────────────────────────────
+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))
+
+# ── Gradio schema patch (HF 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}")
+
+# ── 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
+
+# ── Optional Two-Stage import (exact module path) ────────────────────────────
+TWO_STAGE_AVAILABLE = False
+TWO_STAGE_IMPORT_ORIGIN = ""
+TWO_STAGE_IMPORT_ERROR = ""
+CHROMA_PRESETS: Dict[str, Dict[str, Any]] = {"standard": {}}
+TwoStageProcessor = None  # type: ignore
 
-# ---------------------------------------------------------------------------
-# 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
-
+try:
+    from processing.two_stage.two_stage_processor import TwoStageProcessor, CHROMA_PRESETS  # type: ignore
+    TWO_STAGE_AVAILABLE = True
+    TWO_STAGE_IMPORT_ORIGIN = "processing.two_stage.two_stage_processor"
+    logger.info("Two-stage import OK (%s)", TWO_STAGE_IMPORT_ORIGIN)
+except Exception as e:
+    TWO_STAGE_AVAILABLE = False
+    TWO_STAGE_IMPORT_ERROR = f"{repr(e)}\n{traceback.format_exc()}"
+    logger.warning("Two-stage import FAILED: %s", TWO_STAGE_IMPORT_ERROR.strip().splitlines()[-1])
+
+# ╔══════════════��═══════════════════════════════════════════════════════════╗
+# ║                           VideoProcessor class                           ║
+# ╚══════════════════════════════════════════════════════════════════════════╝
+class VideoProcessor:
+    """
+    Main orchestrator – coordinates all specialised components.
+    """
+
+    def __init__(self):
+        self.config = get_config()
+        self._patch_config_defaults(self.config)  # avoid AttributeError on older configs
+
+        self.device_manager = DeviceManager()
+        self.memory_manager = MemoryManager(self.device_manager.get_optimal_device())
+        self.model_loader = ModelLoader(self.device_manager, self.memory_manager)
+
+        self.audio_processor = AudioProcessor()
+        self.core_processor: CoreVideoProcessor | None = None
+        self.two_stage_processor: Any | None = None
+
+        self.models_loaded = False
+        self.loading_lock = threading.Lock()
+        self.cancel_event = threading.Event()
+        self.progress_tracker: ProgressTracker | None = None
+
+        logger.info(f"VideoProcessor on device: {self.device_manager.get_optimal_device()}")
+
+    # ── Config hardening: add missing fields safely ───────────────────────────
+    @staticmethod
+    def _patch_config_defaults(cfg: Any) -> None:
+        defaults = {
+            # video / i/o
+            "use_nvenc": False,
+            "prefer_mp4": True,
+            "video_codec": "mp4v",
+            "audio_copy": True,
+            "ffmpeg_path": "ffmpeg",
+            # model/resource guards
+            "max_model_size": 0,
+            "max_model_size_bytes": 0,
+            # housekeeping
+            "output_dir": str((Path(__file__).resolve().parent.parent) / "outputs"),
+        }
+        for k, v in defaults.items():
+            if not hasattr(cfg, k):
+                setattr(cfg, k, v)
+        Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
+
+    # ── Progress helper ───────────────────────────────────────────────────────
+    def _init_progress(self, video_path: str, cb: Optional[Callable] = None):
         try:
-            _prog(0.0, "Stage 1: opening video…")
+            import cv2
             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}")
-
-                # 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))
-
-                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 '?'}")
-
+            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             cap.release()
-            writer.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"
 
-            # save mask cache
             try:
-                cache_file = self.mask_cache_dir / (Path(out_path).stem + "_masks.pkl")
-                with open(cache_file, "wb") as f:
-                    pickle.dump(masks, f)
+                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
+                sam2_predictor = getattr(sam2_loaded, "model", None) if sam2_loaded else None
+                mat_model = getattr(mat_loaded, "model", None) if mat_loaded else None
+
+                # single-stage
+                self.core_processor = CoreVideoProcessor(config=self.config, models=self.model_loader)
+
+                # two-stage (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")
+                    except Exception as e:
+                        logger.warning("Two-stage init failed: %r", e)
+                        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
             except Exception as e:
-                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.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(
         self,
-        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,
+        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,
     ) -> 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)
 
-        def _prog(p, d):
-            if progress_callback:
-                try:
-                    progress_callback(float(p), str(d))
-                except Exception:
-                    pass
+        ok, why = validate_video_file(video_path)
+        if not ok:
+            return None, f"Invalid video: {why}"
 
         try:
-            _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"
+            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,
+                )
             else:
-                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)"
+                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}"
         except Exception as e:
-            logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
-            return None, f"Stage 2 failed: {e}"
+            logger.error(f"Unexpected processing error: {e}")
+            logger.debug("Traceback:\n%s", traceback.format_exc())
+            return None, f"Unexpected error: {e}"
 
-    # ---------------------------------------------------------------------
-    # Full pipeline – now passes chosen key into Stage 2
-    # ---------------------------------------------------------------------
-    def process_full_pipeline(
+    # ── Private – single-stage ───────────────────────────────────────────────
+    def _process_single_stage(
         self,
         video_path: 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,
+        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 audio unless preview-only
+        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(
+        self,
+        video_path: str,
+        background_choice: str,
+        custom_background_path: Optional[str],
+        progress_callback: Optional[Callable],
+        chroma_preset: str,
+        key_color_mode: str,
     ) -> Tuple[Optional[str], str]:
-        gs_tmp = tempfile.mktemp(suffix="_gs.mp4")
+        if self.two_stage_processor is None:
+            return None, "Two-stage processor not available"
+
+        import cv2, time
+        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()
+
+        # Prepare background via core processor
         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
+            background = self.core_processor.prepare_background(
+                background_choice, custom_background_path, w, h
             )
-            if gs_info is None:
-                return None, msg1
-
-            # 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
-
-            result, msg2 = self.stage2_greenscreen_to_final(
-                gs_info["path"], background, final_output,
-                chroma_settings=cs, progress_callback=progress_callback, stop_event=stop_event
-            )
-            return result, msg2
-        finally:
-            try:
-                os.remove(gs_tmp)
-            except Exception:
-                pass
-            gc.collect()
-
-    # ---------------------------------------------------------------------
-    # Internal helpers
-    # ---------------------------------------------------------------------
-    def _unwrap_sam2(self, obj):
+        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", chroma_preset, key_color_mode)
+
+        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)
         try:
-            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
+            final_path = self.audio_processor.add_audio_to_video(
+                original_video=video_path, processed_video=result
+            )
         except Exception as e:
-            logger.warning(f"SAM2 unwrap fail: {e}")
-        return None
+            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[:2000] 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
 
-    def _get_mask(self, frame: np.ndarray) -> np.ndarray:
+    def _safe_config_dict(self) -> Dict[str, Any]:
         try:
-            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)
+            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}
 
-    @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]:
+    def _safe_memory_usage(self) -> Dict[str, Any]:
         try:
-            if self.sam2 is None:
-                return None
-            return self._get_mask(frame_bgr)
+            return self.memory_manager.get_memory_usage()
         except Exception:
-            return None
+            return {}
 
-    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:
+    def cancel_processing(self):
+        self.cancel_event.set()
+        logger.info("Cancellation requested")
+
+    def cleanup_resources(self):
         try:
-            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
+            self.memory_manager.cleanup_aggressive()
+        except Exception:
+            pass
+        try:
+            self.model_loader.cleanup()
+        except Exception:
+            pass
+        logger.info("Resources cleaned up")
+
+
+# ── Singleton + thin wrappers (used by UI callbacks) ────────────────────────
+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 entrypoint (must exist; app.py imports main) ─────────────────────────
+def main():
+    try:
+        logger.info("Starting BackgroundFX Pro")
+        logger.info(f"Device: {processor.device_manager.get_optimal_device()}")
+        logger.info("Two-stage available (import): %s", TWO_STAGE_AVAILABLE)
+
+        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()