pipeline.py

"""
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(Pipeline):
    """
    Hugging Face 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 transformers import pipeline

        # Load the pipeline
        pipe = pipeline("image-segmentation",
                       model="your-username/wearit-garment-mask",
                       trust_remote_code=True)

        # Generate masks
        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: str = "chkpt/DensePose",
        schp_atr_ckpt: str = "chkpt/SCHP/exp-schp-201908301523-atr.pth",
        schp_lip_ckpt: str = "chkpt/SCHP/exp-schp-201908261155-lip.pth",
        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)
        """
        if HF_AVAILABLE:
            super().__init__(model=model, **kwargs)

        # 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 model is not None and os.path.exists(model):
            # Model is a local directory
            base_path = model
            densepose_ckpt = os.path.join(base_path, "chkpt/DensePose")
            schp_atr_ckpt = os.path.join(base_path, "chkpt/SCHP/exp-schp-201908301523-atr.pth")
            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
        """
        if HF_AVAILABLE:
            return super().__call__(inputs, **kwargs)
        else:
            # Fallback for non-HF usage
            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)


# Register the pipeline if transformers is available
if HF_AVAILABLE:
    try:
        from transformers.pipelines import PIPELINE_REGISTRY

        PIPELINE_REGISTRY.register_pipeline(
            "garment-mask-generation",
            pipeline_class=GarmentMaskPipeline,
            pt_model=None,  # Custom pipeline without standard model
        )
    except Exception as e:
        if HF_AVAILABLE:
            logger.warning(f"Could not register pipeline: {e}")