litert-community/PIDNet-S-Cityscapes-LiteRT

PIDNet-S — LiteRT (real-time semantic segmentation, GPU)

On-device real-time semantic segmentation running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). PIDNet-S (CVPR 2023) segments a road scene into the 19 Cityscapes classes at ~17 FPS on a Pixel 8a.

• Architecture: PIDNet-S — a three-branch CNN (P: detail, I: context, D: boundary).
• Weights: XuJiacong/PIDNet · MIT · 78.8% mIoU (Cityscapes val).
• Size: 30 MB · ~7.6 M params · pure CNN.

I/O

• Input: [1, 3, 1024, 1024] NCHW, RGB, ImageNet-normalized (mean [0.485,0.456,0.406], std [0.229,0.224,0.225]).
• Output: [1, 19, 128, 128] class logits at 1/8 resolution — argmax over the 19 classes per pixel, then upscale (nearest) to display.

Classes (index order): road, sidewalk, building, wall, fence, pole, traffic light, traffic sign, vegetation, terrain, sky, person, rider, car, truck, bus, train, motorcycle, bicycle.

GPU conversion

PIDNet is a pure CNN — no attention, no dynamic shapes at a fixed input size, and align_corners=False on every bilinear resize. It converts to a fully GPU-compatible graph with zero patches: CONV_2D ×75, RESIZE_BILINEAR ×11 (align_corners=False), AVERAGE_POOL_2D, ADD/MUL/SUB/SUM, LOGISTIC — 0 tensors of rank > 4, 0 GPU-incompatible ops. The converted graph matches the original PyTorch model bit-for-bit on CPU (corr 0.99999999999, 100% argmax); on the Mali GPU (fp16) it agrees with the fp32 reference at 97% of pixels with correct classes.

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

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

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

// argmax over 19 classes per pixel:
val hw = 128 * 128
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="pidnet_s.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
it.set_tensor(inp[0]["index"], x)            # [1,3,1024,1024] float32, ImageNet-norm
it.invoke()
logits = it.get_tensor(out[0]["index"])[0]   # [19,128,128]
label = logits.argmax(0)                      # [128,128] class ids

Conversion

Re-authored/converted with litert-torch (build_pidnet.py): the trained PIDNet-S weights are loaded from an ONNX mirror whose initializer names match the original repo's PyTorch keys, then converted directly — zero GPU patches.

License

MIT (PIDNet / XuJiacong/PIDNet). Cityscapes label taxonomy from the Cityscapes dataset.