Face Parsing

https://github.com/yakhyo/face-parsing
Admin
Face Parsing is a high-precision facial feature segmentation model based on BiSeNet that accurately...
Tokens:7,059
Snippets:49
Trust Score:9.7
License:MIT
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install Dependencies and Download Weights

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Commands to clone the repository, install requirements, and download pretrained model weights.

```bash
# Clone repository
git clone https://github.com/yakhyo/face-parsing.git
cd face-parsing

# Install dependencies
pip install -r requirements.txt
# Requirements:
# torch==2.2.2
# torchvision==0.17.2
# tqdm==4.66.3
# opencv-python==4.9.0.80
# pillow==12.1.1

# For ONNX inference, also install:
pip install onnxruntime-gpu  # or onnxruntime for CPU-only

# Download pretrained weights
mkdir -p weights
wget -P weights https://github.com/yakhyo/face-parsing/releases/download/weights/resnet18.pt
wget -P weights https://github.com/yakhyo/face-parsing/releases/download/weights/resnet34.pt

# Download ONNX models
wget -P weights https://github.com/yakhyo/face-parsing/releases/download/weights/resnet18.onnx
wget -P weights https://github.com/yakhyo/face-parsing/releases/download/weights/resnet34.onnx
```

--------------------------------

### Execute Training

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Commands to start the training process with default or custom parameters.

```bash
python train.py --data-root /path/to/your/dataset
```

```bash
python train.py --data-root /path/to/dataset --backbone resnet34 --batch-size 8 --epochs 200
```

--------------------------------

### BiSeNet Training Setup

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Configures and initializes components for training the BiSeNet model. Includes dataset loading, model instantiation, loss function, optimizer, and learning rate scheduler.

```python
import torch
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import PolynomialLR

from models.bisenet import BiSeNet
from utils.dataset import CelebAMaskHQ
from utils.loss import OhemLossWrapper
from utils.transform import TrainTransform

# Configuration
config = {
    'num_classes': 19,
    'batch_size': 8,
    'epochs': 100,
    'lr_start': 1e-2,
    'momentum': 0.9,
    'weight_decay': 5e-4,
    'image_size': (448, 448),
    'backbone': 'resnet18',
    'data_root': '/path/to/CelebAMask-HQ/',
}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Dataset and DataLoader
images_dir = f"{config['data_root']}/CelebA-HQ-img"
labels_dir = f"{config['data_root']}/mask"
dataset = CelebAMaskHQ(
    images_dir,
    labels_dir,
    transform=TrainTransform(image_size=config['image_size'])
)
data_loader = DataLoader(
    dataset,
    batch_size=config['batch_size'],
    shuffle=True,
    num_workers=12,
    pin_memory=True,
    drop_last=True,
)

# Model
model = BiSeNet(num_classes=config['num_classes'], backbone_name=config['backbone'])
model.to(device)

# Loss with Online Hard Example Mining
n_min = config['batch_size'] * config['image_size'][0] * config['image_size'][1] // 16
criterion = OhemLossWrapper(thresh=0.7, min_kept=n_min)

# Optimizer
optimizer = torch.optim.SGD(
    model.parameters(),
    lr=config['lr_start'],
    momentum=config['momentum'],
    weight_decay=config['weight_decay']
)

# Learning rate scheduler
lr_scheduler = PolynomialLR(
    optimizer,
    total_iters=len(data_loader) * config['epochs'],
    power=0.9
)

```

--------------------------------

### Clone Repository and Install Dependencies

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Clone the face-parsing repository and install the necessary Python packages using pip. Ensure you are in the cloned directory before running the installation command.

```bash
git clone https://github.com/yakhyo/face-parsing.git
cd face-parsing
pip install -r requirements.txt
```

--------------------------------

### Setup PyTorch Inference Pipeline

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Initializes the environment and imports necessary modules for the full inference pipeline.

```python
import os
import numpy as np
import torch
from PIL import Image
from models.bisenet import BiSeNet
from utils.common import vis_parsing_maps

# Setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
num_classes = 19
```

--------------------------------

### Command Line Training Script

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Launch training from the command line with customizable parameters for data root, backbone, batch size, epochs, and learning rate.

```bash
python train.py --data-root /path/to/CelebAMask-HQ
```

```bash
python train.py \
    --data-root /path/to/CelebAMask-HQ \
    --backbone resnet34 \
    --batch-size 8 \
    --epochs 150 \
    --lr-start 0.01 \
    --image-size 512 512
```

```bash
python train.py \
    --data-root /path/to/CelebAMask-HQ \
    --backbone resnet18 \
    --resume
```

--------------------------------

### View Training Configuration

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Displays the available command-line arguments for the training script.

```text
usage: train.py [-h] [--num-classes NUM_CLASSES] [--batch-size BATCH_SIZE] [--num-workers NUM_WORKERS] [--image-size IMAGE_SIZE IMAGE_SIZE] [--data-root DATA_ROOT] [--momentum MOMENTUM] [--weight-decay WEIGHT_DECAY] [--lr-start LR_START]
                [--max-iter MAX_ITER] [--power POWER] [--lr-warmup-epochs LR_WARMUP_EPOCHS] [--warmup-start-lr WARMUP_START_LR] [--score-thres SCORE_THRES] [--epochs EPOCHS] [--backbone BACKBONE] [--print-freq PRINT_FREQ] [--resume]

Argument Parser for Training Configuration

options:
  -h, --help            show this help message and exit
  --num-classes NUM_CLASSES
                        Number of classes in the dataset (default: 19)
  --batch-size BATCH_SIZE
                        Batch size for training (default: 16)
  --num-workers NUM_WORKERS
                        Number of workers for data loading (default: 4)
  --image-size IMAGE_SIZE IMAGE_SIZE
                        Size of input images (default: 512 512)
  --data-root DATA_ROOT
                        Root directory of the dataset
  --momentum MOMENTUM   Momentum for optimizer (default: 0.9)
  --weight-decay WEIGHT_DECAY
                        Weight decay for optimizer (default: 5e-4)
  --lr-start LR_START   Initial learning rate (default: 1e-2)
  --max-iter MAX_ITER   Maximum number of iterations (default: 80000)
  --power POWER         Power for learning rate policy (default: 0.9)
  --lr-warmup-epochs LR_WARMUP_EPOCHS
                        Number of warmup epochs (default: 10)
  --warmup-start-lr WARMUP_START_LR
                        Warmup starting learning rate (default: 1e-5)
  --score-thres SCORE_THRES
                        Score threshold (default: 0.7)
  --epochs EPOCHS       Number of epochs for training (default: 150)
  --backbone BACKBONE   Backbone architecture (default: resnet18)
  --print-freq PRINT_FREQ
                        Print frequency during training (default: 10)
  --resume              Resume training from checkpoint
```

--------------------------------

### Initialize BiSeNet Model

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Instantiates the BiSeNet architecture with a specified backbone and performs a forward pass with dummy input.

```python
import torch
from models.bisenet import BiSeNet

# Initialize the BiSeNet model with ResNet18 backbone
num_classes = 19  # 19 facial component classes
model = BiSeNet(num_classes=num_classes, backbone_name='resnet18')

# Move to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()

# Create dummy input (batch_size=1, channels=3, height=512, width=512)
dummy_input = torch.randn(1, 3, 512, 512).to(device)

# Forward pass returns three outputs: (feat_out, feat_out16, feat_out32)
# During inference, use feat_out (first output) for the final segmentation mask
with torch.no_grad():
    feat_out, feat_out16, feat_out32 = model(dummy_input)
    predicted_mask = feat_out.squeeze(0).cpu().numpy().argmax(0)

print(f"Output shape: {feat_out.shape}")  # torch.Size([1, 19, 512, 512])
print(f"Mask shape: {predicted_mask.shape}")  # (512, 512)
```

