Create hair_segmentation.py

hair_segmentation.py

@@ -0,0 +1,576 @@
+"""
+Professional Hair Segmentation Module
+=====================================
+
+This module provides high-quality hair segmentation for video processing
+using SAM2 + MatAnyone pipeline with comprehensive error handling and fallbacks.
+
+Author: Your Project
+License: MIT
+"""
+
+import os
+import torch
+import cv2
+import numpy as np
+import logging
+from typing import Dict, List, Tuple, Optional, Union
+from pathlib import Path
+import warnings
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+
+# Fix threading issues immediately
+os.environ['OMP_NUM_THREADS'] = '4'
+os.environ['MKL_NUM_THREADS'] = '4'
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class SegmentationResult:
+    """Result container for hair segmentation"""
+    mask: np.ndarray
+    confidence: float
+    coverage_percent: float
+    asymmetry_score: float
+    processing_time: float
+    fallback_used: bool
+    quality_score: float
+    error_message: Optional[str] = None
+
+class BaseSegmentationModel(ABC):
+    """Abstract base class for segmentation models"""
+    
+    @abstractmethod
+    def initialize(self) -> bool:
+        """Initialize the model"""
+        pass
+    
+    @abstractmethod
+    def segment(self, frame: np.ndarray) -> np.ndarray:
+        """Segment hair in frame"""
+        pass
+    
+    @abstractmethod
+    def get_model_name(self) -> str:
+        """Get model name for logging"""
+        pass
+
+class SAM2Model(BaseSegmentationModel):
+    """SAM2 segmentation model wrapper"""
+    
+    def __init__(self, model_path: Optional[str] = None, device: str = 'auto'):
+        self.model_path = model_path
+        self.device = self._get_best_device(device)
+        self.predictor = None
+        self.initialized = False
+    
+    def _get_best_device(self, device: str) -> str:
+        """Determine best available device"""
+        if device == 'auto':
+            return 'cuda' if torch.cuda.is_available() else 'cpu'
+        return device
+    
+    def initialize(self) -> bool:
+        """Initialize SAM2 model"""
+        try:
+            logger.info("🤖 Initializing SAM2 model...")
+            
+            # Import SAM2 (handle different installation methods)
+            try:
+                from sam2.build_sam import build_sam2_video_predictor
+            except ImportError:
+                logger.error("SAM2 not found. Please install SAM2.")
+                return False
+            
+            # Build predictor
+            if self.model_path and Path(self.model_path).exists():
+                self.predictor = build_sam2_video_predictor(self.model_path, device=self.device)
+            else:
+                # Use default model
+                self.predictor = build_sam2_video_predictor("sam2_hiera_large.pt", device=self.device)
+            
+            self.initialized = True
+            logger.info(f"✅ SAM2 initialized on {self.device}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ SAM2 initialization failed: {e}")
+            return False
+    
+    def segment(self, frame: np.ndarray) -> np.ndarray:
+        """Segment using SAM2"""
+        if not self.initialized:
+            raise RuntimeError("SAM2 model not initialized")
+        
+        try:
+            # Convert BGR to RGB
+            if len(frame.shape) == 3:
+                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            else:
+                frame_rgb = frame
+            
+            # Set image for SAM2
+            self.predictor.set_image(frame_rgb)
+            
+            # Auto-detect person in center (you can make this more sophisticated)
+            height, width = frame_rgb.shape[:2]
+            center_point = np.array([[width//2, height//2]])
+            
+            # Predict mask
+            masks, scores, _ = self.predictor.predict(
+                point_coords=center_point,
+                point_labels=np.array([1])
+            )
+            
+            # Return best mask
+            if len(masks) > 0:
+                best_mask_idx = np.argmax(scores)
+                return masks[best_mask_idx].astype(np.float32)
+            else:
+                return np.zeros((height, width), dtype=np.float32)
+                
+        except Exception as e:
+            logger.error(f"SAM2 segmentation failed: {e}")
+            raise
+    
+    def get_model_name(self) -> str:
+        return "SAM2"
+
+class MatAnyoneModel(BaseSegmentationModel):
+    """MatAnyone model wrapper with quality checking"""
+    
+    def __init__(self, use_hf_api: bool = True, hf_token: Optional[str] = None):
+        self.use_hf_api = use_hf_api
+        self.hf_token = hf_token
+        self.client = None
+        self.processor = None
+        self.initialized = False
+        self.quality_threshold = 0.3
+    
+    def initialize(self) -> bool:
+        """Initialize MatAnyone model"""
+        try:
+            logger.info("🎭 Initializing MatAnyone model...")
+            
+            if self.use_hf_api:
+                from gradio_client import Client
+                self.client = Client("PeiqingYang/MatAnyone", hf_token=self.hf_token)
+                logger.info("✅ MatAnyone HF API initialized")
+            else:
+                # Local MatAnyone initialization would go here
+                logger.warning("Local MatAnyone not implemented yet")
+                return False
+            
+            self.initialized = True
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ MatAnyone initialization failed: {e}")
+            return False
+    
+    def segment(self, frame: np.ndarray) -> np.ndarray:
+        """MatAnyone is primarily for matting, not segmentation"""
+        raise NotImplementedError("MatAnyone is used for matting, not direct segmentation")
+    
+    def matte(self, image: np.ndarray, trimap: np.ndarray) -> np.ndarray:
+        """Apply matting using MatAnyone"""
+        if not self.initialized:
+            raise RuntimeError("MatAnyone model not initialized")
+        
+        try:
+            # Save temporary files
+            import tempfile
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as img_file:
+                cv2.imwrite(img_file.name, image)
+                img_path = img_file.name
+            
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tri_file:
+                cv2.imwrite(tri_file.name, trimap)
+                tri_path = tri_file.name
+            
+            # Process with MatAnyone
+            if self.use_hf_api:
+                result = self._process_hf_api(img_path, tri_path)
+            else:
+                result = self._process_local(img_path, tri_path)
+            
+            # Cleanup temp files
+            os.unlink(img_path)
+            os.unlink(tri_path)
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"MatAnyone matting failed: {e}")
+            raise
+    
+    def _process_hf_api(self, image_path: str, trimap_path: str) -> np.ndarray:
+        """Process using HuggingFace API"""
+        try:
+            result = self.client.predict(
+                image=image_path,
+                trimap=trimap_path,
+                api_name="/predict"
+            )
+            
+            # Load result
+            if isinstance(result, str):
+                result_image = cv2.imread(result)
+                return result_image
+            else:
+                return result
+                
+        except Exception as e:
+            logger.error(f"HF API processing failed: {e}")
+            raise
+    
+    def get_model_name(self) -> str:
+        return "MatAnyone"
+
+class TraditionalCVModel(BaseSegmentationModel):
+    """Traditional computer vision fallback"""
+    
+    def __init__(self):
+        self.initialized = False
+    
+    def initialize(self) -> bool:
+        """Initialize traditional CV methods"""
+        self.initialized = True
+        return True
+    
+    def segment(self, frame: np.ndarray) -> np.ndarray:
+        """Traditional hair segmentation using color and texture"""
+        try:
+            # Convert to different color spaces
+            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+            lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)
+            
+            # Hair color detection
+            hair_mask_hsv = self._detect_hair_hsv(hsv)
+            hair_mask_lab = self._detect_hair_lab(lab)
+            
+            # Combine masks
+            combined_mask = cv2.bitwise_or(hair_mask_hsv, hair_mask_lab)
+            
+            # Morphological operations
+            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
+            combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
+            combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
+            
+            return combined_mask.astype(np.float32) / 255.0
+            
+        except Exception as e:
+            logger.error(f"Traditional CV segmentation failed: {e}")
+            raise
+    
+    def _detect_hair_hsv(self, hsv: np.ndarray) -> np.ndarray:
+        """Detect hair in HSV color space"""
+        # Multiple hair color ranges
+        ranges = [
+            # Dark hair
+            ([0, 0, 0], [180, 255, 80]),
+            # Brown hair
+            ([8, 50, 20], [25, 255, 200]),
+            # Blonde hair
+            ([15, 30, 100], [35, 255, 255])
+        ]
+        
+        masks = []
+        for lower, upper in ranges:
+            mask = cv2.inRange(hsv, np.array(lower), np.array(upper))
+            masks.append(mask)
+        
+        # Combine all color ranges
+        final_mask = masks[0]
+        for mask in masks[1:]:
+            final_mask = cv2.bitwise_or(final_mask, mask)
+        
+        return final_mask
+    
+    def _detect_hair_lab(self, lab: np.ndarray) -> np.ndarray:
+        """Detect hair in LAB color space"""
+        l_channel = lab[:, :, 0]
+        hair_mask = cv2.inRange(l_channel, 0, 120)
+        return hair_mask
+    
+    def get_model_name(self) -> str:
+        return "TraditionalCV"
+
+class TemporalSmoother:
+    """Temporal smoothing for video sequences"""
+    
+    def __init__(self, smoothing_factor: float = 0.7, change_threshold: float = 0.05):
+        self.smoothing_factor = smoothing_factor
+        self.change_threshold = change_threshold
+        self.previous_mask = None
+        self.correction_count = 0
+        self.total_frames = 0
+    
+    def smooth(self, current_mask: np.ndarray) -> Tuple[np.ndarray, bool]:
+        """Apply temporal smoothing"""
+        self.total_frames += 1
+        corrected = False
+        
+        if self.previous_mask is not None:
+            # Calculate change
+            diff = np.mean(np.abs(current_mask - self.previous_mask))
+            
+            if diff > self.change_threshold:
+                # Apply smoothing
+                smoothed_mask = (self.smoothing_factor * current_mask + 
+                               (1 - self.smoothing_factor) * self.previous_mask)
+                self.correction_count += 1
+                corrected = True
+            else:
+                smoothed_mask = current_mask
+        else:
+            smoothed_mask = current_mask
+        
+        self.previous_mask = smoothed_mask.copy()
+        return smoothed_mask, corrected
+    
+    def get_correction_ratio(self) -> float:
+        """Get ratio of frames that needed correction"""
+        return self.correction_count / max(self.total_frames, 1)
+
+class HairSegmentationPipeline:
+    """Main hair segmentation pipeline with multiple models and fallbacks"""
+    
+    def __init__(self, config: Optional[Dict] = None):
+        self.config = config or {}
+        self.models = {}
+        self.active_model = None
+        self.fallback_models = []
+        self.temporal_smoother = TemporalSmoother()
+        self.initialized = False
+        
+        # Setup models
+        self._setup_models()
+    
+    def _setup_models(self):
+        """Setup available models"""
+        try:
+            # Primary model: SAM2
+            sam2_model = SAM2Model(
+                model_path=self.config.get('sam2_model_path'),
+                device=self.config.get('device', 'auto')
+            )
+            self.models['sam2'] = sam2_model
+            
+            # MatAnyone for matting
+            matanyone_model = MatAnyoneModel(
+                use_hf_api=self.config.get('use_hf_api', True),
+                hf_token=self.config.get('hf_token')
+            )
+            self.models['matanyone'] = matanyone_model
+            
+            # Fallback: Traditional CV
+            traditional_model = TraditionalCVModel()
+            self.models['traditional'] = traditional_model
+            
+        except Exception as e:
+            logger.error(f"Model setup failed: {e}")
+    
+    def initialize(self, preferred_model: str = 'sam2') -> bool:
+        """Initialize the pipeline"""
+        logger.info("🚀 Initializing Hair Segmentation Pipeline...")
+        
+        # Try to initialize preferred model
+        if preferred_model in self.models:
+            if self.models[preferred_model].initialize():
+                self.active_model = preferred_model
+                logger.info(f"✅ Primary model {preferred_model} initialized")
+            else:
+                logger.warning(f"⚠️ Primary model {preferred_model} failed")
+        
+        # Initialize fallback models
+        for model_name, model in self.models.items():
+            if model_name != self.active_model:
+                if model.initialize():
+                    self.fallback_models.append(model_name)
+                    logger.info(f"✅ Fallback model {model_name} ready")
+        
+        # Check if we have at least one working model
+        if self.active_model or self.fallback_models:
+            self.initialized = True
+            logger.info(f"🎯 Pipeline ready - Active: {self.active_model}, Fallbacks: {self.fallback_models}")
+            return True
+        else:
+            logger.error("❌ No working models available")
+            return False
+    
+    def segment_frame(self, frame: np.ndarray, 
+                     apply_temporal_smoothing: bool = True) -> SegmentationResult:
+        """Segment hair in a single frame"""
+        if not self.initialized:
+            raise RuntimeError("Pipeline not initialized")
+        
+        import time
+        start_time = time.time()
+        
+        # Try active model first
+        mask, model_used, error_msg = self._try_segment_with_model(frame, self.active_model)
+        
+        # If failed, try fallback models
+        if mask is None:
+            for fallback_model in self.fallback_models:
+                mask, model_used, error_msg = self._try_segment_with_model(frame, fallback_model)
+                if mask is not None:
+                    break
+        
+        if mask is None:
+            # Complete failure - return empty mask
+            h, w = frame.shape[:2]
+            mask = np.zeros((h, w), dtype=np.float32)
+            model_used = "none"
+            error_msg = "All models failed"
+        
+        # Apply temporal smoothing
+        corrected = False
+        if apply_temporal_smoothing:
+            mask, corrected = self.temporal_smoother.smooth(mask)
+        
+        # Calculate metrics
+        processing_time = time.time() - start_time
+        confidence = self._calculate_confidence(mask)
+        coverage = self._calculate_coverage(mask)
+        asymmetry = self._calculate_asymmetry(mask)
+        quality = self._calculate_quality(mask)
+        
+        return SegmentationResult(
+            mask=mask,
+            confidence=confidence,
+            coverage_percent=coverage,
+            asymmetry_score=asymmetry,
+            processing_time=processing_time,
+            fallback_used=(model_used != self.active_model),
+            quality_score=quality,
+            error_message=error_msg
+        )
+    
+    def _try_segment_with_model(self, frame: np.ndarray, model_name: str) -> Tuple[Optional[np.ndarray], str, Optional[str]]:
+        """Try to segment with a specific model"""
+        if model_name not in self.models:
+            return None, model_name, f"Model {model_name} not available"
+        
+        try:
+            mask = self.models[model_name].segment(frame)
+            return mask, model_name, None
+        except Exception as e:
+            error_msg = f"Model {model_name} failed: {str(e)}"
+            logger.warning(error_msg)
+            return None, model_name, error_msg
+    
+    def _calculate_confidence(self, mask: np.ndarray) -> float:
+        """Calculate mask confidence"""
+        # Edge sharpness
+        edges = cv2.Canny((mask * 255).astype(np.uint8), 50, 150)
+        edge_ratio = np.sum(edges > 0) / mask.size
+        
+        # Mask smoothness
+        gradient = np.gradient(mask)
+        smoothness = 1.0 / (1.0 + np.std(gradient))
+        
+        return min(edge_ratio * 0.3 + smoothness * 0.7, 1.0)
+    
+    def _calculate_coverage(self, mask: np.ndarray) -> float:
+        """Calculate hair coverage percentage"""
+        return (np.sum(mask > 0.5) / mask.size) * 100
+    
+    def _calculate_asymmetry(self, mask: np.ndarray) -> float:
+        """Calculate left-right asymmetry score"""
+        h, w = mask.shape[:2]
+        center_x = w // 2
+        
+        left_half = mask[:, :center_x]
+        right_half = np.fliplr(mask[:, center_x:])
+        
+        min_width = min(left_half.shape[1], right_half.shape[1])
+        left_half = left_half[:, :min_width]
+        right_half = right_half[:, :min_width]
+        
+        return np.mean(np.abs(left_half - right_half))
+    
+    def _calculate_quality(self, mask: np.ndarray) -> float:
+        """Calculate overall mask quality"""
+        # Combine multiple quality metrics
+        confidence = self._calculate_confidence(mask)
+        coverage = self._calculate_coverage(mask) / 100.0
+        asymmetry_penalty = 1.0 - min(self._calculate_asymmetry(mask), 1.0)
+        
+        return (confidence * 0.5 + coverage * 0.3 + asymmetry_penalty * 0.2)
+    
+    def get_pipeline_stats(self) -> Dict:
+        """Get pipeline performance statistics"""
+        return {
+            'active_model': self.active_model,
+            'fallback_models': self.fallback_models,
+            'temporal_correction_ratio': self.temporal_smoother.get_correction_ratio(),
+            'total_frames_processed': self.temporal_smoother.total_frames,
+            'corrections_applied': self.temporal_smoother.correction_count
+        }
+
+# Convenience functions
+def create_pipeline(config: Optional[Dict] = None) -> HairSegmentationPipeline:
+    """Create and initialize hair segmentation pipeline"""
+    pipeline = HairSegmentationPipeline(config)
+    pipeline.initialize()
+    return pipeline
+
+def segment_image(image_path: str, config: Optional[Dict] = None) -> SegmentationResult:
+    """Segment hair in a single image"""
+    pipeline = create_pipeline(config)
+    frame = cv2.imread(image_path)
+    return pipeline.segment_frame(frame)
+
+def segment_video_frames(video_frames: List[np.ndarray], 
+                        config: Optional[Dict] = None) -> List[SegmentationResult]:
+    """Segment hair in multiple video frames"""
+    pipeline = create_pipeline(config)
+    results = []
+    
+    for frame in video_frames:
+        result = pipeline.segment_frame(frame)
+        results.append(result)
+    
+    return results
+
+# Example usage
+if __name__ == "__main__":
+    # Example configuration
+    config = {
+        'sam2_model_path': None,  # Use default
+        'device': 'auto',
+        'use_hf_api': True,
+        'hf_token': None  # Set your token if needed
+    }
+    
+    # Create pipeline
+    pipeline = create_pipeline(config)
+    
+    # Test with example frame (you would load your actual frame)
+    test_frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+    
+    # Segment frame
+    result = pipeline.segment_frame(test_frame)
+    
+    # Print results
+    print(f"Segmentation Results:")
+    print(f"  Coverage: {result.coverage_percent:.1f}%")
+    print(f"  Confidence: {result.confidence:.3f}")
+    print(f"  Quality: {result.quality_score:.3f}")
+    print(f"  Processing time: {result.processing_time:.2f}s")
+    print(f"  Fallback used: {result.fallback_used}")
+    
+    # Get pipeline stats
+    stats = pipeline.get_pipeline_stats()
+    print(f"\nPipeline Stats:")
+    print(f"  Active model: {stats['active_model']}")
+    print(f"  Fallbacks: {stats['fallback_models']}")
+    print(f"  Correction ratio: {stats['temporal_correction_ratio']:.3f}")