# 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:

1. **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

2. **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)

3. **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
```python
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
1. Add type mapping in `GarmentMaskProcessor.TYPE_TO_PART`
2. Define body parts to protect in `mappings.py:PROTECT_BODY_PARTS`
3. Define cloth parts to protect in `mappings.py:PROTECT_CLOTH_PARTS`
4. Define mask cloth parts in `mappings.py:MASK_CLOTH_PARTS`
5. 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