File size: 3,429 Bytes
bd2db50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | ---
license: mit
library_name: litert
pipeline_tag: image-segmentation
tags:
- litert
- tflite
- android
- on-device
- gpu
- semantic-segmentation
- cityscapes
- real-time
- pidnet
---
# 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](https://arxiv.org/abs/2206.02066)
(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](https://github.com/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)
```kotlin
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)
```python
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.
|