Update utilities.py

utilities.py

@@ -400,7 +400,7 @@ def refine_mask_hq(image: np.ndarray, mask: np.ndarray, matanyone_processor: Any
             raise MaskRefinementError(f"Unexpected error: {e}")
 
 def _matanyone_refine(image: np.ndarray, mask: np.ndarray, processor: Any) -> Optional[np.ndarray]:
-    """Attempt MatAnyone mask refinement"""
+    """Attempt MatAnyone mask refinement - Python 3.10 compatible"""
     try:
         # Different possible MatAnyone interfaces
         if hasattr(processor, 'infer'):
@@ -419,12 +419,121 @@ def _matanyone_refine(image: np.ndarray, mask: np.ndarray, processor: Any) -> Op
         # Process the refined mask
         refined_mask = _process_mask(refined_mask)
         
+        logger.debug("MatAnyone refinement successful")
         return refined_mask
         
     except Exception as e:
         logger.warning(f"MatAnyone processing error: {e}")
         return None
 
+def _background_matting_v2_refine(image: np.ndarray, mask: np.ndarray) -> Optional[np.ndarray]:
+    """Use BackgroundMattingV2 for mask refinement"""
+    try:
+        # Import BackgroundMattingV2 if available
+        from inference_images import inference_img
+        import torch
+        
+        # Convert inputs to proper format
+        image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+        mask_tensor = torch.from_numpy(mask).float() / 255.0
+        
+        # Create trimap from mask
+        trimap = _create_trimap_from_mask(mask)
+        trimap_tensor = torch.from_numpy(trimap).float()
+        
+        # Run inference
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        with torch.no_grad():
+            alpha = inference_img(
+                image_tensor.unsqueeze(0).to(device),
+                trimap_tensor.unsqueeze(0).unsqueeze(0).to(device)
+            )
+        
+        # Convert back to numpy
+        refined_mask = alpha.cpu().squeeze().numpy()
+        refined_mask = (refined_mask * 255).astype(np.uint8)
+        
+        logger.info("BackgroundMattingV2 refinement successful")
+        return refined_mask
+        
+    except ImportError:
+        logger.warning("BackgroundMattingV2 not available")
+        return None
+    except Exception as e:
+        logger.warning(f"BackgroundMattingV2 error: {e}")
+        return None
+
+def _rembg_refine(image: np.ndarray, mask: np.ndarray) -> Optional[np.ndarray]:
+    """Use rembg for mask refinement"""
+    try:
+        from rembg import remove, new_session
+        
+        # Use rembg to get a high-quality mask
+        session = new_session('u2net')
+        
+        # Convert image to PIL
+        from PIL import Image
+        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+        
+        # Remove background
+        output = remove(pil_image, session=session)
+        
+        # Extract alpha channel as mask
+        if output.mode == 'RGBA':
+            alpha = np.array(output)[:, :, 3]
+        else:
+            # Fallback: convert to grayscale
+            alpha = np.array(output.convert('L'))
+        
+        # Combine with original mask using weighted average
+        original_mask_norm = mask.astype(np.float32) / 255.0
+        rembg_mask_norm = alpha.astype(np.float32) / 255.0
+        
+        # Weighted combination: 70% rembg, 30% original
+        combined = 0.7 * rembg_mask_norm + 0.3 * original_mask_norm
+        combined = np.clip(combined * 255, 0, 255).astype(np.uint8)
+        
+        logger.info("Rembg refinement successful")
+        return combined
+        
+    except ImportError:
+        logger.warning("Rembg not available")
+        return None
+    except Exception as e:
+        logger.warning(f"Rembg error: {e}")
+        return None
+
+def _create_trimap_from_mask(mask: np.ndarray, erode_size: int = 10, dilate_size: int = 20) -> np.ndarray:
+    """Create trimap from binary mask for BackgroundMattingV2"""
+    try:
+        # Ensure mask is binary
+        _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+        
+        # Create trimap: 0 = background, 128 = unknown, 255 = foreground
+        trimap = np.zeros_like(mask, dtype=np.uint8)
+        
+        # Erode mask to get sure foreground
+        kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erode_size, erode_size))
+        sure_fg = cv2.erode(binary_mask, kernel_erode, iterations=1)
+        
+        # Dilate mask to get unknown region
+        kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_size, dilate_size))
+        unknown_region = cv2.dilate(binary_mask, kernel_dilate, iterations=1)
+        
+        # Set trimap values
+        trimap[sure_fg == 255] = 255  # Sure foreground
+        trimap[(unknown_region == 255) & (sure_fg == 0)] = 128  # Unknown
+        # Background remains 0
+        
+        return trimap
+        
+    except Exception as e:
+        logger.warning(f"Trimap creation failed: {e}")
+        # Return simple trimap based on original mask
+        trimap = np.where(mask > 127, 255, 0).astype(np.uint8)
+        return trimap
+
 def enhance_mask_opencv_advanced(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     """
     Advanced OpenCV-based mask enhancement with multiple techniques