Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -235,94 +235,250 @@ def _create_trimap(self, mask_path: str) -> str:
     
     def _process_video_to_greenscreen(self, video_path: str, trimap_path: str, trim_seconds: float) -> str:
         """
-        Process video to create green screen version.
+        Process video to create green screen version using robust MatAnyone adapter.
         
-        KEY: This creates RGB video with GREEN BACKGROUND, not alpha matte.
+        Uses the ChatGPT-provided adapter that handles all interface variations.
         """
         
-        logger.info("Processing video to green screen format...")
+        logger.info("Processing video with robust MatAnyone adapter...")
+        
+        # Import the adapter (assuming it's in the same module or available)
+        from .matanyone_adapter import run_matanyone_with_reference_mask
+        
+        # Trim video first
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        target_frames = min(int(trim_seconds * fps), total_frames)
+        cap.release()
+        
+        # Create trimmed video
+        trimmed_video_path = os.path.join(self.temp_dir, "trimmed_for_matanyone.mp4")
+        with VideoFileClip(video_path) as clip:
+            trimmed_clip = clip.subclip(0, trim_seconds)
+            trimmed_clip.write_videofile(
+                trimmed_video_path,
+                codec='libx264',
+                audio_codec='aac',
+                verbose=False,
+                logger=None
+            )
+        
+        # Prepare MatAnyone output directory
+        matanyone_output_dir = os.path.join(self.temp_dir, "matanyone_output")
+        os.makedirs(matanyone_output_dir, exist_ok=True)
+        
+        try:
+            logger.info("Using robust MatAnyone adapter...")
+            logger.info(f"Video: {trimmed_video_path}")
+            logger.info(f"Trimap: {trimap_path}")
+            logger.info(f"Output dir: {matanyone_output_dir}")
+            
+            # Use the robust adapter that handles all interface variations
+            alpha_dir = run_matanyone_with_reference_mask(
+                inference_core=self.matanyone_processor,
+                video_path=trimmed_video_path,
+                reference_mask_path=trimap_path,
+                output_dir=matanyone_output_dir,
+                write_png=True
+            )
+            
+            if not alpha_dir or not os.path.exists(alpha_dir):
+                raise RuntimeError(f"MatAnyone adapter failed to create alpha directory: {alpha_dir}")
+            
+            # Convert PNG alpha sequence to green screen video
+            green_screen_path = self._convert_png_alpha_to_greenscreen(
+                trimmed_video_path, alpha_dir
+            )
+            
+            logger.info(f"Green screen video created: {green_screen_path}")
+            return green_screen_path
+            
+        except Exception as e:
+            logger.error(f"MatAnyone adapter failed: {e}")
+            logger.error(f"Falling back to SAM2-only green screen...")
+            
+            # Fallback: create green screen using just trimap
+            return self._create_sam2_only_greenscreen(trimmed_video_path, trimap_path)
+    
+    def _convert_png_alpha_to_greenscreen(self, original_video_path: str, alpha_dir: str) -> str:
+        """Convert PNG alpha sequence to green screen video."""
+        
+        logger.info("Converting PNG alpha sequence to green screen format...")
+        
+        # Get PNG files
+        png_files = sorted([f for f in os.listdir(alpha_dir) if f.startswith("alpha_") and f.endswith(".png")])
+        if not png_files:
+            raise RuntimeError(f"No alpha PNG files found in {alpha_dir}")
+        
+        # Open original video
+        original_cap = cv2.VideoCapture(original_video_path)
+        if not original_cap.isOpened():
+            raise ValueError("Could not open original video for conversion")
+        
+        # Get video properties
+        fps = original_cap.get(cv2.CAP_PROP_FPS)
+        width = int(original_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(original_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        # Create output video writer
+        green_screen_path = os.path.join(self.temp_dir, "greenscreen_video.mp4")
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(green_screen_path, fourcc, fps, (width, height))
+        
+        frame_count = 0
+        while True:
+            ret, original_frame = original_cap.read()
+            if not ret:
+                break
+            
+            # Load corresponding alpha PNG
+            if frame_count < len(png_files):
+                alpha_path = os.path.join(alpha_dir, png_files[frame_count])
+                alpha_frame = cv2.imread(alpha_path, cv2.IMREAD_GRAYSCALE)
+                
+                if alpha_frame is not None:
+                    # Resize alpha to match video if needed
+                    if alpha_frame.shape[:2] != (height, width):
+                        alpha_frame = cv2.resize(alpha_frame, (width, height))
+                    
+                    alpha_matte = alpha_frame.astype(np.float32) / 255.0
+                else:
+                    # Fallback alpha if PNG read fails
+                    alpha_matte = np.ones((height, width), dtype=np.float32)
+            else:
+                # Use last alpha if we run out of PNG files
+                alpha_matte = np.ones((height, width), dtype=np.float32)
+            
+            # Create green screen frame
+            green_screen_frame = self._create_greenscreen_frame(original_frame, alpha_matte)
+            out.write(green_screen_frame)
+            
+            frame_count += 1
+            if frame_count % 25 == 0:
+                logger.info(f"Converted {frame_count} frames to green screen")
+        
+        original_cap.release()
+        out.release()
+        
+        logger.info(f"PNG alpha to green screen conversion complete: {green_screen_path}")
+        return green_screen_path
+    
+    def _find_alpha_video(self, output_dir: str, input_video_path: str) -> Optional[str]:
+        """Find the alpha video file created by MatAnyone."""
+        input_name = Path(input_video_path).stem
+        
+        patterns = [
+            f"{input_name}_pha.mp4",
+            f"{input_name}_alpha.mp4", 
+            "alpha.mp4",
+            "output_alpha.mp4",
+            "pha.mp4"
+        ]
+        
+        for pattern in patterns:
+            potential_path = os.path.join(output_dir, pattern)
+            if os.path.exists(potential_path):
+                return potential_path
+        
+        # Fallback: find any .mp4 file in output directory
+        for file in os.listdir(output_dir):
+            if file.endswith('.mp4'):
+                return os.path.join(output_dir, file)
+        
+        return None
+    
+    def _convert_alpha_to_greenscreen(self, original_video_path: str, alpha_video_path: str) -> str:
+        """Convert MatAnyone alpha video to green screen video."""
+        
+        logger.info("Converting alpha video to green screen format...")
+        
+        original_cap = cv2.VideoCapture(original_video_path)
+        alpha_cap = cv2.VideoCapture(alpha_video_path)
+        
+        if not original_cap.isOpened() or not alpha_cap.isOpened():
+            raise ValueError("Could not open video files for conversion")
+        
+        # Get video properties
+        fps = original_cap.get(cv2.CAP_PROP_FPS)
+        width = int(original_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(original_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        # Create output video writer
+        green_screen_path = os.path.join(self.temp_dir, "greenscreen_video.mp4")
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(green_screen_path, fourcc, fps, (width, height))
+        
+        frame_count = 0
+        while True:
+            ret_orig, original_frame = original_cap.read()
+            ret_alpha, alpha_frame = alpha_cap.read()
+            
+            if not ret_orig or not ret_alpha:
+                break
+            
+            # Process alpha
+            if len(alpha_frame.shape) == 3:
+                alpha_matte = np.mean(alpha_frame, axis=2)
+            else:
+                alpha_matte = alpha_frame
+            
+            alpha_matte = alpha_matte.astype(np.float32) / 255.0
+            alpha_matte = np.clip(alpha_matte, 0, 1)
+            
+            # Create green screen frame
+            green_screen_frame = self._create_greenscreen_frame(original_frame, alpha_matte)
+            out.write(green_screen_frame)
+            
+            frame_count += 1
+            if frame_count % 25 == 0:
+                logger.info(f"Converted {frame_count} frames to green screen")
+        
+        original_cap.release()
+        alpha_cap.release()
+        out.release()
+        
+        logger.info(f"Alpha to green screen conversion complete: {green_screen_path}")
+        return green_screen_path
+    
+    def _create_sam2_only_greenscreen(self, video_path: str, trimap_path: str) -> str:
+        """Fallback: create green screen video using only SAM2 trimap."""
+        
+        logger.info("Creating SAM2-only green screen video (fallback)...")
         
-        # Load trimap
         trimap = cv2.imread(trimap_path, cv2.IMREAD_GRAYSCALE)
         if trimap is None:
             raise ValueError(f"Could not load trimap: {trimap_path}")
         
-        # Open input video
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise ValueError(f"Could not open video: {video_path}")
         
-        # Get video properties
         fps = cap.get(cv2.CAP_PROP_FPS)
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        target_frames = min(int(trim_seconds * fps), total_frames)
         
-        # Resize trimap to match video
         trimap_resized = cv2.resize(trimap, (width, height))
+        alpha_matte = (trimap_resized > 128).astype(np.float32)
         
-        # Create output video writer
         green_screen_path = os.path.join(self.temp_dir, "greenscreen_video.mp4")
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(green_screen_path, fourcc, fps, (width, height))
         
-        logger.info(f"Creating green screen video: {target_frames} frames")
-        
         frame_count = 0
-        while frame_count < target_frames:
+        while True:
             ret, frame = cap.read()
             if not ret:
                 break
             
-            try:
-                # Convert to RGB for MatAnyone
-                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                
-                # FIXED: Use correct MatAnyone InferenceCore API
-                try:
-                    # Try different MatAnyone interface methods
-                    if hasattr(self.matanyone_processor, 'process'):
-                        alpha_matte = self.matanyone_processor.process(frame_rgb, trimap_resized)
-                    elif hasattr(self.matanyone_processor, 'predict'):
-                        alpha_matte = self.matanyone_processor.predict(frame_rgb, trimap_resized)
-                    elif hasattr(self.matanyone_processor, '__call__'):
-                        alpha_matte = self.matanyone_processor(frame_rgb, trimap_resized)
-                    else:
-                        # If no known method, try the most common ones
-                        logger.error(f"MatAnyone methods: {[m for m in dir(self.matanyone_processor) if not m.startswith('_')]}")
-                        raise AttributeError("No known MatAnyone processing method found")
-                        
-                except Exception as matanyone_error:
-                    logger.error(f"MatAnyone processing failed: {matanyone_error}")
-                    # Use trimap as fallback
-                    alpha_matte = (trimap_resized > 128).astype(np.float32)
-                
-                if len(alpha_matte.shape) == 3:
-                    alpha_matte = np.mean(alpha_matte, axis=2)
-                
-                alpha_matte = np.clip(alpha_matte, 0, 1)
-                
-                # Create green screen frame
-                green_screen_frame = self._create_greenscreen_frame(frame, alpha_matte)
-                
-                out.write(green_screen_frame)
-                
-            except Exception as e:
-                logger.error(f"Frame processing failed on frame {frame_count}: {e}")
-                # Fallback: use trimap
-                fallback_alpha = (trimap_resized > 128).astype(np.float32)
-                green_screen_frame = self._create_greenscreen_frame(frame, fallback_alpha)
-                out.write(green_screen_frame)
-            
+            green_screen_frame = self._create_greenscreen_frame(frame, alpha_matte)
+            out.write(green_screen_frame)
             frame_count += 1
-            if frame_count % 25 == 0:
-                logger.info(f"Processed {frame_count}/{target_frames} frames")
         
         cap.release()
         out.release()
         
-        logger.info(f"Green screen video created: {green_screen_path}")
+        logger.info(f"SAM2-only green screen created: {green_screen_path}")
         return green_screen_path
     
     def _create_greenscreen_frame(self, original_frame: np.ndarray, alpha_matte: np.ndarray) -> np.ndarray: