a-ml/yolo26-face

YOLO26 Face Parsing — 19-class, realtime, Core ML

Two YOLO26 semantic segmentation models fine-tuned on CelebAMask-HQ for 19-class face parsing, plus Core ML exports that run in ~1.1 ms on the Apple Neural Engine.

Test = the 2 824-image held-out split, never seen during training or model selection. Test ≈ val for both models, so neither overfits. The nano reaches 96.3 % of the large model's mIoU with 11× fewer parameters — the gap is concentrated almost entirely in the rarest classes.

Latency measured on an Apple M5 Max ANE with computeUnits = .cpuAndNeuralEngine. Expect an iPhone to be slower; both still leave large headroom against a 33 ms/30 fps budget.

Demo

Live front-camera overlay on iPhone, nano model, ~30 fps.

<video controls src="https://huggingface.co/a-ml/yolo26-face/resolve/main/demo.mp4"></video>

Contents

checkpoints/  yolo26{n,l}-sem-celebamaskhq.pt    Ultralytics checkpoints (resume/fine-tune/val)
coreml/       yolo26n-sem-face-fp16.mlpackage    the shipping realtime model
              yolo26l-sem-face-int8.mlpackage    the accuracy variant
scripts/      the full pipeline, dataset prep -> training -> export -> evaluation
assets/       per-epoch metrics + qualitative contact sheets
demo.mp4      iPhone screen recording

Classes

0 background · 1 skin · 2 nose · 3 eyeglasses · 4 left_eye · 5 right_eye · 6 left_eyebrow · 7 right_eyebrow · 8 left_ear · 9 right_ear · 10 mouth · 11 upper_lip · 12 lower_lip · 13 hair · 14 hat · 15 earring · 16 necklace · 17 neck · 18 cloth

Index order matches the canonical CelebAMask-HQ g_mask.py ordering, so later classes overwrite earlier ones where the source masks overlap.

Results

Per-class mIoU on the CelebAMask-HQ validation split:

On the classes that dominate a face overlay — skin, hair, background — the two models are within 0.015 of each other. Choose the nano unless you specifically need jewellery.

Cross-dataset generalisation

Neither model saw these images. No ground truth, so this reports mean pixel confidence (qualitative sheets in assets/):

Confidence is consistently highest on AI-generated faces (unusually prototypical) and lowest on the low-resolution img_align_celeba crops.

Fine-tuning methodology

Initialised from the official Cityscapes-pretrained yolo26{n,l}-sem.pt and fine-tuned for 12 epochs at 512 px on the canonical CelebA partition (24 183 train / 2 993 val / 2 824 test), batch 32, AdamW (optimizer=auto, lr ≈ 4.3e-4), on a single Apple M5 Max via MPS. Large: 7.77 h. Nano: 5.55 h.

python scripts/prepare_celebamaskhq_semantic.py --mask-size 512 --img-mode resize
python scripts/train.py --model yolo26l-sem.pt --epochs 12 --imgsz 512 --batch 32 \
       --cls-pw 1.0 --loss-size 256 --workers 12 --mosaic 0.5 --no-deterministic --amp

Four choices did most of the work. They are documented because three of them are silent failures — the defaults produce a plausible-looking model that is quietly wrong.

1. `cls_pw=1.0` — the single biggest quality win. Ultralytics computes ENet inverse-log class weights as (1/ln(1.02+p))**cls_pw but defaults the exponent to `0.0`, which makes every weight exactly 1.0; set_class_weights then returns early, so class weighting is silently disabled. On a dataset this imbalanced (necklace is 0.02 % of pixels, hair is 32 %) that is fatal for rare classes. Setting cls_pw=1.0 (the maximum an internal assert allows) gives rare classes ~14× the weight of hair/skin:

(Left column is an early short run, so it is indicative rather than a controlled ablation — but classes scoring exactly zero are attributable to the missing weighting.)

2. `fliplr=0.0` — mandatory for this label set. Ultralytics' RandomFlip mirrors geometry without swapping left/right class labels, so a horizontally flipped left_eye is still labelled left_eye. Nine of the 19 classes are laterally paired (eyes, brows, ears). Any fliplr > 0 teaches the model to confuse them. Rotation (degrees=10), translation, scale and mosaic are used instead.

3. `--loss-size 256` — ~5× faster training for −0.002 mIoU. The stock semantic loss upsamples the stride-8 logits to full mask resolution before CE/Dice — a 159 M-element tensor at 512 px / batch 32 — then softmaxes it and boolean-index-gathers it (a data-dependent shape that forces a device sync), twice counting the auxiliary head. The loss ends up costing more than the entire forward+backward pass. scripts/loss_patch.py evaluates it at a fixed lower resolution instead. A/B at 320 px:

Profiling first mattered: the dataloader was not the bottleneck (it sustains 1900–4000 img/s). See scripts/isolate_bottleneck.py.

