rmsandu/ffhq-wrinkle-unet

FFHQ-Wrinkle U-Net (ResNet50 encoder)

A U-Net for pixel-level facial wrinkle segmentation, with a pretrained ResNet50 encoder and optional attention-gated skip connections. Trained on the manually-labeled subset of the FFHQ-Wrinkle dataset.

Code, training scripts, and the Gradio demo: rmsandu/FFHQ-detect-face-wrinkles. Background and methodology write-up: "Segmentation of Fine Facial Wrinkles with U-Net" by Raluca-Maria Sandu.

Model description

• Architecture: U-Net with a torchvision ResNet50 encoder (resnet50 layers 1–4 as the four downsampling stages) and a custom decoder (DoubleConv + transposed-convolution upsampling blocks). Attention gates (AttentionGate) can be enabled on the skip connections.
• Input: 512×512 RGB face crop, ImageNet-normalized (mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)).
• Output: single-channel logit map at 512×512; apply sigmoid then threshold (0.5 in the reference demo) for a binary wrinkle mask.
• Loss: binary_focal_loss (α=0.9, γ=1.0) + Dice loss, chosen to handle the extreme class imbalance in wrinkle pixels (~0.03% of all pixels in the training set).
• Checkpoint: wrinkle_model.safetensors, converted from the training checkpoint at epoch 31 (best validation IoU: 0.3254). Only model weights are included — optimizer/scheduler state from the original .pth checkpoint was dropped in conversion.

Intended use

Research and experimentation with facial wrinkle segmentation — e.g. cosmetic dermatology research, aging-related computer vision work, or as a component upstream of a wrinkle-severity metric. See the linked code repo's app.py for a full inference pipeline (face detection → face parsing/masking → this model → thresholded overlay).

Not intended for: medical diagnosis, clinical decision-making, or any use where segmentation errors could cause harm. This is a research checkpoint with a validation IoU of ~0.33 — it is not a high-precision instrument.

Limitations & bias

• Trained on FFHQ-derived faces; performance on populations, lighting conditions, or camera setups not well represented in FFHQ (or in the 1,000 manually-labeled wrinkle masks specifically) is unverified.
• The extreme class imbalance (wrinkles are a tiny fraction of pixels) makes the model's practical recall/precision tradeoff sensitive to the choice of decision threshold — the default 0.5 threshold is not necessarily optimal for every use case.
• "Wrinkle" itself is a label with inherent annotator disagreement (see the project's README for discussion of this); ground truth reflects one team's manual annotation, not a clinical consensus.

Usage

from safetensors.torch import load_file
from unet import UNet  # from the code repo: rmsandu/FFHQ-detect-face-wrinkles

state_dict = load_file("wrinkle_model.safetensors", device="cpu")
model = UNet(n_channels=3, n_classes=1, bilinear=False, pretrained=False, freeze_encoder=True)
model.load_state_dict(state_dict)
model.eval()

Or fetch it programmatically via the code repo's scripts/download_weights.py.

License

The code in the linked repository is MIT-licensed. These weights were trained on the FFHQ-Wrinkle dataset and are released here under CC BY-NC-SA 4.0 — non-commercial use only, share-alike, attribution required — consistent with the dataset's own license.

Citation

This model and the manually-labeled training set it was trained on were put together by Raluca-Maria Sandu. If you use this model, please cite:

@misc{sandu2025wrinkle,
  title={Segmentation of Fine Facial Wrinkles with U-Net},
  author={Sandu, Raluca-Maria},
  howpublished={\url{https://rmsandu.net/blog/2025-04-18-wrinkle-segmentation.html}},
  year={2025}
}

The underlying FFHQ-Wrinkle dataset itself should also be cited:

@article{moon2024facial,
  title={Facial Wrinkle Segmentation for Cosmetic Dermatology: Pretraining with Texture Map-Based Weak Supervision},
  author={Moon, Junho and Chung, Haejun and Jang, Ikbeom},
  journal={arXiv preprint arXiv:2408.10060},
  year={2024}
}