Update processing/video/video_processor.py

processing/video/video_processor.py

@@ -2,105 +2,144 @@
 """
 Compatibility shim: CoreVideoProcessor
 
-Bridges the legacy import `from processing.video.video_processor import CoreVideoProcessor`
-to the modern pipeline functions living in `utils.cv_processing` and models in `core.models`.
+Bridges the legacy import
+    from processing.video.video_processor import CoreVideoProcessor
+to the modern pipeline functions in utils.cv_processing, using whatever
+models provider is passed in (e.g., models.loaders.ModelLoader).
+
+Requirements for the models provider:
+- get_sam2() -> predictor or None
+- get_matanyone() -> processor or None
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
 from typing import Optional, Dict, Any, Tuple, Callable
+import time
+import threading
 
 import cv2
 import numpy as np
-import time
-import threading
 
-from utils.logger import get_logger
-from core.models import ModelManager
+# Try project logger; fall back to std logging
+try:
+    from utils.logger import get_logger
+    _log = get_logger(__name__)
+except Exception:
+    import logging
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
+    _log = logging.getLogger(__name__)
 
-# ← these funcs are the ones you showed (in utils/cv_processing.py)
+# CV pipeline helpers
 from utils.cv_processing import (
     segment_person_hq,
     refine_mask_hq,
     replace_background_hq,
     create_professional_background,
     validate_video_file,
+    PROFESSIONAL_BACKGROUNDS,
 )
 
+
 @dataclass
 class ProcessorConfig:
     background_preset: str = "minimalist"   # key in PROFESSIONAL_BACKGROUNDS
     write_fps: Optional[float] = None       # None -> keep source fps
 
+
 class CoreVideoProcessor:
     """
     Minimal, safe implementation used by core/app.py.
-    It relies on ModelManager (SAM2 + MatAnyone) and your cv_processing helpers.
-    Now supports live progress + cancel/stop.
+    It relies on a models provider (e.g., ModelLoader) that implements:
+        - get_sam2()
+        - get_matanyone()
+    and uses utils.cv_processing for the pipeline.
+
+    Supports progress callback and cancellation via stop_event.
     """
 
-    def __init__(self, config: Optional[ProcessorConfig] = None, models: Optional[ModelManager] = None):
-        self.log = get_logger(f"{__name__}.CoreVideoProcessor")
+    def __init__(self, config: Optional[ProcessorConfig] = None, models: Optional[Any] = None):
+        self.log = _log
         self.config = config or ProcessorConfig()
-        self.models = models or ModelManager()
-        try:
-            self.models.load_all()
-        except Exception as e:
-            self.log.warning(f"Model load issue (will use fallbacks if needed): {e}")
+        self.models = models  # do NOT load here; core/app handles loading
+        if self.models is None:
+            self.log.warning("CoreVideoProcessor initialized without a models provider; will use fallbacks.")
 
-    # --- single-frame API (useful for images or per-frame video loop) ---
+    # ---------- Single frame ----------
     def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str, Any]:
-        """Return dict with composited frame + mask; always succeeds with fallbacks."""
+        """Return dict with composited frame + mask; always attempts fallbacks."""
         predictor = None
         try:
-            sam2_model = self.models.get_sam2()
-            if sam2_model is not None:
-                if hasattr(sam2_model, 'predictor'):
-                    predictor = sam2_model.predictor
-                elif hasattr(sam2_model, 'set_image'):
-                    predictor = sam2_model
-                elif isinstance(sam2_model, dict) and 'model' in sam2_model:
-                    self.log.warning("SAM2 loaded as dict format, not directly usable")
-                    predictor = None
-            if predictor is None:
-                self.log.debug("SAM2 predictor not available, will use fallback")
+            if self.models and hasattr(self.models, "get_sam2"):
+                predictor = self.models.get_sam2()
+                # Some wrappers expose predictor directly, others are already usable
+                # segment_person_hq checks for set_image/predict itself.
         except Exception as e:
             self.log.warning(f"SAM2 predictor unavailable: {e}")
 
-        # 1) segment
+        # 1) segmentation (with fallbacks inside)
         mask = segment_person_hq(frame, predictor, fallback_enabled=True)
 
-        # 2) refine
+        # 2) refinement (MatAnyOne if available, else robust OpenCV path)
         matanyone = None
         try:
-            matanyone_model = self.models.get_matanyone()
-            if matanyone_model is not None:
-                matanyone = matanyone_model
+            if self.models and hasattr(self.models, "get_matanyone"):
+                matanyone = self.models.get_matanyone()
         except Exception as e:
-            self.log.warning(f"MatAnyone unavailable: {e}")
+            self.log.warning(f"MatAnyOne unavailable: {e}")
 
         mask_refined = refine_mask_hq(frame, mask, matanyone, fallback_enabled=True)
 
-        # 3) composite
+        # 3) compositing
         out = replace_background_hq(frame, mask_refined, background, fallback_enabled=True)
 
         return {"frame": out, "mask": mask_refined}
 
-    # --- simple video API (covers typical usage in older core/app.py code) ---
+    # ---------- Build background once per video ----------
+    def _prepare_background_from_config(
+        self,
+        bg_config: Optional[Dict[str, Any]],
+        width: int,
+        height: int
+    ) -> np.ndarray:
+        """
+        Accepts either:
+          - {"custom_path": "/path/to/image.png"} → load image
+          - {"background_choice": "minimalist"} → preset
+          - None → use self.config.background_preset
+        """
+        # 1) custom image?
+        if bg_config and bg_config.get("custom_path"):
+            path = bg_config["custom_path"]
+            img = cv2.imread(path, cv2.IMREAD_COLOR)
+            if img is None:
+                self.log.warning(f"Custom background at '{path}' could not be read. Falling back to preset.")
+            else:
+                return cv2.resize(img, (width, height), interpolation=cv2.INTER_LANCZOS4)
+
+        # 2) preset (explicit choice or default)
+        choice = None
+        if bg_config and "background_choice" in bg_config:
+            choice = bg_config["background_choice"]
+        if not choice:
+            choice = self.config.background_preset
+
+        cfg = PROFESSIONAL_BACKGROUNDS.get(choice, PROFESSIONAL_BACKGROUNDS["minimalist"])
+        return create_professional_background(cfg, width, height)
+
+    # ---------- Full video ----------
     def process_video(
         self,
         input_path: str,
         output_path: str,
         bg_config: Optional[Dict[str, Any]] = None,
-        progress_callback: Optional[Callable[[int, int, float], None]] = None,  # <-- ADDED
-        stop_event: Optional[threading.Event] = None   # <-- ADDED
+        progress_callback: Optional[Callable[[int, int, float], None]] = None,
+        stop_event: Optional[threading.Event] = None
     ) -> Dict[str, Any]:
         """
