Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -8,12 +8,13 @@
 - New: Return green screen path for monitoring
 - Fix: Force green key color
 - Fix: Use RobustVideoWriter to prevent frame loss
+- Fix: Add logging for background preparation issue
 """
 from __future__ import annotations
 import cv2, numpy as np, os, gc, pickle, logging, tempfile, traceback, threading
 from pathlib import Path
 from utils.cv_processing import segment_person_hq, refine_mask_hq
-from quality_manager import quality_manager  # New quality manager import
+from quality_manager import quality_manager # New quality manager import
 # Project logger if available
 try:
     from utils.logger import get_logger
@@ -52,22 +53,22 @@ def create_video_writer(output_path: str, fps: float, width: int, height: int, p
 # ---------------------------------------------------------------------------
 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:
@@ -78,7 +79,7 @@ def write(self, frame):
                 logger.warning(f"Frame write failed: {e}")
                 return False
         return True
-   
+  
     def release(self):
         """Flush remaining frames and close"""
         # Write any remaining buffered frames
@@ -89,7 +90,7 @@ def release(self):
                 self.frames_written += 1
             except Exception:
                 break
-       
+      
         # Duplicate last frame 3 times to force flush
         if self.frames_written > 0 and self.writer:
             last_frame = self.frame_buffer[-1] if self.frame_buffer else None
@@ -106,14 +107,14 @@ def release(self):
                 for _ in range(3):
                     self.writer.write(last_frame)
                     self.frames_written += 1
-       
+      
         # 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)
@@ -178,11 +179,11 @@ 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
-       
+      
         # Load quality profile
         self.q = quality_manager.get_params()
         logger.info(f"TwoStageProcessor init – SAM2: {self.sam2 is not None} | MatAnyOne: {self.matanyone is not None}")
@@ -307,11 +308,11 @@ 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
-           
+          
         return quality_manager.should_refine_frame(frame_idx)
     # ---------------------------------------------------------------------
     # Stage 1 – Original → keyed (green/blue/…) -- chooses colour on 1st frame
@@ -333,7 +334,7 @@ def _prog(p, d):
                     pass
         try:
             # pick up any new quality selection
-            quality_manager.load_profile()  # Refresh
+            quality_manager.load_profile() # Refresh
             self.q = quality_manager.get_params()
             _prog(0.0, "Stage 1: opening video…")
             cap = cv2.VideoCapture(video_path)
@@ -347,7 +348,7 @@ def _prog(p, d):
             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
@@ -414,11 +415,11 @@ def _prog(p, d):
             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, {self.frames_refined} refined)"
@@ -462,8 +463,9 @@ def _prog(p, d):
             if isinstance(background, str):
                 bg = cv2.imread(background)
                 if bg is None:
+                    logger.error(f"Failed to load background image from {background}")
                     cap.release()
-                    return None, "Could not load background image"
+                    return None, "Could not load background image - check file path and permissions"
                 bg = cv2.resize(bg, (w, h))
             else:
                 bg = cv2.resize(background, (w, h))
@@ -476,7 +478,7 @@ def _prog(p, d):
             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
@@ -532,11 +534,11 @@ def _prog(p, d):
             cap.release()
             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:
             logger.error(f"Stage 2 error: {e}\n{traceback.format_exc()}")
@@ -631,7 +633,7 @@ def _combined_progress(pct, desc):
             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()