Update processing/video/video_processor.py

processing/video/video_processor.py

@@ -9,10 +9,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, Dict, Any, Tuple
+from typing import Optional, Dict, Any, Tuple, Callable
 
 import cv2
 import numpy as np
+import time
+import threading
 
 from utils.logger import get_logger
 from core.models import ModelManager
@@ -31,11 +33,11 @@ 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.
     """
 
     def __init__(self, config: Optional[ProcessorConfig] = None, models: Optional[ModelManager] = None):
@@ -53,21 +55,14 @@ def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str,
         predictor = None
         try:
             sam2_model = self.models.get_sam2()
-            # Check if we have a working SAM2 predictor
-            # SAM2ImagePredictor has set_image and predict methods
             if sam2_model is not None:
-                # Check if it's wrapped (has .predictor attribute)
                 if hasattr(sam2_model, 'predictor'):
                     predictor = sam2_model.predictor
-                # Or if it IS the predictor (has set_image method)
                 elif hasattr(sam2_model, 'set_image'):
                     predictor = sam2_model
-                # Or if it's a dict with model and processor (from transformers)
                 elif isinstance(sam2_model, dict) and 'model' in sam2_model:
-                    # For now, we can't use this format easily
                     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")
         except Exception as e:
@@ -80,7 +75,6 @@ def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str,
         matanyone = None
         try:
             matanyone_model = self.models.get_matanyone()
-            # Just check if we have a MatAnyone model at all
             if matanyone_model is not None:
                 matanyone = matanyone_model
         except Exception as e:
@@ -94,8 +88,20 @@ def process_frame(self, frame: np.ndarray, background: np.ndarray) -> Dict[str,
         return {"frame": out, "mask": mask_refined}
 
     # --- simple video API (covers typical usage in older core/app.py code) ---
-    def process_video(self, input_path: str, output_path: str, bg_config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
-        """Process a full video; returns basic stats."""
+    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
+    ) -> 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.
+        """
         ok, msg = validate_video_file(input_path)
         if not ok:
             raise ValueError(f"Invalid video: {msg}")
@@ -108,32 +114,55 @@ def process_video(self, input_path: str, output_path: str, bg_config: Optional[D
         width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         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)
 
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         writer = cv2.VideoWriter(output_path, fourcc, fps_out, (width, height))
 
         # Build background (once)
-        from utils.cv_processing import PROFESSIONAL_BACKGROUNDS  # local import to avoid circulars
+        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)
 
         frame_count = 0
+        start_time = time.time()
         try:
             while True:
                 ret, frame = cap.read()
                 if not ret:
                     break
+
+                # --- CANCEL SUPPORT ---
+                if stop_event is not None and stop_event.is_set():
+                    self.log.info("Processing stopped by user request")  # <-- CHANGED
+                    break
+
                 res = self.process_frame(frame, background)
                 writer.write(res["frame"])
                 frame_count += 1
+
+                # --- LIVE PROGRESS ---
+                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
+                    )
         finally:
             cap.release()
             writer.release()
 
         self.log.info(f"Processed {frame_count} frames → {output_path}")
-        return {"frames": frame_count, "width": width, "height": height, "fps_out": fps_out}
+        return {
+            "frames": frame_count,
+            "width": width,
+            "height": height,
+            "fps_out": fps_out
+        }
 
 # Backward-compat export name
-VideoProcessor = CoreVideoProcessor
+VideoProcessor = CoreVideoProcessor