Update utils/refinement/mask_refiner.py

utils/refinement/mask_refiner.py

@@ -0,0 +1,259 @@
+"""
+Complete utils/__init__.py with all required functions
+Provides direct implementations to avoid import recursion
+"""
+
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import logging
+from typing import Optional, Tuple, Dict, Any, List
+import tempfile
+import os
+
+logger = logging.getLogger(__name__)
+
+# Professional backgrounds configuration
+PROFESSIONAL_BACKGROUNDS = {
+    "office": {"color": (240, 248, 255), "gradient": True},
+    "studio": {"color": (32, 32, 32), "gradient": False},
+    "nature": {"color": (34, 139, 34), "gradient": True},
+    "abstract": {"color": (75, 0, 130), "gradient": True},
+    "white": {"color": (255, 255, 255), "gradient": False},
+    "black": {"color": (0, 0, 0), "gradient": False}
+}
+
+def validate_video_file(video_path: str) -> bool:
+    """Validate if video file is readable"""
+    try:
+        if not os.path.exists(video_path):
+            return False
+        
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return False
+            
+        ret, frame = cap.read()
+        cap.release()
+        return ret and frame is not None
+        
+    except Exception as e:
+        logger.error(f"Video validation failed: {e}")
+        return False
+
+def segment_person_hq(frame: np.ndarray, use_sam2: bool = True) -> Optional[np.ndarray]:
+    """High-quality person segmentation using SAM2 or fallback methods"""
+    try:
+        if use_sam2:
+            # Try SAM2 segmentation
+            try:
+                from sam2.sam2_image_predictor import SAM2ImagePredictor
+                from sam2.build_sam import build_sam2
+                from huggingface_hub import hf_hub_download
+                
+                # Load SAM2 model
+                sam_checkpoint = hf_hub_download("facebook/sam2-hiera-base-plus", "sam2_hiera_b+.pt")
+                sam_model = build_sam2(model_name='sam2_hiera_base_plus_t', ckpt_path=sam_checkpoint)
+                predictor = SAM2ImagePredictor(sam_model)
+                
+                # Set image and predict
+                predictor.set_image(frame)
+                
+                # Use center point as prompt (assuming person is in center)
+                h, w = frame.shape[:2]
+                center_point = np.array([[w//2, h//2]])
+                center_label = np.array([1])
+                
+                masks, scores, _ = predictor.predict(
+                    point_coords=center_point,
+                    point_labels=center_label,
+                    multimask_output=False
+                )
+                
+                return masks[0] if len(masks) > 0 else None
+                
+            except Exception as e:
+                logger.warning(f"SAM2 segmentation failed: {e}, falling back to simple method")
+        
+        # Fallback: Simple person detection using background subtraction
+        return _simple_person_segmentation(frame)
+        
+    except Exception as e:
+        logger.error(f"Person segmentation failed: {e}")
+        return None
+
+def _simple_person_segmentation(frame: np.ndarray) -> np.ndarray:
+    """Simple person segmentation using color-based methods"""
+    # Convert to HSV for better color detection
+    hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
+    
+    # Create mask for common background colors (green screen, white, etc.)
+    # Green screen detection
+    lower_green = np.array([40, 40, 40])
+    upper_green = np.array([80, 255, 255])
+    green_mask = cv2.inRange(hsv, lower_green, upper_green)
+    
+    # White background detection
+    lower_white = np.array([0, 0, 200])
+    upper_white = np.array([180, 30, 255])
+    white_mask = cv2.inRange(hsv, lower_white, upper_white)
+    
+    # Combine masks
+    bg_mask = cv2.bitwise_or(green_mask, white_mask)
+    
+    # Invert to get person mask
+    person_mask = cv2.bitwise_not(bg_mask)
+    
+    # Clean up mask with morphological operations
+    kernel = np.ones((5, 5), np.uint8)
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_CLOSE, kernel)
+    person_mask = cv2.morphologyEx(person_mask, cv2.MORPH_OPEN, kernel)
+    
+    # Convert to float and normalize
+    return person_mask.astype(np.float32) / 255.0
+
+def refine_mask_hq(mask: np.ndarray, frame: np.ndarray, use_matanyone: bool = True) -> np.ndarray:
+    """High-quality mask refinement using MatAnyone or fallback methods"""
+    try:
+        if use_matanyone:
+            try:
+                from matanyone import InferenceCore
+                
+                # Initialize MatAnyone
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                processor = InferenceCore(model_name="PeiqingYang/MatAnyone-v1.0", device=device)
+                
+                # Convert inputs to PIL Images
+                frame_pil = Image.fromarray(frame.astype(np.uint8))
+                mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
+                
+                # Refine mask
+                refined_mask = processor.infer(frame_pil, mask_pil)
+                
+                # Convert back to numpy
+                return np.array(refined_mask).astype(np.float32) / 255.0
+                
+            except Exception as e:
+                logger.warning(f"MatAnyone refinement failed: {e}, using simple refinement")
+        
+        # Fallback: Simple mask refinement
+        return _simple_mask_refinement(mask, frame)
+        
+    except Exception as e:
+        logger.error(f"Mask refinement failed: {e}")
+        return mask
+
+def _simple_mask_refinement(mask: np.ndarray, frame: np.ndarray) -> np.ndarray:
+    """Simple mask refinement using OpenCV operations"""
+    # Convert mask to uint8
+    mask_uint8 = (mask * 255).astype(np.uint8)
+    
+    # Apply Gaussian blur for smoother edges
+    mask_blurred = cv2.GaussianBlur(mask_uint8, (5, 5), 0)
+    
+    # Apply bilateral filter to preserve edges while smoothing
+    mask_refined = cv2.bilateralFilter(mask_blurred, 9, 75, 75)
+    
+    # Convert back to float
+    return mask_refined.astype(np.float32) / 255.0
+
+def replace_background_hq(frame: np.ndarray, mask: np.ndarray, background: np.ndarray) -> np.ndarray:
+    """High-quality background replacement with proper compositing"""
+    try:
+        # Ensure all inputs are the same size
+        h, w = frame.shape[:2]
+        background_resized = cv2.resize(background, (w, h))
+        
+        # Ensure mask has 3 channels
+        if len(mask.shape) == 2:
+            mask_3d = np.stack([mask] * 3, axis=-1)
+        else:
+            mask_3d = mask
+        
+        # Apply feathering to mask edges for smoother blending
+        mask_feathered = _apply_feathering(mask_3d)
+        
+        # Composite the image
+        result = frame * mask_feathered + background_resized * (1 - mask_feathered)
+        
+        return result.astype(np.uint8)
+        
+    except Exception as e:
+        logger.error(f"Background replacement failed: {e}")
+        return frame
+
+def _apply_feathering(mask: np.ndarray, feather_amount: int = 3) -> np.ndarray:
+    """Apply feathering to mask edges for smoother blending"""
+    if len(mask.shape) == 3:
+        # Work with single channel
+        mask_single = mask[:, :, 0]
+    else:
+        mask_single = mask
+    
+    # Apply Gaussian blur for feathering
+    mask_feathered = cv2.GaussianBlur(mask_single, (feather_amount*2+1, feather_amount*2+1), 0)
+    
+    # Restore 3 channels if needed
+    if len(mask.shape) == 3:
+        mask_feathered = np.stack([mask_feathered] * 3, axis=-1)
+    
+    return mask_feathered
+
+def create_professional_background(bg_type: str, width: int, height: int) -> np.ndarray:
+    """Create professional background of specified type and size"""
+    try:
+        if bg_type not in PROFESSIONAL_BACKGROUNDS:
+            bg_type = "office"  # Default fallback
+        
+        config = PROFESSIONAL_BACKGROUNDS[bg_type]
+        color = config["color"]
+        use_gradient = config["gradient"]
+        
+        if use_gradient:
+            # Create gradient background
+            background = _create_gradient_background(color, width, height)
+        else:
+            # Create solid color background
+            background = np.full((height, width, 3), color, dtype=np.uint8)
+        
+        return background
+        
+    except Exception as e:
+        logger.error(f"Background creation failed: {e}")
+        # Return white background as fallback
+        return np.full((height, width, 3), (255, 255, 255), dtype=np.uint8)
+
+def _create_gradient_background(base_color: Tuple[int, int, int], width: int, height: int) -> np.ndarray:
+    """Create a gradient background from base color"""
+    # Create gradient from darker to lighter version of base color
+    r, g, b = base_color
+    
+    # Create darker version (multiply by 0.7)
+    dark_color = (int(r * 0.7), int(g * 0.7), int(b * 0.7))
+    
+    # Create gradient
+    background = np.zeros((height, width, 3), dtype=np.uint8)
+    
+    for y in range(height):
+        # Calculate blend factor (0 to 1)
+        blend = y / height
+        
+        # Interpolate between dark and light color
+        current_r = int(dark_color[0] * (1 - blend) + r * blend)
+        current_g = int(dark_color[1] * (1 - blend) + g * blend)
+        current_b = int(dark_color[2] * (1 - blend) + b * blend)
+        
+        background[y, :] = [current_r, current_g, current_b]
+    
+    return background
+
+# Export all functions
+__all__ = [
+    "segment_person_hq",
+    "refine_mask_hq", 
+    "replace_background_hq",
+    "create_professional_background",
+    "PROFESSIONAL_BACKGROUNDS",
+    "validate_video_file"
+]