Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -4,22 +4,7 @@
 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
-
-Additional safety:
-  * Ensures MatAnyone receives a valid first-frame mask (bootstraps the session
-    with the first SAM2 mask). This prevents "First frame arrived without a mask"
-    warnings and shape mismatches inside the stateful refiner.
-
-Quality profiles (set via env BFX_QUALITY = speed | balanced | max):
-  * refine cadence, spill suppression, edge softness
-  * hybrid matte mix (segmentation vs chroma), small dilate/blur on mask
-  * optional tiny background blur to hide seams on very flat backgrounds
+UPDATED: Enhanced quality profiles, improved frame handling, better status reporting
 """
 from __future__ import annotations
 
@@ -37,9 +22,10 @@
     logger = logging.getLogger(__name__)
 
 # ---------------------------------------------------------------------------
-# Local video-writer helper
+# Local video-writer helper with frame count guarantee
 # ---------------------------------------------------------------------------
 def create_video_writer(output_path: str, fps: float, width: int, height: int, prefer_mp4: bool = True):
+    """Creates video writer with fallback options"""
     try:
         ext = ".mp4" if prefer_mp4 else ".avi"
         if not output_path:
@@ -64,6 +50,64 @@ def create_video_writer(output_path: str, fps: float, width: int, height: int, p
         logger.error(f"create_video_writer failed: {e}")
         return None, output_path
 
+# ---------------------------------------------------------------------------
+# Robust video writer wrapper to prevent frame loss
+# ---------------------------------------------------------------------------
+class RobustVideoWriter:
+    """Wrapper that ensures all frames are written"""
+    
+    def __init__(self, writer, output_path: str):
+        self.writer = writer
+        self.output_path = output_path
+        self.frame_buffer = []
+        self.frames_written = 0
+        self.frames_attempted = 0
+        
+    def write(self, frame):
+        """Buffer and write frame"""
+        if frame is None:
+            return False
+        
+        self.frames_attempted += 1
+        self.frame_buffer.append(frame.copy())
+        
+        # Try to write buffered frames
+        while self.frame_buffer and self.writer:
+            try:
+                self.writer.write(self.frame_buffer[0])
+                self.frame_buffer.pop(0)
+                self.frames_written += 1
+            except Exception as e:
+                logger.warning(f"Frame write failed: {e}")
+                return False
+        return True
+    
+    def release(self):
+        """Flush remaining frames and close"""
+        # Write any remaining buffered frames
+        while self.frame_buffer and self.writer:
+            try:
+                self.writer.write(self.frame_buffer[0])
+                self.frame_buffer.pop(0)
+                self.frames_written += 1
+            except Exception:
+                break
+        
+        # Close writer
+        if self.writer:
+            self.writer.release()
+            
+        # Log statistics
+        logger.info(f"Video writer closed: {self.frames_written}/{self.frames_attempted} frames written")
+        
+        # Verify output exists
+        if os.path.exists(self.output_path):
+            size = os.path.getsize(self.output_path)
+            if size == 0:
+                logger.error("WARNING: Output file is empty!")
+            else:
+                logger.info(f"Output file size: {size:,} bytes")
+
 # ---------------------------------------------------------------------------
 # Key-colour helpers (fast, no external deps)
 # ---------------------------------------------------------------------------
@@ -72,13 +116,11 @@ def _bgr_to_hsv_hue_deg(bgr: np.ndarray) -> np.ndarray:
     # 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},
@@ -87,7 +129,6 @@ def _key_candidates_bgr() -> dict:
         "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:
@@ -110,7 +151,6 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
     except Exception:
         return _key_candidates_bgr()["green"]
 
-
 # ---------------------------------------------------------------------------
 # Chroma presets
 # ---------------------------------------------------------------------------
@@ -121,12 +161,36 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 }
 
 # ---------------------------------------------------------------------------
-# Quality profiles (env: BFX_QUALITY = speed | balanced | max)
+# ENHANCED Quality profiles with clear differentiation
 # ---------------------------------------------------------------------------
 QUALITY_PROFILES: Dict[str, Dict[str, Any]] = {
-    "speed":    dict(refine_stride=4, spill=0.30, edge_softness=2, mix=0.60, dilate=0, blur=0, bg_sigma=0.0),
-    "balanced": dict(refine_stride=2, spill=0.40, edge_softness=2, mix=0.75, dilate=1, blur=1, bg_sigma=0.6),
-    "max":      dict(refine_stride=1, spill=0.45, edge_softness=3, mix=0.85, dilate=2, blur=2, bg_sigma=1.0),
+    "speed": dict(
+        refine_stride=10,  # Refine every 10th frame only
+        spill=0.15,        # Minimal spill work
+        edge_softness=1,   # Basic edges
+        mix=0.50,          # 50/50 chroma/segmentation
+        dilate=1,          # Minimal morphology
+        blur=0,            # No blur
+        bg_sigma=0.0       # No background blur
+    ),
+    "balanced": dict(
+        refine_stride=3,   # Refine every 3rd frame
+        spill=0.35,        # Moderate spill removal
+        edge_softness=2,   # Smooth edges
+        mix=0.70,          # Favor segmentation (70%)
+        dilate=2,          # Some hole filling
+        blur=1,            # Light feathering
+        bg_sigma=0.8       # Subtle background blur
+    ),
+    "max": dict(
+        refine_stride=1,   # Refine EVERY frame
+        spill=0.50,        # Strong spill removal
+        edge_softness=3,   # Very smooth edges
+        mix=0.85,          # Heavy segmentation bias (85%)
+        dilate=3,          # Strong hole filling
+        blur=2,            # More feathering
+        bg_sigma=1.5       # Visible background blur
+    ),
 }
 
 # ---------------------------------------------------------------------------
@@ -134,7 +198,7 @@ def _choose_best_key_color(frame_bgr: np.ndarray, mask_uint8: np.ndarray) -> dic
 # ---------------------------------------------------------------------------
 class TwoStageProcessor:
     def __init__(self, sam2_predictor=None, matanyone_model=None):
-        self.sam2  = self._unwrap_sam2(sam2_predictor)
+        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)
@@ -142,6 +206,10 @@ def __init__(self, sam2_predictor=None, matanyone_model=None):
         # Internal flags/state
         self._mat_bootstrapped = False
         self._alpha_prev: Optional[np.ndarray] = None  # temporal smoothing
+        
+        # Frame tracking
+        self.total_frames_processed = 0
+        self.frames_refined = 0
 
         # Quality selection at construction
         qname = os.getenv("BFX_QUALITY", "balanced").strip().lower()
@@ -149,8 +217,10 @@ def __init__(self, sam2_predictor=None, matanyone_model=None):
             qname = "balanced"
         self.quality = qname
         self.q = QUALITY_PROFILES[qname]
-        logger.info(f"TwoStageProcessor quality='{self.quality}' ⇒ {self.q}")
-
+        
+        # Log quality details
+        logger.info(f"TwoStageProcessor quality='{self.quality}' ⇒ refine_every={self.q['refine_stride']}, "
+                   f"spill={self.q['spill']:.2f}, mix={self.q['mix']:.2f}, bg_blur={self.q['bg_sigma']:.1f}")
         logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
 
     # --------------------------- internal utils ---------------------------
@@ -169,9 +239,12 @@ def _refresh_quality_from_env(self):
         if qname not in QUALITY_PROFILES:
             qname = "balanced"
         if qname != getattr(self, "quality", None) or not hasattr(self, "q"):
+            old_quality = self.quality
             self.quality = qname
             self.q = QUALITY_PROFILES[qname]
-            logger.info(f"Quality switched to '{self.quality}' ⇒ {self.q}")
+            logger.info(f"Quality switched from '{old_quality}' to '{self.quality}' ⇒ "
+                       f"refine_every={self.q['refine_stride']}, spill={self.q['spill']:.2f}, "
+                       f"mix={self.q['mix']:.2f}, bg_blur={self.q['bg_sigma']:.1f}")
 
     def _get_mask(self, frame: np.ndarray) -> np.ndarray:
         """Get segmentation mask using SAM2 (delegates to project helper)."""
@@ -276,7 +349,7 @@ def _soft_key_mask(self, frame_bgr: np.ndarray, key_bgr: np.ndarray, tol: int =
         alpha = 255.0 * np.clip((d - tol) / (tol*1.7), 0.0, 1.0)  # far from key = keep (255)
         return alpha.astype(np.uint8)
 
-    # --------------------- NEW: MatAnyone bootstrap ----------------------
+    # --------------------- MatAnyone bootstrap ----------------------
     def _bootstrap_matanyone_if_needed(self, frame_bgr: np.ndarray, coarse_mask: np.ndarray):
         """
         Call the MatAnyone session ONCE with the first coarse mask to initialize