--------------------------------

### Python Inference with BiSeNet

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Load a pre-trained BiSeNet model, preprocess an image, run inference, and visualize the segmentation mask. Ensure the model weights and device are correctly specified.

```python
model = BiSeNet(num_classes, backbone_name='resnet18')
model.load_state_dict(torch.load('./weights/resnet18.pt', map_location=device))
model.to(device)
model.eval()

image = Image.open('./assets/images/1.jpg').convert('RGB')
original_size = image.size  # (width, height)

transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
image_tensor = transform(image).unsqueeze(0).to(device)

with torch.no_grad():
    output = model(image_tensor)[0]  # Use first output for inference
    predicted_mask = output.squeeze(0).cpu().numpy().argmax(0)

mask_pil = Image.fromarray(predicted_mask.astype(np.uint8))
restored_mask = np.array(mask_pil.resize(original_size, resample=Image.NEAREST))

vis_parsing_maps(
    image,
    restored_mask,
    save_image=True,
    save_path='./result.jpg'
)
print("Segmentation complete!")
```

--------------------------------

### Load Pretrained Model Weights

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Initializes the model and loads weights from a file, ensuring the model is set to evaluation mode.

```python
import torch
from models.bisenet import BiSeNet

def load_model(model_name: str, num_classes: int, weight_path: str, device: torch.device):
    """Load and initialize the BiSeNet model with pretrained weights."""
    model = BiSeNet(num_classes, backbone_name=model_name)
    model.to(device)

    if not os.path.exists(weight_path):
        raise ValueError(f'Weights not found: {weight_path}')

    model.load_state_dict(torch.load(weight_path, map_location=device))
    model.eval()
    return model

# Usage example
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = load_model(
    model_name='resnet18',      # or 'resnet34'
    num_classes=19,
    weight_path='./weights/resnet18.pt',
    device=device
)
print(f"Model loaded on {device}")
```

--------------------------------

### View PyTorch Inference Configuration

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Displays the available command-line arguments for the PyTorch inference script.

```text
usage: inference.py [-h] [--model MODEL] [--weight WEIGHT] [--input INPUT] [--output OUTPUT]

Face parsing inference

options:
  -h, --help       show this help message and exit
  --model MODEL    model name, i.e resnet18, resnet34
  --weight WEIGHT  path to trained model, i.e resnet18/34
  --input INPUT    path to an image or a folder of images
  --output OUTPUT  path to save model outputs
```

--------------------------------

### View ONNX Export Configuration

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Displays the available command-line arguments for the ONNX export script.

```text
usage: onnx_export.py [-h] [--model MODEL] [--weight WEIGHT]

Convert PyTorch model to ONNX format

options:
  -h, --help       show this help message and exit
  --model MODEL    model name, i.e resnet18, resnet34 (default: resnet18)
  --weight WEIGHT  path to trained PyTorch model (default: ./weights/resnet18.pt)
```

--------------------------------

### View ONNX Inference Configuration

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Displays the available command-line arguments for the ONNX inference script.

```text
usage: onnx_inference.py [-h] --model MODEL [--input INPUT] [--output OUTPUT]

Face parsing inference with ONNX

options:
  -h, --help       show this help message and exit
  --model MODEL    path to ONNX model file
  --input INPUT    path to an image or a folder of images
  --output OUTPUT  path to save model outputs
```

--------------------------------

### Execute ONNX Export

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Commands to convert PyTorch models to ONNX format.

```bash
# Export ResNet18 model
python onnx_export.py --model resnet18 --weight ./weights/resnet18.pt

# Export ResNet34 model
python onnx_export.py --model resnet34 --weight ./weights/resnet34.pt
```

--------------------------------

### Command Line Inference with BiSeNet

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Perform face parsing inference using the command line. Supports single image processing and batch processing of entire directories. Specify the model, weights, input path, and output directory.

