File size: 9,107 Bytes
436df5c f40e7e5 436df5c e1f4ba1 436df5c e1f4ba1 436df5c e1f4ba1 436df5c e1f4ba1 436df5c 6090de7 436df5c 6090de7 436df5c e1f4ba1 436df5c e1f4ba1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 | """
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")
|