@@ -294,6 +367,18 @@ def _bootstrap_matanyone_if_needed(self, frame_bgr: np.ndarray, coarse_mask: np.
         except Exception as e:
             logger.warning(f"MatAnyone bootstrap failed (continuing without): {e}")
 
+    def _should_refine_frame(self, frame_idx: int) -> bool:
+        """Check if current frame should be refined based on quality profile"""
+        if not self.matanyone:
+            return False
+        
+        # Always refine first frame for bootstrap
+        if frame_idx == 0:
+            return True
+            
+        stride = max(1, int(self.q.get("refine_stride", 3)))
+        return (frame_idx % stride) == 0
+
     # ---------------------------------------------------------------------
     # Stage 1 – Original → keyed (green/blue/…)  -- chooses colour on 1st frame
     # ---------------------------------------------------------------------
@@ -328,16 +413,20 @@ def _prog(p, d):
             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:
+            base_writer, out_path = create_video_writer(output_path, fps, w, h)
+            if base_writer is None:
                 cap.release()
                 return None, "Could not create output writer"
+            
+            # Use robust wrapper
+            writer = RobustVideoWriter(base_writer, out_path)
 
             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
+            self.frames_refined = 0
 
             solid_bg = np.zeros((h, w, 3), np.uint8)  # overwritten per-frame
 
@@ -373,12 +462,15 @@ def _prog(p, d):
                     logger.info(f"[TwoStage] Using key colour: {key_color_mode} → {chosen_bgr.tolist()}")
 
                 # --- Optional refinement via MatAnyone (profile cadence) ---
-                stride = int(self.q.get("refine_stride", 3))
-                if self.matanyone and (frame_idx % max(1, stride) == 0):
+                if self._should_refine_frame(frame_idx):
                     try:
                         mask = refine_mask_hq(frame, mask, self.matanyone, fallback_enabled=True)
+                        self.frames_refined += 1
+                        logger.debug(f"Frame {frame_idx}: Refined (quality={self.quality})")
                     except Exception as e:
                         logger.warning(f"MatAnyOne refine fail f={frame_idx}: {e}")
+                else:
+                    logger.debug(f"Frame {frame_idx}: Skipped refinement (cadence={self.q['refine_stride']})")
 
                 # --- Composite onto solid key colour ---
                 solid_bg[:] = chosen_bgr
@@ -389,23 +481,30 @@ def _prog(p, d):
 
                 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 '?'}")
