litert-community/BiSeNet-Face-Parsing-LiteRT

BiSeNet Face Parsing — LiteRT (GPU)

On-device real-time face parsing running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). BiSeNet (zllrunning/face-parsing.PyTorch) segments a face into the 19 CelebAMask-HQ classes (skin, brows, eyes, nose, lips, ears, hair, hat, glasses, neck, cloth, …) — for AR / beauty / makeup. ~22 ms/frame on a Pixel 8a.

• Architecture: BiSeNet (ResNet18 backbone + context path + feature-fusion) — pure CNN.
• Weights: zllrunning/face-parsing.PyTorch · MIT · ~13.3 M params.
• Size: 53 MB.

I/O

• Input: [1, 3, 512, 512] NCHW, RGB, ImageNet-normalized (mean [0.485,0.456,0.406], std [0.229,0.224,0.225]).
• Output: [1, 19, 512, 512] class logits — argmax over the 19 classes per pixel.

Classes: background, skin, l_brow, r_brow, l_eye, r_eye, eyeglass, l_ear, r_ear, earring, nose, mouth, u_lip, l_lip, neck, necklace, cloth, hair, hat.

GPU conversion

BiSeNet is a pure CNN; three re-authoring patches make it a fully GPU-compatible graph — 74/74 nodes on the delegate, 1 partition (device corr 0.99999, argmax 99.96% vs PyTorch):

1. align_corners=True → False — the output upsamples use align_corners=True, which the GPU delegate rejects (1.6% argmax change vs original).
2. global avg_pool2d(x, x.size()[2:]) → mean([2,3]) — the context/attention modules pool with a full-spatial kernel, which the Mali delegate rejects as an AVERAGE_POOL_2D; a MEAN reduce is supported.
3. zero-pad maxpool — the ResNet stem MaxPool2d(padding=1) lowers to a PADV2 with -inf padding (PADV2: src has wrong size on Mali); an explicit 0-pad + unpadded maxpool is exact (input is post-ReLU ≥ 0).

CPU-exact vs PyTorch (corr 0.99999999999).

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val options = CompiledModel.Options(Accelerator.GPU)
val model = CompiledModel.create(context.assets, "faceparsing.tflite", options, null)
val inBufs = model.createInputBuffers()
val outBufs = model.createOutputBuffers()

inBufs[0].writeFloat(inputNCHW)          // [1,3,512,512], RGB, ImageNet-norm
model.run(inBufs, outBufs)
val logits = outBufs[0].readFloat()      // [19,512,512] (NCHW, batch dropped)

val hw = 512 * 512
val label = IntArray(hw) { i ->
    var best = 0; var bv = logits[i]
    for (c in 1 until 19) { val v = logits[c * hw + i]; if (v > bv) { bv = v; best = c } }
    best
}

Python (LiteRT / ai-edge-litert)

from ai_edge_litert.interpreter import Interpreter
import numpy as np

it = Interpreter(model_path="faceparsing.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
it.set_tensor(inp[0]["index"], x)        # [1,3,512,512] float32, ImageNet-norm
it.invoke()
logits = it.get_tensor(out[0]["index"])[0]   # [19,512,512]
label = logits.argmax(0)                      # [512,512] class ids

Conversion

Converted with litert-torch (build_faceparsing.py): loads the trained BiSeNet weights, applies the three patches, and exports.

License

MIT (BiSeNet / zllrunning/face-parsing.PyTorch). CelebAMask-HQ label taxonomy.