```bash
python inference.py \
    --model resnet18 \
    --weight ./weights/resnet18.pt \
    --input ./assets/images/1.jpg \
    --output ./results
```

```bash
python inference.py \
    --model resnet34 \
    --weight ./weights/resnet34.pt \
    --input ./assets/images \
    --output ./assets/results
```

--------------------------------

### Configure Data Augmentation Transforms

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Training and validation transforms for image preprocessing and augmentation.

```python
from utils.transform import TrainTransform, DefaultTransform
from PIL import Image

# Training transform with augmentation
train_transform = TrainTransform(image_size=(448, 448))

# Augmentation pipeline includes:
# 1. ColorJitter: brightness=0.5, contrast=0.5, saturation=0.5
# 2. HorizontalFlip: p=0.5, with label swapping (l_eye <-> r_eye, etc.)
# 3. RandomScale: scales=[0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
# 4. RandomCrop: crops to target image_size
# 5. ToTensor: converts to PyTorch tensor
# 6. Normalize: ImageNet mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)

# Default transform (inference/validation)
default_transform = DefaultTransform()
# Only includes ToTensor + Normalize

# Usage example
image = Image.open('face.jpg')
label = Image.open('mask.png').convert('P')

# Apply training transforms
augmented_image, augmented_label = train_transform(image, label)
print(f"Augmented image shape: {augmented_image.shape}")  # torch.Size([3, 448, 448])
```

--------------------------------

### Export BiSeNet to ONNX

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Convert a PyTorch BiSeNet model to ONNX format for deployment. This function supports dynamic batch size and requires specifying the model name and weight path.

```python
import torch
from models.bisenet import BiSeNet

def torch2onnx_export(model_name: str, weight_path: str):
    """Export PyTorch model to ONNX format."""
    num_classes = 19

    # Load model
    model = BiSeNet(num_classes, backbone_name=model_name)
    model.load_state_dict(torch.load(weight_path))
    model.eval()

    # Output path
    onnx_path = weight_path.replace('.pt', '.onnx')

    # Create dummy input
    dummy_input = torch.randn(1, 3, 512, 512, requires_grad=True)

    # Export to ONNX
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=20,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'},
        },
    )
    print(f"Model exported to {onnx_path}")

# Export ResNet18 model
torch2onnx_export('resnet18', './weights/resnet18.pt')

# Command line usage:
# python onnx_export.py --model resnet18 --weight ./weights/resnet18.pt
# python onnx_export.py --model resnet34 --weight ./weights/resnet34.pt
```

--------------------------------

### Preprocess Images for Inference

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Resizes and normalizes PIL images using ImageNet statistics to prepare them for model input.

```python
import torch
from PIL import Image
import torchvision.transforms as transforms

def prepare_image(image: Image.Image, input_size=(512, 512)) -> torch.Tensor:
    """Prepare an image for inference by resizing and normalizing."""
    # Resize the image
    resized_image = image.resize(input_size, resample=Image.BILINEAR)

    # Define transformation pipeline with ImageNet normalization
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(
            mean=(0.485, 0.456, 0.406),
            std=(0.229, 0.224, 0.225)
        ),
    ])

    # Apply transformations and add batch dimension
    image_tensor = transform(resized_image)
    image_batch = image_tensor.unsqueeze(0)

    return image_batch

# Usage example
image = Image.open('./assets/images/1.jpg').convert('RGB')
original_size = image.size  # Store for later restoration
image_batch = prepare_image(image)
print(f"Input tensor shape: {image_batch.shape}")  # torch.Size([1, 3, 512, 512])
```

--------------------------------

### PyTorch Training Loop

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Standard PyTorch training loop for iterating through epochs, processing batches, calculating loss, and updating model weights. Includes saving model checkpoints.

