Update core/app.py

core/app.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 """
-BackgroundFX Pro - Main Application Entry Point
-Refactored modular architecture - orchestrates specialized components
+BackgroundFX Pro – Main Application Entry Point
+Refactored modular architecture – orchestrates specialised components
 """
 
-# 0) Early env/threading hygiene (must be first)
+# ─────────────────────────────────────────────────────────────────────────────
+# 0) Early env/threading hygiene (must run first)
+# ─────────────────────────────────────────────────────────────────────────────
 import early_env  # sets OMP/MKL/OPENBLAS + torch threads safely
 
 import logging
@@ -12,16 +14,21 @@
 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"
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 )
 logger = logging.getLogger("core.app")
 
+# ─────────────────────────────────────────────────────────────────────────────
 # 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):
@@ -40,7 +47,9 @@ def _patched_get_type(schema):
 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
@@ -52,35 +61,41 @@ def _patched_get_type(schema):
 
 # Optional two-stage processor
 try:
-    from processing.two_stage.two_stage_processor import TwoStageProcessor, CHROMA_PRESETS
+    from processing.two_stage.two_stage_processor import (
+        TwoStageProcessor,
+        CHROMA_PRESETS,
+    )
+
     TWO_STAGE_AVAILABLE = True
 except Exception:
     TWO_STAGE_AVAILABLE = False
     CHROMA_PRESETS = {"standard": {}}
 
-# Validation helper for inputs (lives with CV utils)
+# Validation helper
 from utils.cv_processing import validate_video_file
 
 
+# ╔══════════════════════════════════════════════════════════════════════════╗
+# ║                           VideoProcessor class                           ║
+# ╚══════════════════════════════════════════════════════════════════════════╝
 class VideoProcessor:
     """
-    Main video processing orchestrator - coordinates all specialized components.
+    Main orchestrator – coordinates all specialised components.
     """
+
+    # ─────────────────────────────────────────────────────────────────────
+    # Init
+    # ─────────────────────────────────────────────────────────────────────
     def __init__(self):
-        self.config = get_config()  # singleton-style app config
+        self.config = get_config()
         self.device_manager = DeviceManager()
-
-        # Memory manager now requires a device object/string
         self.memory_manager = MemoryManager(self.device_manager.get_optimal_device())
-
-        # Model loader takes device + memory managers
         self.model_loader = ModelLoader(self.device_manager, self.memory_manager)
 
         self.audio_processor = AudioProcessor()
         self.core_processor: CoreVideoProcessor | None = None
         self.two_stage_processor: TwoStageProcessor | None = None
 
-        # State
         self.models_loaded = False
         self.loading_lock = threading.Lock()
         self.cancel_event = threading.Event()
@@ -88,10 +103,13 @@ def __init__(self):
 
         logger.info(f"VideoProcessor on device: {self.device_manager.get_optimal_device()}")
 
-    # ---------- Progress ----------
+    # ─────────────────────────────────────────────────────────────────────
+    # 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()
@@ -102,7 +120,9 @@ def _init_progress(self, video_path: str, cb: Optional[Callable] = None):
             logger.warning(f"Progress init failed: {e}")
             self.progress_tracker = ProgressTracker(100, cb)
 
-    # ---------- Load Models ----------
+    # ─────────────────────────────────────────────────────────────────────
+    # Model loading
+    # ─────────────────────────────────────────────────────────────────────
     def load_models(self, progress_callback: Optional[Callable] = None) -> str:
         with self.loading_lock:
             if self.models_loaded:
@@ -111,57 +131,44 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
             try:
                 self.cancel_event.clear()
                 if progress_callback:
-                    progress_callback(0.0, f"Loading on {self.device_manager.get_optimal_device()}")
-
-                try:
-                    sam2_loaded, mat_loaded = self.model_loader.load_all_models(
-                        progress_callback=progress_callback,
-                        cancel_event=self.cancel_event
+                    progress_callback(
+                        0.0, f"Loading on {self.device_manager.get_optimal_device()}"
                     )
-                except IndexError as e:
-                    import traceback
-                    tb = traceback.extract_tb(e.__traceback__)
-                    where = f"{tb[-1].filename}:{tb[-1].lineno}" if tb else "unknown"
-                    logger.error(f"IndexError in load_all_models at {where}: {e}")
-                    raise ModelLoadingError(f"Model loading failed (IndexError @ {where}): {e}")
-                except Exception as e:
-                    import traceback
-                    logger.error(f"Unexpected error in load_all_models: {e}\n{traceback.format_exc()}")
-                    raise
+
+                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 model refs for two-stage
+                # Unwrap actual predictor / model objects
                 sam2_predictor = sam2_loaded.model if sam2_loaded else None
                 mat_model = mat_loaded.model if mat_loaded else None
-                if (sam2_predictor is None) and (mat_model is None):
-                    return "Model loading failed - see logs"
 
-                # Core processor expects a "models" provider (we pass the loader itself)
+                # Core single-stage processor
                 self.core_processor = CoreVideoProcessor(
-                    config=self.config,
-                    models=self.model_loader
+                    config=self.config, models=self.model_loader
                 )
 
-                # Optional 2-stage
+                # Two-stage processor (optional)
                 if TWO_STAGE_AVAILABLE and (sam2_predictor or mat_model):
                     try:
                         self.two_stage_processor = TwoStageProcessor(
-                            sam2_predictor=sam2_predictor,
-                            matanyone_model=mat_model
+                            sam2_predictor=sam2_predictor, matanyone_model=mat_model
                         )
-                        logger.info("Two-stage processor initialized")
+                        logger.info("Two-stage processor initialised")
                     except Exception as e:
                         logger.warning(f"Two-stage init failed: {e}")
                         self.two_stage_processor = None
 
                 self.models_loaded = True
                 msg = self.model_loader.get_load_summary()
-                if self.two_stage_processor:
-                    msg += "\n✅ Two-stage processor ready"
-                else:
-                    msg += "\n⚠️ Two-stage processor not available"
+                msg += (
+                    "\n✅ Two-stage processor ready"
+                    if self.two_stage_processor
+                    else "\n⚠️ Two-stage processor not available"
+                )
                 logger.info(msg)
                 return msg
 
@@ -176,7 +183,9 @@ def load_models(self, progress_callback: Optional[Callable] = None) -> str:
                 logger.error(err)
                 return err
 
-    # ---------- Process ----------
+    # ─────────────────────────────────────────────────────────────────────
+    # Public entry – process video
+    # ─────────────────────────────────────────────────────────────────────
     def process_video(
         self,
         video_path: str,
@@ -185,13 +194,15 @@ def process_video(
         progress_callback: Optional[Callable] = None,
         use_two_stage: bool = False,
         chroma_preset: str = "standard",
+        key_color_mode: str = "auto",          # NEW
         preview_mask: bool = False,
         preview_greenscreen: bool = False,
     ) -> Tuple[Optional[str], str]:
-
+        """
+        Dispatch to single-stage or two-stage pipeline.
+        """
         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"
 
@@ -206,10 +217,24 @@ def process_video(
                 if not TWO_STAGE_AVAILABLE:
                     return None, "Two-stage processing not available on this build"
                 if not self.two_stage_processor:
-                    return None, "Two-stage processor not initialized (models not ready?)"
-                return self._process_two_stage(video_path, background_choice, custom_background_path, progress_callback, chroma_preset)
+                    return None, "Two-stage processor not initialised"
+                return self._process_two_stage(
+                    video_path,
+                    background_choice,
+                    custom_background_path,
+                    progress_callback,
+                    chroma_preset,
+                    key_color_mode,           # NEW
+                )
             else:
-                return self._process_single_stage(video_path, background_choice, custom_background_path, progress_callback, preview_mask, preview_greenscreen)
+                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}")
@@ -218,6 +243,9 @@ def process_video(
             logger.error(f"Unexpected processing error: {e}")
             return None, f"Unexpected error: {e}"
 
+    # ─────────────────────────────────────────────────────────────────────
+    # Private – single-stage
+    # ─────────────────────────────────────────────────────────────────────
     def _process_single_stage(
         self,
         video_path: str,
@@ -227,8 +255,8 @@ def _process_single_stage(
         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)
@@ -237,13 +265,18 @@ def _process_single_stage(
         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}
+            bg_config={
+                "background_choice": background_choice,
+                "custom_path": custom_background_path,
+            },
         )
         if not result:
             return None, "Video processing failed"
 
         if not (preview_mask or preview_greenscreen):
-            final_path = self.audio_processor.add_audio_to_video(original_video=video_path, processed_video=out_path)
+            final_path = self.audio_processor.add_audio_to_video(
+                original_video=video_path, processed_video=out_path
+            )
         else:
             final_path = out_path
 
@@ -256,6 +289,9 @@ def _process_single_stage(
         )
         return final_path, msg
 
+    # ─────────────────────────────────────────────────────────────────────
+    # Private – two-stage
+    # ─────────────────────────────────────────────────────────────────────
     def _process_two_stage(
         self,
         video_path: str,
@@ -263,17 +299,21 @@ def _process_two_stage(
         custom_background_path: Optional[str],
         progress_callback: Optional[Callable],
         chroma_preset: str,
+        key_color_mode: str,               # NEW
     ) -> Tuple[Optional[str], str]:
         if self.two_stage_processor is None:
             return None, "Two-stage processor not available"
 
         import cv2, time
+
         cap = cv2.VideoCapture(video_path)
         w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         cap.release()
 
-        background = self.core_processor.prepare_background(background_choice, custom_background_path, w, h)
+        background = self.core_processor.prepare_background(
+            background_choice, custom_background_path, w, h
+        )
         if background is None:
             return None, "Failed to prepare background"
 
@@ -282,11 +322,16 @@ def _process_two_stage(
         out_dir.mkdir(parents=True, exist_ok=True)
         final_out = str(out_dir / f"final_{ts}.mp4")
 
-        chroma = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS["standard"])
-        logger.info(f"Two-stage with preset: {chroma_preset}")
+        chroma_cfg = CHROMA_PRESETS.get(chroma_preset, CHROMA_PRESETS["standard"])
+        logger.info(f"Two-stage with preset: {chroma_preset} and key_color_mode={key_color_mode}")
 
         result, message = self.two_stage_processor.process_full_pipeline(
-            video_path, background, final_out, chroma_settings=chroma, progress_callback=progress_callback
+            video_path,
+            background,
+            final_out,
+            key_color_mode=key_color_mode,   # NEW
+            chroma_settings=chroma_cfg,
+            progress_callback=progress_callback,
         )
         if result is None:
             return None, message
@@ -299,11 +344,14 @@ def _process_two_stage(
         )
         return result, msg
 
-    # ---------- Status / Control ----------
+    # ─────────────────────────────────────────────────────────────────────
+    # Status helpers
+    # ─────────────────────────────────────────────────────────────────────
     def get_status(self) -> Dict[str, Any]:
         status = {
             "models_loaded": self.models_loaded,
-            "two_stage_available": TWO_STAGE_AVAILABLE and (self.two_stage_processor is not None),
+            "two_stage_available": TWO_STAGE_AVAILABLE
+            and (self.two_stage_processor is not None),
             "device": str(self.device_manager.get_optimal_device()),
             "memory_usage": self.memory_manager.get_memory_usage(),
             "config": self.config.to_dict(),
@@ -311,7 +359,9 @@ def get_status(self) -> Dict[str, Any]:
         }
         try:
             status["sam2_loaded"] = self.model_loader.get_sam2() is not None
-            status["matanyone_loaded"] = self.model_loader.get_matanyone() is not None
+            status["matanyone_loaded"] = (
+                self.model_loader.get_matanyone() is not None
+            )
         except Exception:
             status["sam2_loaded"] = False
             status["matanyone_loaded"] = False
@@ -326,15 +376,15 @@ def cancel_processing(self):
 
     def cleanup_resources(self):
         self.memory_manager.cleanup_aggressive()
-        if self.model_loader:
-            self.model_loader.cleanup()
+        self.model_loader.cleanup()
         logger.info("Resources cleaned up")
 
 
-# Singleton for UI callbacks
+# ╔══════════════════════════════════════════════════════════════════════════╗
+# ║                         Singleton instance + wrappers                    ║
+# ╚══════════════════════════════════════════════════════════════════════════╝
 processor = VideoProcessor()
 
-# Back-compat wrappers used by ui/callbacks.py
 def load_models_with_validation(progress_callback: Optional[Callable] = None) -> str:
     return processor.load_models(progress_callback)
 
@@ -345,43 +395,51 @@ def process_video_fixed(
     progress_callback: Optional[Callable] = None,
     use_two_stage: bool = False,
     chroma_preset: str = "standard",
+    key_color_mode: str = "auto",          # NEW
     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, preview_mask, preview_greenscreen
+        video_path,
+        background_choice,
+        custom_background_path,
+        progress_callback,
+        use_two_stage,
+        chroma_preset,
+        key_color_mode,                     # NEW
+        preview_mask,
+        preview_greenscreen,
     )
 
 def get_model_status() -> Dict[str, Any]:
     return processor.get_status()
 
 def get_cache_status() -> Dict[str, Any]:
-    # For now same as status; can add cache metrics later
+    # Placeholder – could expose FS cache size, etc.
     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(f"Two-stage available: {TWO_STAGE_AVAILABLE}")
 
-        # NOTE: UI was split into ui/components.py
+        # UI lives in ui/components.py
         from ui.components import create_interface
-        demo = create_interface()
 
+        demo = create_interface()
         demo.queue().launch(
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            debug=False
+            debug=False,
         )
-    except Exception as e:
-        logger.error(f"Startup failed: {e}")
-        raise
     finally:
         processor.cleanup_resources()