""" Hugging Face Pipeline wrapper for GarmentMaskProcessor """ import os from typing import Dict, List, Union, Optional, Any import numpy as np from PIL import Image try: from transformers import Pipeline from transformers.utils import logging HF_AVAILABLE = True except ImportError: HF_AVAILABLE = False # Fallback base class if transformers not installed class Pipeline: pass from garment_mask_processor import GarmentMaskProcessor if HF_AVAILABLE: logger = logging.get_logger(__name__) class GarmentMaskPipeline: """ Standalone Pipeline for generating garment masks for image inpainting. This pipeline combines DensePose and SCHP models to detect human body parts and garment regions, then generates variable-shaped masks suitable for inpainting. Example: ```python from huggingface_hub import snapshot_download import sys # Download from HF Hub repo_path = snapshot_download("your-username/wearit-garment-mask") sys.path.insert(0, repo_path) # Import and use from pipeline import GarmentMaskPipeline pipe = GarmentMaskPipeline(model=repo_path, device="cuda") results = pipe("person.jpg", garment_types=["upper", "lower"]) # Access masks for result in results: result["masks"]["upper"]["person_mask"].save("upper_mask.png") ``` """ def __init__( self, model: Optional[str] = None, device: str = "cuda:0", densepose_ckpt: Optional[str] = None, schp_atr_ckpt: Optional[str] = None, schp_lip_ckpt: Optional[str] = None, output_height: int = 1024, process_size: int = 512, use_convex_hull: bool = True, schp_batch_size: int = 12, allowed_strategies: Optional[List[str]] = None, save_images: bool = False, **kwargs ): """ Initialize the GarmentMask pipeline. Args: model: Model identifier (for HF Hub compatibility) device: Device to run models on ("cuda:0", "cpu") densepose_ckpt: Path to DensePose checkpoint schp_atr_ckpt: Path to SCHP-ATR checkpoint schp_lip_ckpt: Path to SCHP-LIP checkpoint output_height: Height of standardized output (default: 1024) process_size: Size for intermediate processing (default: 512) use_convex_hull: Apply convex hull to final mask (default: True) schp_batch_size: Batch size for SCHP processing (default: 12) allowed_strategies: List of allowed mask strategies ["ellipse", "box", "poly"] save_images: Whether to save intermediate images (default: False) """ # Store config self.device = device self.output_height = output_height self.process_size = process_size self.use_convex_hull = use_convex_hull self.schp_batch_size = schp_batch_size self.allowed_strategies = allowed_strategies self.save_images = save_images # Resolve checkpoint paths (support both local and HF Hub) if densepose_ckpt is None or schp_atr_ckpt is None or schp_lip_ckpt is None: # Try to load from model directory (HF Hub or local) if model is not None: if os.path.exists(model): # Local directory base_path = model else: # Try HF Hub download try: from huggingface_hub import snapshot_download base_path = snapshot_download(repo_id=model) except: base_path = "." else: base_path = "." # Set default paths if not provided if densepose_ckpt is None: densepose_ckpt = os.path.join(base_path, "chkpt/DensePose") if schp_atr_ckpt is None: schp_atr_ckpt = os.path.join(base_path, "chkpt/SCHP/exp-schp-201908301523-atr.pth") if schp_lip_ckpt is None: schp_lip_ckpt = os.path.join(base_path, "chkpt/SCHP/exp-schp-201908261155-lip.pth") # Initialize the processor self.processor = GarmentMaskProcessor( device=device, densepose_ckpt=densepose_ckpt, schp_atr_ckpt=schp_atr_ckpt, schp_lip_ckpt=schp_lip_ckpt, output_height=output_height, process_size=process_size, use_convex_hull=use_convex_hull, schp_batch_size=schp_batch_size, allowed_strategies=allowed_strategies, save_images=save_images ) def _sanitize_parameters( self, garment_types: Union[str, List[str]] = None, image_ids: Optional[List[str]] = None, output_dir: Optional[str] = None, save_mask_s: bool = False, save_strong_protect: bool = False, **kwargs ): """ Sanitize and organize parameters for preprocessing, forward, and postprocessing. Args: garment_types: Type(s) of garment to generate masks for image_ids: Optional IDs for images (for deterministic seed) output_dir: Directory to save outputs save_mask_s: Save tight mask before expansion save_strong_protect: Save protected areas """ preprocess_params = {} forward_params = { "garment_types": garment_types or "upper", "image_ids": image_ids, "output_dir": output_dir, "save_mask_s": save_mask_s, "save_strong_protect": save_strong_protect, } postprocess_params = {} return preprocess_params, forward_params, postprocess_params def preprocess(self, inputs: Union[str, Image.Image, List[Union[str, Image.Image]]]) -> Dict[str, Any]: """ Preprocess inputs to prepare for model inference. Args: inputs: Image path(s), PIL Image(s), or list of either Returns: Dictionary with preprocessed images """ # Ensure inputs is a list if not isinstance(inputs, list): inputs = [inputs] # Load images if they're paths images = [] for img in inputs: if isinstance(img, str): images.append(Image.open(img).convert("RGB")) elif isinstance(img, Image.Image): images.append(img.convert("RGB")) else: raise ValueError(f"Unsupported input type: {type(img)}") return {"images": images} def _forward(self, model_inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: """ Run the actual model inference. Args: model_inputs: Dictionary with preprocessed images forward_params: Additional parameters for processing Returns: Dictionary with model outputs """ images = model_inputs["images"] # Call the processor results = self.processor.process_batch( images=images, garment_types=forward_params.get("garment_types", "upper"), image_ids=forward_params.get("image_ids"), output_dir=forward_params.get("output_dir"), save_mask_s=forward_params.get("save_mask_s", False), save_strong_protect=forward_params.get("save_strong_protect", False) ) return {"results": results} def postprocess(self, model_outputs: Dict[str, Any]) -> List[Dict[str, Any]]: """ Postprocess model outputs to final format. Args: model_outputs: Dictionary with raw model results Returns: List of dictionaries, one per image, with: - image_id: Unique identifier - image_standardized: Processed PIL Image - masks: Dict with garment type keys, each containing: - person_mask: Variable mask for inpainting (PIL Image) """ return model_outputs["results"] def __call__( self, inputs: Union[str, Image.Image, List[Union[str, Image.Image]]], **kwargs ) -> List[Dict[str, Any]]: """ Main entry point for the pipeline. Args: inputs: Image path(s) or PIL Image(s) **kwargs: Additional parameters (garment_types, output_dir, etc.) Returns: List of results with masks for each image """ # Direct processing without Pipeline inheritance preprocess_params, forward_params, postprocess_params = self._sanitize_parameters(**kwargs) model_inputs = self.preprocess(inputs) model_outputs = self._forward(model_inputs, **forward_params) return self.postprocess(model_outputs) # Note: This pipeline is standalone and does not inherit from transformers.Pipeline # Use it directly: pipe = GarmentMaskPipeline(model="repo_path", device="cuda")