```python
for epoch in range(config['epochs']):
    model.train()
    for batch_idx, (images, labels) in enumerate(data_loader):
        images = images.to(device)
        labels = labels.to(device)

        output = model(images)
        loss = criterion(output, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()

        if (batch_idx + 1) % 50 == 0:
            print(f'Epoch [{epoch}][{batch_idx+1}/{len(data_loader)}] Loss: {loss.item():.4f}')

    # Save checkpoint
    torch.save(model.state_dict(), f'./weights/{config["backbone"]}.pt')
```

--------------------------------

### Execute PyTorch Inference

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Commands to run inference on single images or batches using PyTorch.

```bash
python inference.py --model resnet18 --weight ./weights/resnet18.pt --input ./assets/images/1.jpg --output ./results
```

```bash
python inference.py --model resnet18 --weight ./weights/resnet18.pt --input ./assets/images --output ./assets/results
```

--------------------------------

### Execute ONNX Inference

Source: https://github.com/yakhyo/face-parsing/blob/main/README.md

Command to run inference using an ONNX model.

```bash
python onnx_inference.py --model ./weights/resnet18.onnx --input ./assets/images --output ./assets/results/resnet18onnx
```

--------------------------------

### Command Line ONNX Inference

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Performs batch ONNX inference on a directory of images from the command line. Specify model, input, and output paths.

```bash
# ONNX inference on a directory of images
python onnx_inference.py \
    --model ./weights/resnet18.onnx \
    --input ./assets/images \
    --output ./assets/results

# Expected output:
# 2024-03-20 10:30:00 - ONNX model loaded successfully from ./weights/resnet18.onnx
# 2024-03-20 10:30:00 - Using execution provider: CUDAExecutionProvider
# 2024-03-20 10:30:00 - Found 9 files to process
# Processing images: 100%|████████| 9/9 [00:01<00:00, 6.50it/s]
# 2024-03-20 10:30:02 - Processing complete. Results saved to ./assets/results/resnet18
```

--------------------------------

### ONNX Runtime Face Parsing Inference

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Infers face parsing masks using ONNX Runtime. Auto-detects GPU/CPU. Requires ONNX model and image preprocessing.

```python
import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image
from utils.common import vis_parsing_maps

class FaceParsingONNX:
    """Face parsing inference using ONNXRuntime."""

    def __init__(self, model_path: str):
        # Auto-detect GPU availability
        providers = (
            ['CUDAExecutionProvider', 'CPUExecutionProvider']
            if ort.get_device() == 'GPU'
            else ['CPUExecutionProvider']
        )
        self.session = ort.InferenceSession(model_path, providers=providers)

        self.input_size = (512, 512)
        self.input_mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        self.input_std = np.array([0.229, 0.224, 0.225], dtype=np.float32)

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]

    def preprocess(self, image: np.ndarray) -> np.ndarray:
        """Preprocess BGR image for inference."""
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, self.input_size, interpolation=cv2.INTER_LINEAR)
        image = image.astype(np.float32) / 255.0
        image = (image - self.input_mean) / self.input_std
        image = np.transpose(image, (2, 0, 1))  # HWC -> CHW
        return np.expand_dims(image, axis=0).astype(np.float32)

    def predict(self, image: np.ndarray) -> np.ndarray:
        """Run inference and return segmentation mask."""
        original_size = (image.shape[1], image.shape[0])
        input_tensor = self.preprocess(image)
        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
        mask = outputs[0].squeeze(0).argmax(0).astype(np.uint8)
        return cv2.resize(mask, original_size, interpolation=cv2.INTER_NEAREST)

# Usage example
engine = FaceParsingONNX('./weights/resnet18.onnx')
image = cv2.imread('./assets/images/1.jpg')
mask = engine.predict(image)

# Visualize result
image_pil = Image.open('./assets/images/1.jpg').convert('RGB')
vis_parsing_maps(image_pil, mask, save_image=True, save_path='./onnx_result.jpg')
print(f"Mask shape: {mask.shape}, Unique classes: {np.unique(mask)}")
```

