Update two_stage_processor.py

two_stage_processor.py

@@ -11,6 +11,7 @@
 import pickle
 import logging
 from pathlib import Path
+import tempfile
 
 logger = logging.getLogger(__name__)
 
@@ -83,8 +84,11 @@ def _prog(pct: float, desc: str):
             
             # Save masks for potential reuse
             mask_file = self.mask_cache_dir / f"{Path(output_path).stem}_masks.pkl"
-            with open(mask_file, 'wb') as f:
-                pickle.dump(masks, f)
+            try:
+                with open(mask_file, 'wb') as f:
+                    pickle.dump(masks, f)
+            except Exception as e:
+                logger.warning(f"Failed to save masks: {e}")
             
             _prog(1.0, "Stage 1 complete: Green screen created")
             return output_path, f"Green screen created: {frame_count} frames"
@@ -209,6 +213,7 @@ def _extract_person_mask(self, frame):
     def _refine_mask(self, frame, mask):
         """Refine mask using MatAnyone if available"""
         if self.matanyone_model is None:
+            logger.warning("MatAnyone not available, using original mask")
             return mask
         
         try:
@@ -216,6 +221,7 @@ def _refine_mask(self, frame, mask):
             # This would depend on your MatAnyone implementation
             return mask
         except:
+            logger.warning("MatAnyone refinement failed, using original mask")
             return mask
     
     def _apply_greenscreen_hard(self, frame, mask, green_bg):
@@ -277,8 +283,6 @@ def process_full_pipeline(self, video_path, background, final_output,
         """
         Run the complete two-stage pipeline
         """
-        import tempfile
-        
         # Stage 1: Create green screen
         greenscreen_path = tempfile.mktemp(suffix='_greenscreen.mp4')
         gs_result, gs_msg = self.stage1_extract_to_greenscreen(