clothes_segmentation.py

import os
import torch
import torch.nn as nn
import numpy as np
from typing import Union, Tuple
from PIL import Image, ImageFilter
import cv2
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
from huggingface_hub import hf_hub_download
import shutil

# Device configuration
device = "cuda" if torch.cuda.is_available() else "cpu"

# Model configuration
AVAILABLE_MODELS = {
    "segformer_b2_clothes": "1038lab/segformer_clothes"
}

# Model paths
current_dir = os.path.dirname(os.path.abspath(__file__))
models_dir = os.path.join(current_dir, "models")


def pil2tensor(image: Image.Image) -> torch.Tensor:
    """Convert PIL Image to tensor."""
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]


def tensor2pil(image: torch.Tensor) -> Image.Image:
    """Convert tensor to PIL Image."""
    return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))


def image2mask(image: Image.Image) -> torch.Tensor:
    """Convert image to mask tensor."""
    if isinstance(image, Image.Image):
        image = pil2tensor(image)
    return image.squeeze()[..., 0]


def mask2image(mask: torch.Tensor) -> Image.Image:
    """Convert mask tensor to PIL Image."""
    if len(mask.shape) == 2:
        mask = mask.unsqueeze(0)
    return tensor2pil(mask)


class ClothesSegmentation:
    """
    Standalone clothing segmentation using Segformer model.
    """
    
    def __init__(self):
        self.processor = None
        self.model = None
        self.cache_dir = os.path.join(models_dir, "RMBG", "segformer_clothes")
        
        # Class mapping for segmentation - consistent with latest repo
        self.class_map = {
            "Background": 0, "Hat": 1, "Hair": 2, "Sunglasses": 3, 
            "Upper-clothes": 4, "Skirt": 5, "Pants": 6, "Dress": 7,
            "Belt": 8, "Left-shoe": 9, "Right-shoe": 10, "Face": 11,
            "Left-leg": 12, "Right-leg": 13, "Left-arm": 14, "Right-arm": 15,
            "Bag": 16, "Scarf": 17
        }
    
    def check_model_cache(self):
        """Check if model files exist in cache."""
        if not os.path.exists(self.cache_dir):
            return False, "Model directory not found"
        
        required_files = [
            'config.json',
            'model.safetensors',
            'preprocessor_config.json'
        ]
        
        missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
        if missing_files:
            return False, f"Required model files missing: {', '.join(missing_files)}"
        return True, "Model cache verified"

    def clear_model(self):
        """Clear model from memory - improved version."""
        if self.model is not None:
            self.model.cpu()
            del self.model
            self.model = None
            self.processor = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def download_model_files(self):
        """Download model files from Hugging Face - improved version."""
        model_id = AVAILABLE_MODELS["segformer_b2_clothes"]
        model_files = {
            'config.json': 'config.json',
            'model.safetensors': 'model.safetensors',
            'preprocessor_config.json': 'preprocessor_config.json'
        }
        
        os.makedirs(self.cache_dir, exist_ok=True)
        print(f"Downloading Clothes Segformer model files...")
        
        try:
            for save_name, repo_path in model_files.items():
                print(f"Downloading {save_name}...")
                downloaded_path = hf_hub_download(
                    repo_id=model_id,
                    filename=repo_path,
                    local_dir=self.cache_dir,
                    local_dir_use_symlinks=False
                )
                
                if os.path.dirname(downloaded_path) != self.cache_dir:
                    target_path = os.path.join(self.cache_dir, save_name)
                    shutil.move(downloaded_path, target_path)
            return True, "Model files downloaded successfully"
        except Exception as e:
            return False, f"Error downloading model files: {str(e)}"

    def load_model(self):
        """Load the clothing segmentation model - improved version."""
        try:
            # Check and download model if needed
            cache_status, message = self.check_model_cache()
            if not cache_status:
                print(f"Cache check: {message}")
                download_status, download_message = self.download_model_files()
                if not download_status:
                    print(f"❌ {download_message}")
                    return False
            
            # Load model if needed
            if self.processor is None:
                print("Loading clothes segmentation model...")
                self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
                self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
                self.model.eval()
                for param in self.model.parameters():
                    param.requires_grad = False
                self.model.to(device)
                print("✅ Clothes segmentation model loaded successfully")
            
            return True
            
        except Exception as e:
            print(f"❌ Error loading clothes model: {e}")
            self.clear_model()  # Cleanup on error
            return False

    def segment_clothes(self, image_path: str, target_classes: list = None, process_res: int = 512) -> np.ndarray:
        """
        Segment clothing from an image - improved version with process_res parameter.
        
        Args:
            image_path: Path to the image
            target_classes: List of clothing classes to segment (default: ["Upper-clothes"])
            process_res: Processing resolution (default: 512)
            
        Returns:
            Binary mask as numpy array
        """
        if target_classes is None:
            target_classes = ["Upper-clothes"]
            
        if not self.load_model():
            print("❌ Cannot load clothes segmentation model")
            return None

        try:
            # Load and preprocess image
            image = cv2.imread(image_path)
            if image is None:
                print(f"❌ Could not load image: {image_path}")
                return None
                
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            original_size = image_rgb.shape[:2]
            
            # Preprocess image with custom resolution
            pil_image = Image.fromarray(image_rgb)
            
            # Resize for processing if needed
            if process_res != 512:
                pil_image = pil_image.resize((process_res, process_res), Image.Resampling.LANCZOS)
            
            inputs = self.processor(images=pil_image, return_tensors="pt")
            inputs = {k: v.to(device) for k, v in inputs.items()}
            
            # Run inference
            with torch.no_grad():
                outputs = self.model(**inputs)
                logits = outputs.logits.cpu()
                
                # Resize logits to original image size
                upsampled_logits = nn.functional.interpolate(
                    logits,
                    size=original_size,
                    mode="bilinear",
                    align_corners=False,
                )
                pred_seg = upsampled_logits.argmax(dim=1)[0]

                # Combine selected class masks
                combined_mask = None
                for class_name in target_classes:
                    if class_name in self.class_map:
                        mask = (pred_seg == self.class_map[class_name]).float()
                        if combined_mask is None:
                            combined_mask = mask
                        else:
                            combined_mask = torch.clamp(combined_mask + mask, 0, 1)
                    else:
                        print(f"⚠️  Unknown class: {class_name}")

                if combined_mask is None:
                    print(f"❌ No valid classes found in: {target_classes}")
                    return None

                # Convert to numpy
                mask_np = combined_mask.numpy().astype(np.float32)
                
                return mask_np

        except Exception as e:
            print(f"❌ Error in clothes segmentation: {e}")
            return None
        finally:
            # Clean up model if not training (consistent with updated repo)
            if self.model is not None and not self.model.training:
                self.clear_model()

    def segment_clothes_with_filters(self, image_path: str, target_classes: list = None, 
                                   mask_blur: int = 0, mask_offset: int = 0, 
                                   process_res: int = 512) -> np.ndarray:
        """
        Segment clothing with additional filtering options - new method from updated repo.
        
        Args:
            image_path: Path to the image
            target_classes: List of clothing classes to segment
            mask_blur: Blur amount for mask edges
            mask_offset: Expand/Shrink mask boundary
            process_res: Processing resolution
            
        Returns:
            Filtered binary mask as numpy array
        """
        # Get initial mask
        mask = self.segment_clothes(image_path, target_classes, process_res)
        if mask is None:
            return None
        
        try:
            # Convert to PIL for filtering
            mask_image = Image.fromarray((mask * 255).astype(np.uint8))
            
            # Apply blur if specified
            if mask_blur > 0:
                mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
            
            # Apply offset if specified
            if mask_offset != 0:
                if mask_offset > 0:
                    mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
                else:
                    mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
            
            # Convert back to numpy
            filtered_mask = np.array(mask_image).astype(np.float32) / 255.0
            return filtered_mask
            
        except Exception as e:
            print(f"❌ Error applying filters: {e}")
            return mask


# Standalone function for easy usage
def segment_upper_clothes(image_path: str) -> np.ndarray:
    """
    Convenience function to segment upper clothes from an image.
    
    Args:
        image_path: Path to the image
        
    Returns:
        Binary mask as numpy array or None if failed
    """
    segmenter = ClothesSegmentation()
    return segmenter.segment_clothes(image_path, ["Upper-clothes"])