CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

This is a garment mask generation system for image inpainting. The project uses deep learning models (DensePose and SCHP) to detect human body parts and garment regions, then generates variable-shaped masks suitable for inpainting tasks. The masks protect certain body areas (face, hands, feet) while creating randomized mask shapes around garments.

Architecture

Core Processing Pipeline

The system follows a multi-stage pipeline:

Image Preprocessing (resize_image_processor.py)
- Uses DensePose to detect person and crop intelligently around them
- Adjusts aspect ratios with padding/cropping
- Standardizes image dimensions
Segmentation (via SCHP/ and DensePose/)
- DensePose: Body part segmentation (torso, arms, legs, hands, feet, face)
- SCHP-ATR: Clothing segmentation (18 classes)
- SCHP-LIP: Refined clothing segmentation (20 classes)
Mask Generation (garment_mask_processor.py)
- Creates "strong protect" zones (face, hands, feet - never to be masked)
- Creates "weak protect" zones (other body parts based on garment type)
- Generates tight garment mask (mask_S)
- Expands to variable mask using one of three strategies:
  - Ellipse: Morphological dilation with elliptical kernel
  - Box: Jittered bounding box around garment
  - Poly: Polygonal approximation of dilated mask
- Strategy selection is deterministic based on image ID seed

Key Components

GarmentMaskProcessor (main entry point):

One-shot processor for batch garment mask generation
Supports garment types: "upper", "lower", "dress" (full body)
Configurable output dimensions, processing size, and mask strategies
Returns standardized images and masks for each garment type

ResizeImageProcessor:

Combines DensePose-based intelligent cropping with aspect ratio adjustment
Two-stage process: crop around person → adjust to target dimensions

mask_utils.py:

Low-level mask operations: dilation, morphology, connected components
Protect area computation (strong vs weak)
Mask area calculation per garment part

mappings.py:

Defines body part indices for DensePose (24 classes)
Defines garment class mappings for ATR and LIP
Maps garment types to body/cloth parts for masking and protection

Model Wrappers

SCHP/ (Self-Correction Human Parsing):

ResNet101-based segmentation network
Two variants: ATR dataset (18 classes) and LIP dataset (20 classes)
Input: PIL Image or path → Output: Palette-indexed PIL Image

DensePose/ (from Detectron2):

R_50_FPN_s1x architecture
Input: PIL Image or path → Output: Grayscale PIL Image with body part indices
Uses temporary files for processing

Checkpoint Structure

Model checkpoints should be organized as:

chkpt/
├── DensePose/
│   ├── model_final_162be9.pkl
│   ├── densepose_rcnn_R_50_FPN_s1x.yaml
│   └── Base-DensePose-RCNN-FPN.yaml
└── SCHP/
    ├── exp-schp-201908301523-atr.pth
    └── exp-schp-201908261155-lip.pth

Usage Patterns

Basic Usage

from garment_mask_processor import GarmentMaskProcessor

processor = GarmentMaskProcessor(
    device="cuda:0",
    densepose_ckpt="chkpt/DensePose",
    schp_atr_ckpt="chkpt/SCHP/exp-schp-201908301523-atr.pth",
    schp_lip_ckpt="chkpt/SCHP/exp-schp-201908261155-lip.pth"
)

results = processor.process_batch(
    images=["img1.jpg", "img2.jpg"],
    garment_types=["upper", "lower"],  # Generates both types for each image
    output_dir="./output"
)

Result Structure

Each result dict contains:

image_id: Unique identifier
image_standardized: Processed PIL Image (RGB)
masks: Dict with keys for each garment type
- person_mask: Variable-shaped mask for inpainting (PIL Image, mode L)

Important Implementation Details

Mask Strategy Selection

Deterministic based on MD5 hash of {image_id}_{garment_type}
Default probabilities: ellipse (50%), box (30%), poly (20%)
Can be restricted via allowed_strategies parameter

Protected Areas

Strong protect (always protected): Face, hands, feet
Weak protect (depends on garment type):
- Upper body: Protects legs/pants
- Lower body: Protects arms/upper-clothes
- Full body: Minimal protection

Coordinate Systems

DensePose outputs index values 0-24
SCHP outputs are already at input image resolution
All masks are resized to match standardized output dimensions

Dependencies

PyTorch (CUDA recommended)
Detectron2 (for DensePose)
OpenCV (cv2)
PIL/Pillow
numpy

File Structure Notes

detectron2/: Full Detectron2 framework (likely vendored)
densepose/: DensePose project from Detectron2
SCHP/: Human parsing model wrapper and network definitions
DensePose/: Wrapper class for DensePose inference
Root-level Python files: Main processing logic

Common Tasks

Adding a New Garment Type

Add type mapping in GarmentMaskProcessor.TYPE_TO_PART
Define body parts to protect in mappings.py:PROTECT_BODY_PARTS
Define cloth parts to protect in mappings.py:PROTECT_CLOTH_PARTS
Define mask cloth parts in mappings.py:MASK_CLOTH_PARTS
Define mask dense parts in mappings.py:MASK_DENSE_PARTS

Adjusting Mask Expansion

Modify _rho_params_from_area() in garment_mask_processor.py:162 to change dilation ratios based on garment area
Adjust strategy probabilities via constructor parameters or class constants

Debugging Mask Generation

Set save_images=True and provide output_dir to save intermediate results
Use save_mask_s=True to see tight mask before expansion
Use save_strong_protect=True to visualize protected areas