4. Apple-Silicon throughput. deterministic=True (a default) more than doubles loss cost on MPS by forcing slow scatter kernels. AMP helps on MPS (8.1 vs 12.6 s/iter). Batch 64 is worse per-image than batch 32. Ultralytics also hard-forces workers=0 on MPS (trainer.py:162, a CPU-era assumption); scripts/fast_trainer.py restores them and adds validation throttling, since the 2 993-image val split otherwise costs ~18 min/epoch.

Core ML export

yolo export format=coreml and a naive manual conversion both fail with coremltools 9 + numpy ≥ 2:

TypeError: only 0-dimensional arrays can be converted to Python scalars

coremltools' _cast op handler does mb.const(val=int(x.val)), and numpy 2 forbids int() on a size-1, >0-d array. YOLO26's attention block emits exactly that cast. scripts/coreml_patch.py re-registers _cast with .item() coercion — import it before ct.convert. scripts/export_coreml.py then traces the raw module (fp16 logits, no baked argmax) and compares quantization variants.

Quantize-then-convert vs convert-then-quantize, measured on the large model:

Convert-then-quantize wins — better fidelity, smaller and faster than the reverse. Judge palettization by mIoU rather than argmax agreement: 99.1 % agreement hides a drop to 0.924 mIoU, because the disagreements land on small classes.

Reproduce with scripts/compare_quant_order.py.

Using the Core ML models

• —input image — 512×512 RGB (CVPixelBuffer; Core ML handles BGRA→RGB)
• —output logits — MLMultiArray float16, (1, 19, 64, 64), NCHW

The output is a stride-8 logit grid, not an argmax map — bilinearly upsample and argmax over the 19 channels on the GPU. Doing it in a fragment shader keeps the whole network on the ANE and costs nothing measurable:

half best = logitsTex.sample(bilinear, modelUV, 0).r;
ushort cls = 0;
for (ushort c = 1; c < 19; c++) {
    half v = logitsTex.sample(bilinear, modelUV, c).r;
    if (v > best) { best = v; cls = c; }
}

Set configuration.computeUnits = .cpuAndNeuralEngine. Avoid .all: it aborts in MPSGraph (MLIR pass manager failed) for these shapes on some machines, and keeping the model off the GPU leaves it free for rendering.

⚠️ Feed the model an upright face. It was trained exclusively on upright, roughly centred faces, so a rotated input degrades output badly enough to look like a broken model. AVCaptureConnection.videoRotationAngle is silently ignored on some device/format combinations — verify the buffer dimensions rather than trusting it.

Letterbox the whole frame, don't centre-crop it. The input is square while camera frames are not. Cropping an inscribed square (min(w, h)) is tempting because it keeps the subject at full resolution, but it produces no logits at all outside that square — a face near the top or bottom of a portrait frame simply gets no mask. Scale the whole frame to fit and pad the short axis (max(w, h)) instead. A 9:16 frame then fills ~56 % of the input width, which costs some fine detail but covers every visible pixel.

Using the PyTorch checkpoints

from ultralytics import YOLO
model = YOLO("checkpoints/yolo26n-sem-celebamaskhq.pt")
result = model("face.jpg")[0]
class_map = result.semantic_mask.data          # HxW int class indices

Validate against your own copy of the dataset (edit path in scripts/data.yaml):

yolo semantic val model=checkpoints/yolo26l-sem-celebamaskhq.pt data=scripts/data.yaml \
     imgsz=512 split=test

Limitations

• —`necklace` is weak (0.131 nano / 0.256 large) and earring is moderate. Only 143 validation images contain a necklace, totalling ~163 k pixels. This is a dataset limit, not a training one.
• —Upright, roughly centred faces only — see the warning above.
• —CelebAMask-HQ demographic bias carries over; the training distribution is celebrity photography and is not representative. Validate on your own population before relying on it, and do not use it for identification, surveillance, or any consequential decision.
• —The 64×64 logit grid is coarse for very thin structures; boundaries are smoothed by the bilinear upsample.
• —Latency figures are from an M5 Max ANE, not a phone.

Licensing — read before use

• —Ultralytics is AGPL-3.0. These weights are derived from Ultralytics code and pretrained checkpoints, so the AGPL's copyleft terms apply. Shipping them inside a closed-source application may require a commercial licence from Ultralytics.
• —CelebAMask-HQ is non-commercial: research and educational use only. The dataset authors do not own the underlying image copyrights.

Both apply to anything you build from these artifacts. This repository is published for research and educational purposes.

Citation

@article{CelebAMask-HQ,
  title   = {MaskGAN: Towards Diverse and Interactive Facial Image Manipulation},
  author  = {Lee, Cheng-Han and Liu, Ziwei and Wu, Lingyun and Luo, Ping},
  journal = {Technical Report},
  year    = {2019}
}