--------------------------------

### Python Visualization Function

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Visualizes segmentation masks overlaid on original images using predefined colors for facial components. Requires OpenCV and Pillow.

```python
import cv2
import numpy as np
from PIL import Image

# Facial component labels (19 classes)
ATTRIBUTES = [
    'skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g',
    'l_ear', 'r_ear', 'ear_r', 'nose', 'mouth', 'u_lip',
    'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat'
]

# Color mapping for visualization (BGR format)
COLOR_LIST = [
    [0, 0, 0],       # background
    [255, 85, 0],    # skin
    [255, 170, 0],   # left eyebrow
    [255, 0, 85],    # right eyebrow
    [255, 0, 170],   # left eye
    [0, 255, 0],     # right eye
    [85, 255, 0],    # glasses
    [170, 255, 0],   # left ear
    [0, 255, 85],    # right ear
    [0, 255, 170],   # earring
    [0, 0, 255],     # nose
    [85, 0, 255],    # mouth
    [170, 0, 255],   # upper lip
    [0, 85, 255],    # lower lip
    [0, 170, 255],   # neck
    [255, 255, 0],   # necklace
    [255, 255, 85],  # cloth
    [255, 255, 170], # hair
    [255, 0, 255],   # hat
]

def vis_parsing_maps(image, segmentation_mask, save_image=False, save_path='result.png'):
    """Visualize segmentation mask overlaid on original image."""
    image = np.array(image).copy().astype(np.uint8)
    segmentation_mask = segmentation_mask.copy().astype(np.uint8)

    # Create colored mask
    mask_color = np.zeros((segmentation_mask.shape[0], segmentation_mask.shape[1], 3))

    for class_idx in range(1, np.max(segmentation_mask) + 1):
        mask_color[segmentation_mask == class_idx] = COLOR_LIST[class_idx]

    mask_color = mask_color.astype(np.uint8)
    bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # Blend original image with mask (60% image, 40% mask)
    blended = cv2.addWeighted(bgr_image, 0.6, mask_color, 0.4, 0)

    if save_image:
        cv2.imwrite(save_path, blended, [int(cv2.IMWRITE_JPEG_QUALITY), 100])

    return blended

# Usage
image = Image.open('./assets/images/1.jpg').convert('RGB')
mask = np.load('predicted_mask.npy')  # Loaded from inference
result = vis_parsing_maps(image, mask, save_image=True, save_path='./visualized.jpg')
```

--------------------------------

### Define CelebAMaskHQ Dataset Class

Source: https://context7.com/yakhyo/face-parsing/llms.txt

Custom PyTorch Dataset class for loading CelebAMask-HQ face parsing data with paired images and segmentation masks.

```python
from torch.utils.data import Dataset
from PIL import Image
import numpy as np
import os

class CelebAMaskHQ(Dataset):
    """Dataset class for CelebAMask-HQ face parsing dataset."""

    def __init__(self, images_dir, labels_dir, transform=None):
        self.images_dir = images_dir
        self.labels_dir = labels_dir
        self.transform = transform or DefaultTransform()

        self.image_files = []
        self.label_files = []

        # Find matching image-label pairs
        for filename in os.listdir(self.images_dir):
            if filename.endswith(('.jpg', '.jpeg', '.png')):
                image_path = os.path.join(self.images_dir, filename)
                label_path = os.path.join(self.labels_dir, f'{filename[:-4]}.png')

                if os.path.isfile(label_path):
                    self.image_files.append(image_path)
                    self.label_files.append(label_path)

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        image = Image.open(self.image_files[idx])
        image = image.resize((512, 512), Image.BILINEAR)

        label = Image.open(self.label_files[idx]).convert('P')

        image, label = self.transform(image, label)
        label = np.array(label).astype(np.int64)

        return image, label
```
Tokens:
Raw