+                _prog(pct, f"Stage 1: {frame_idx}/{total or '?'} (refined: {self.frames_refined})")
 
             cap.release()
             writer.release()
+            self.total_frames_processed = frame_idx
 
             # 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)
+                logger.info(f"Cached {len(masks)} masks to {cache_file}")
             except Exception as e:
                 logger.warning(f"mask cache save fail: {e}")
 
             _prog(1.0, "Stage 1: complete")
+            
+            # Log quality impact
+            logger.info(f"Stage 1 complete: {frame_idx} frames, {self.frames_refined} refined "
+                       f"({100*self.frames_refined/max(1,frame_idx):.1f}%)")
+            
             return (
                 {"path": out_path, "frames": frame_idx, "key_bgr": chosen_bgr.tolist()},
-                f"Green-screen video created ({frame_idx} frames)"
+                f"Green-screen video created ({frame_idx} frames, {self.frames_refined} refined)"
             )
 
         except Exception as e:
@@ -462,11 +561,15 @@ def _prog(p, d):
             sigma = float(self.q.get("bg_sigma", 0.0))
             if sigma > 0:
                 bg = cv2.GaussianBlur(bg, (0, 0), sigmaX=sigma, sigmaY=sigma)
+                logger.debug(f"Applied background blur: sigma={sigma:.1f}")
 
-            writer, out_path = create_video_writer(output_path, fps, w, h)
-            if writer is None:
+            base_writer, out_path = create_video_writer(output_path, fps, w, h)
+            if base_writer is None:
                 cap.release()
                 return None, "Could not create output writer"
+            
+            # Use robust wrapper
+            writer = RobustVideoWriter(base_writer, out_path)
 
             # Load cached masks if available
             masks = None
@@ -530,6 +633,11 @@ def _prog(p, d):
             writer.release()
 
             _prog(1.0, "Stage 2: complete")
+            
+            # Verify frame counts match
+            if total > 0 and frame_idx != total:
+                logger.warning(f"Frame count mismatch: processed {frame_idx}, expected {total}")
+            
             return out_path, f"Final composite created ({frame_idx} frames)"
 
         except Exception as e:
@@ -637,6 +745,9 @@ def _combined_progress(pct, desc):
             # Reset per-video state
             self._mat_bootstrapped = False
             self._alpha_prev = None
+            self.total_frames_processed = 0
+            self.frames_refined = 0
+            
             if self.matanyone is not None and hasattr(self.matanyone, "reset"):
                 try:
                     self.matanyone.reset()
@@ -670,8 +781,13 @@ def _combined_progress(pct, desc):
             except Exception:
                 pass
 
+            # Report quality impact
+            logger.info(f"Pipeline complete with quality='{self.quality}': "
+                       f"{self.total_frames_processed} frames, "
+                       f"{self.frames_refined} refined ({100*self.frames_refined/max(1,self.total_frames_processed):.1f}%)")
+
             return final_path, stage2_msg
 
         except Exception as e:
             logger.error(f"Full pipeline error: {e}\n{traceback.format_exc()}")
-            return None, f"Pipeline failed: {e}"
+            return None, f"Pipeline failed: {e}"