-        Process a full video with live progress and optional stop.
-        progress_callback: function(current_frame, total_frames, fps)
-        stop_event: threading.Event() - if set(), abort processing.
-        Returns: dict with stats.
+        Process a full video with live progress and optional cancel.
+        progress_callback(current_frame, total_frames, fps_live)
         """
         ok, msg = validate_video_file(input_path)
         if not ok:
@@ -115,16 +154,16 @@ def process_video(
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         fps    = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        fps_out = self.config.write_fps or (fps if fps and fps > 0 else 25.0)
 
+        fps_out = self.config.write_fps or (fps if fps and fps > 0 else 25.0)
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        writer = cv2.VideoWriter(output_path, fourcc, fps_out, (width, height))
+        writer = cv2.VideoWriter(output_path, fourcc, float(fps_out), (width, height))
+        if not writer.isOpened():
+            cap.release()
+            raise RuntimeError(f"Could not open writer for: {output_path}")
 
-        # Build background (once)
-        from utils.cv_processing import PROFESSIONAL_BACKGROUNDS
-        preset = self.config.background_preset
-        cfg = bg_config or PROFESSIONAL_BACKGROUNDS.get(preset, PROFESSIONAL_BACKGROUNDS["minimalist"])
-        background = create_professional_background(cfg, width, height)
+        # Build background once
+        background = self._prepare_background_from_config(bg_config, width, height)
 
         frame_count = 0
         start_time = time.time()
@@ -134,24 +173,25 @@ def process_video(
                 if not ret:
                     break
 
-                # --- CANCEL SUPPORT ---
+                # Cancel support
                 if stop_event is not None and stop_event.is_set():
-                    self.log.info("Processing stopped by user request")  # <-- CHANGED
+                    self.log.info("Processing stopped by user request.")
                     break
 
-                res = self.process_frame(frame, background)
-                writer.write(res["frame"])
+                # Process single frame
+                result = self.process_frame(frame, background)
+                writer.write(result["frame"])
                 frame_count += 1
 
-                # --- LIVE PROGRESS ---
+                # Progress callback
                 if progress_callback:
                     elapsed = time.time() - start_time
-                    fps_live = frame_count / elapsed if elapsed > 0 else 0
-                    progress_callback(
-                        frame_count,
-                        total_frames,
-                        fps_live
-                    )
+                    fps_live = frame_count / elapsed if elapsed > 0 else 0.0
+                    try:
+                        progress_callback(frame_count, total_frames, fps_live)
+                    except Exception:
+                        # Don’t break processing due to a UI callback error
+                        pass
         finally:
             cap.release()
             writer.release()
@@ -161,8 +201,10 @@ def process_video(
             "frames": frame_count,
             "width": width,
             "height": height,
-            "fps_out": fps_out
+            "fps_out": float(fps_out),
+            "output_path": output_path,
         }
 
-# Backward-compat export name
+
+# Backward-compat alias used elsewhere
 VideoProcessor = CoreVideoProcessor