Update utils/refinement.py

utils/refinement.py

@@ -5,7 +5,7 @@
 """
 
 from __future__ import annotations
-from typing import Any, Optional, Tuple
+from typing import Any, Optional, Tuple, List
 import logging
 
 import cv2
@@ -26,11 +26,12 @@ class MaskRefinementError(Exception):
 # ============================================================================
 __all__ = [
     "refine_mask_hq",
+    "refine_masks_batch",
     "MaskRefinementError",
 ]
 
 # ============================================================================
-# MAIN API
+# MAIN API - SINGLE FRAME
 # ============================================================================
 def refine_mask_hq(
     image: np.ndarray,
@@ -77,7 +78,50 @@ def refine_mask_hq(
     return mask
 
 # ============================================================================
-# AI-BASED REFINEMENT
+# BATCH PROCESSING FOR TEMPORAL CONSISTENCY
+# ============================================================================
+def refine_masks_batch(
+    frames: List[np.ndarray],
+    masks: List[np.ndarray],
+    matanyone_model: Optional[Any] = None,
+    fallback_enabled: bool = True
+) -> List[np.ndarray]:
+    """
+    Refine multiple masks using MatAnyone's temporal consistency.
+    
+    Args:
+        frames: List of BGR images
+        masks: List of initial binary masks
+        matanyone_model: MatAnyone InferenceCore model
+        fallback_enabled: Whether to use fallback methods
+    
+    Returns:
+        List of refined binary masks
+    """
+    if not frames or not masks:
+        return masks
+    
+    if len(frames) != len(masks):
+        raise MaskRefinementError(f"Frame count {len(frames)} doesn't match mask count {len(masks)}")
+    
+    if matanyone_model is not None:
+        try:
+            refined = _refine_batch_with_matanyone(frames, masks, matanyone_model)
+            # Validate all masks
+            if all(_validate_refined_mask(r, m) for r, m in zip(refined, masks)):
+                return refined
+            log.warning("Batch MatAnyone refinement failed validation")
+        except Exception as e:
+            log.warning(f"Batch MatAnyone refinement failed: {e}")
+    
+    # Fallback to frame-by-frame classical refinement
+    if fallback_enabled:
+        return [_classical_refinement(f, m) for f, m in zip(frames, masks)]
+    
+    return masks
+
+# ============================================================================
+# AI-BASED REFINEMENT - SINGLE FRAME
 # ============================================================================
 def _refine_with_matanyone(
     image: np.ndarray,
@@ -86,7 +130,6 @@ def _refine_with_matanyone(
 ) -> np.ndarray:
     """Use MatAnyone model for mask refinement."""
     try:
-        # MatAnyone's InferenceCore expects torch tensors
         # Convert BGR to RGB and normalize
         image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         h, w = image_rgb.shape[:2]
@@ -105,74 +148,42 @@ def _refine_with_matanyone(
         mask_tensor = torch.from_numpy(mask).float() / 255.0
         mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
         
-        # MatAnyone InferenceCore workflow for single frame
-        # The model should have been initialized as InferenceCore(matanyone_model)
+        # Try different methods on InferenceCore
         result = None
         
-        if hasattr(model, 'process_frame'):
-            # Single frame processing method
-            with torch.no_grad():
-                result = model.process_frame(image_tensor, mask_tensor)
-        elif hasattr(model, 'step'):
-            # Step method for iterative processing
-            with torch.no_grad():
-                # Initialize memory with first frame
-                model.reset()
-                # Process frame with mask
+        # Log available methods for debugging
+        methods = [m for m in dir(model) if not m.startswith('_')]
+        log.debug(f"MatAnyone InferenceCore methods: {methods}")
+        
+        with torch.no_grad():
+            if hasattr(model, 'step'):
+                # Step method for iterative processing (don't call reset)
                 result = model.step(image_tensor, mask_tensor)
-        elif hasattr(model, 'forward'):
-            # Direct forward pass
-            with torch.no_grad():
+            elif hasattr(model, 'process_frame'):
+                result = model.process_frame(image_tensor, mask_tensor)
+            elif hasattr(model, 'forward'):
                 result = model.forward(image_tensor, mask_tensor)
-        elif hasattr(model, 'predict'):
-            # Predict method
-            with torch.no_grad():
-                result = model.predict(image_tensor, mask_tensor)
-        elif hasattr(model, '__call__'):
-            # Callable model
-            with torch.no_grad():
+            elif hasattr(model, '__call__'):
                 result = model(image_tensor, mask_tensor)
-        else:
-            # Try to find any method that might work
-            methods = [m for m in dir(model) if not m.startswith('_')]
-            processing_methods = [m for m in methods if any(keyword in m.lower() 
-                                 for keyword in ['process', 'refine', 'matte', 'alpha', 'predict'])]
-            if processing_methods:
-                method = getattr(model, processing_methods[0])
-                with torch.no_grad():
-                    result = method(image_tensor, mask_tensor)
             else:
-                raise MaskRefinementError(f"MatAnyone model has no recognized processing method. Available methods: {methods}")
+                raise MaskRefinementError(f"No recognized method. Available: {methods}")
         
         if result is None:
             raise MaskRefinementError("MatAnyone returned None")
         
-        # Handle different return types
-        if isinstance(result, tuple) or isinstance(result, list):
-            # Extract alpha matte from tuple/list result
-            alpha = result[0] if len(result) > 0 else None
-        elif isinstance(result, dict):
-            # Extract from dictionary result
-            alpha = result.get('alpha', result.get('matte', result.get('mask', None)))
-        else:
-            alpha = result
-        
-        if alpha is None:
-            raise MaskRefinementError("Could not extract alpha matte from MatAnyone result")
+        # Extract alpha matte from result
+        alpha = _extract_alpha_from_result(result)
         
-        # Convert back to numpy
+        # Convert back to numpy and resize if needed
         if isinstance(alpha, torch.Tensor):
-            alpha = alpha.squeeze().cpu().numpy()  # Remove batch dimensions
-            
-        # Ensure proper shape
+            alpha = alpha.squeeze().cpu().numpy()
+        
         if alpha.ndim == 3:
             alpha = alpha[0] if alpha.shape[0] == 1 else alpha.mean(axis=0)
-            
-        # Convert to uint8
+        
         if alpha.dtype != np.uint8:
             alpha = (alpha * 255).clip(0, 255).astype(np.uint8)
-            
-        # Resize if needed
+        
         if alpha.shape != (h, w):
             alpha = cv2.resize(alpha, (w, h), interpolation=cv2.INTER_LINEAR)
         
@@ -183,44 +194,121 @@ def _refine_with_matanyone(
         raise MaskRefinementError(f"MatAnyone processing failed: {str(e)}")
 
 # ============================================================================
-# CLASSICAL REFINEMENT
+# AI-BASED REFINEMENT - BATCH
 # ============================================================================
-def _classical_refinement(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
-    """Apply classical CV techniques for mask refinement."""
-    refined = mask.copy()
-    
-    # 1. Morphological operations to clean up
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-    refined = cv2.morphologyEx(refined, cv2.MORPH_CLOSE, kernel)
-    refined = cv2.morphologyEx(refined, cv2.MORPH_OPEN, kernel)
+def _refine_batch_with_matanyone(
+    frames: List[np.ndarray],
+    masks: List[np.ndarray],
+    model: Any
+) -> List[np.ndarray]:
+    """Process batch of frames through MatAnyone for temporal consistency."""
+    try:
+        batch_size = len(frames)
+        h, w = frames[0].shape[:2]
+        
+        # Convert frames to tensor batch
+        frame_tensors = []
+        for frame in frames:
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).float() / 255.0
+            frame_tensors.append(tensor)
+        
+        # Stack into batch (N, C, H, W)
+        batch_tensor = torch.stack(frame_tensors)
+        
+        # Prepare first mask for initialization
+        first_mask = masks[0]
+        if first_mask.dtype != np.uint8:
+            first_mask = (first_mask * 255).astype(np.uint8)
+        if first_mask.ndim == 3:
+            first_mask = cv2.cvtColor(first_mask, cv2.COLOR_BGR2GRAY)
+        
+        # Convert first mask to tensor
+        first_mask_tensor = torch.from_numpy(first_mask).float() / 255.0
+        first_mask_tensor = first_mask_tensor.unsqueeze(0).unsqueeze(0)
+        
+        refined_masks = []
+        
+        with torch.no_grad():
+            # Check for batch processing methods
+            if hasattr(model, 'process_batch'):
+                # Direct batch processing
+                results = model.process_batch(batch_tensor, first_mask_tensor)
+                for result in results:
+                    alpha = _extract_alpha_from_result(result)
+                    refined_masks.append(_tensor_to_mask(alpha, h, w))
+                    
+            elif hasattr(model, 'step'):
+                # Process frames sequentially with memory
+                for i, frame_tensor in enumerate(frame_tensors):
+                    if i == 0:
+                        # First frame with mask
+                        result = model.step(frame_tensor.unsqueeze(0), first_mask_tensor)
+                    else:
+                        # Subsequent frames use memory from previous
+                        result = model.step(frame_tensor.unsqueeze(0), None)
+                    
+                    alpha = _extract_alpha_from_result(result)
+                    refined_masks.append(_tensor_to_mask(alpha, h, w))
+                    
+            else:
+                # Fallback to processing each frame with its mask
+                log.warning("MatAnyone batch processing not available, using frame-by-frame")
+                for frame_tensor, mask in zip(frame_tensors, masks):
+                    mask_tensor = torch.from_numpy(mask).float() / 255.0
+                    mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
+                    
+                    result = model(frame_tensor.unsqueeze(0), mask_tensor)
+                    alpha = _extract_alpha_from_result(result)
+                    refined_masks.append(_tensor_to_mask(alpha, h, w))
+        
+        return refined_masks
+        
+    except Exception as e:
+        log.error(f"Batch MatAnyone processing error: {str(e)}")
+        raise MaskRefinementError(f"Batch processing failed: {str(e)}")
+
+# ============================================================================
+# HELPER FUNCTIONS
+# ============================================================================
+def _extract_alpha_from_result(result):
+    """Extract alpha matte from various result formats."""
+    if isinstance(result, (tuple, list)):
+        return result[0] if len(result) > 0 else None
+    elif isinstance(result, dict):
+        return result.get('alpha', result.get('matte', result.get('mask', None)))
+    else:
+        return result
+
+def _tensor_to_mask(tensor, target_h, target_w):
+    """Convert tensor to numpy mask with proper sizing."""
+    if isinstance(tensor, torch.Tensor):
+        mask = tensor.squeeze().cpu().numpy()
+    else:
+        mask = tensor
     
-    # 2. Edge-aware smoothing
-    refined = _edge_aware_smooth(image, refined)
+    if mask.ndim == 3:
+        mask = mask[0] if mask.shape[0] == 1 else mask.mean(axis=0)
     
-    # 3. Feather edges slightly
-    refined = _feather_edges(refined, radius=3)
+    if mask.dtype != np.uint8:
+        mask = (mask * 255).clip(0, 255).astype(np.uint8)
     
-    # 4. Remove small disconnected components
-    refined = _remove_small_components(refined, min_area_ratio=0.005)
+    if mask.shape != (target_h, target_w):
+        mask = cv2.resize(mask, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
     
-    return refined
+    return mask
 
-# ============================================================================
-# HELPER FUNCTIONS
-# ============================================================================
 def _validate_refined_mask(refined: np.ndarray, original: np.ndarray) -> bool:
     """Check if refined mask is reasonable."""
     if refined is None or refined.size == 0:
         return False
     
-    # Check if mask has reasonable coverage
     refined_area = np.sum(refined > 127)
     original_area = np.sum(original > 127)
     
     if refined_area == 0:
         return False
     
-    # Allow some variation but not extreme changes
     ratio = refined_area / max(original_area, 1)
     return 0.5 <= ratio <= 2.0
 
@@ -239,41 +327,45 @@ def _process_mask(mask: np.ndarray) -> np.ndarray:
     _, binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
     return binary
 
+# ============================================================================
+# CLASSICAL REFINEMENT
+# ============================================================================
+def _classical_refinement(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """Apply classical CV techniques for mask refinement."""
+    refined = mask.copy()
+    
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    refined = cv2.morphologyEx(refined, cv2.MORPH_CLOSE, kernel)
+    refined = cv2.morphologyEx(refined, cv2.MORPH_OPEN, kernel)
+    refined = _edge_aware_smooth(image, refined)
+    refined = _feather_edges(refined, radius=3)
+    refined = _remove_small_components(refined, min_area_ratio=0.005)
+    
+    return refined
+
 def _edge_aware_smooth(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     """Apply edge-aware smoothing using guided filter."""
-    # Convert to float for processing
     mask_float = mask.astype(np.float32) / 255.0
-    
-    # Simple guided filter approximation
     radius = 5
     eps = 0.01
     
-    # Use image as guide
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
     
-    # Box filter for mean
     mean_I = cv2.boxFilter(gray, -1, (radius, radius))
     mean_p = cv2.boxFilter(mask_float, -1, (radius, radius))
     mean_Ip = cv2.boxFilter(gray * mask_float, -1, (radius, radius))
     
-    # Covariance
     cov_Ip = mean_Ip - mean_I * mean_p
-    
-    # Variance
     mean_II = cv2.boxFilter(gray * gray, -1, (radius, radius))
     var_I = mean_II - mean_I * mean_I
     
-    # Coefficients
     a = cov_Ip / (var_I + eps)
     b = mean_p - a * mean_I
     
-    # Filter
     mean_a = cv2.boxFilter(a, -1, (radius, radius))
     mean_b = cv2.boxFilter(b, -1, (radius, radius))
     
     refined = mean_a * gray + mean_b
-    
-    # Convert back to binary
     return (refined * 255).clip(0, 255).astype(np.uint8)
 
 def _feather_edges(mask: np.ndarray, radius: int = 3) -> np.ndarray:
@@ -281,10 +373,8 @@ def _feather_edges(mask: np.ndarray, radius: int = 3) -> np.ndarray:
     if radius <= 0:
         return mask
     
-    # Blur then threshold to maintain binary nature
     blurred = cv2.GaussianBlur(mask, (radius*2+1, radius*2+1), radius/2)
     _, binary = cv2.threshold(blurred, 127, 255, cv2.THRESH_BINARY)
-    
     return binary
 
 def _remove_small_components(mask: np.ndarray, min_area_ratio: float = 0.005) -> np.ndarray:
@@ -294,18 +384,15 @@ def _remove_small_components(mask: np.ndarray, min_area_ratio: float = 0.005) ->
     if num_labels <= 1:
         return mask
     
-    # Calculate minimum area
     total_area = mask.shape[0] * mask.shape[1]
     min_area = int(total_area * min_area_ratio)
     
-    # Find largest component (excluding background)
     areas = stats[1:, cv2.CC_STAT_AREA]
     if len(areas) == 0:
         return mask
     
     max_label = np.argmax(areas) + 1
     
-    # Keep only components above threshold or the largest one
     cleaned = np.zeros_like(mask)
     for label in range(1, num_labels):
         if stats[label, cv2.CC_STAT_AREA] >= min_area or label == max_label: