litert-community/MODNet-LiteRT

MODNet — LiteRT (trimap-free portrait matting, GPU)

On-device real-time portrait matting running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). MODNet (AAAI 2022) predicts a soft alpha matte for a person — no trimap, no green screen — for background blur/replace (video calls, virtual backgrounds). ~79 ms/frame on a Pixel 8a.

• Architecture: MODNet — MobileNetV2 low-res branch + high-res + fusion branches (pure CNN).
• Weights: ZHKKKe/MODNet · Apache-2.0 · ~6.5 M params.
• Size: 26 MB.

I/O

• Input: [1, 3, 512, 512] NCHW, RGB, normalized to [-1, 1] ((x/255 - 0.5) / 0.5).
• Output: [1, 1, 512, 512] soft alpha matte in [0, 1] (composite: fg·α + bg·(1-α)).

GPU conversion

MODNet is a pure CNN with align_corners=False interpolation. Two re-authoring patches make it a fully GPU-compatible graph — 0 tensors of rank > 4, 0 banned ops:

1. SE block Linear → 1×1 conv — the stock squeeze-excite pool → Linear → view(b,c,1,1) → x*w confuses the NCHW↔NHWC layout (mul broadcast mismatch); 1×1 convs on the pooled tensor are identical and NCHW-clean.
2. fp16-safe hierarchical-mean InstanceNorm — MODNet's IBNorm runs InstanceNorm2d over up to 512×512 spatial; on the Mali GPU (fp16) the variance sum(dd²) overflows (≫ 65504) and the matte degrades (halos, blotchy interior, corr 0.94). Computing the spatial mean via a cascade of /2 average-pools (magnitude-bounded, exact for power-of-2) + dd·rsqrt(mean(dd²)+eps) restores it to GPU corr 0.99994 with clean edges.

CPU-exact vs PyTorch (corr 0.99999999999); device Mali GPU corr 0.99994.

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

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

inBufs[0].writeFloat(inputNCHW)          // [1,3,512,512], RGB, (x/255-0.5)/0.5
model.run(inBufs, outBufs)
val alpha = outBufs[0].readFloat()       // [512*512] soft matte in [0,1]
// composite:  out = fg*alpha + bg*(1-alpha)

Python (LiteRT / ai-edge-litert)

from ai_edge_litert.interpreter import Interpreter
import numpy as np

it = Interpreter(model_path="modnet.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
x = ((img[None].transpose(0,3,1,2) / 255.0 - 0.5) / 0.5).astype(np.float32)  # [1,3,512,512]
it.set_tensor(inp[0]["index"], x); it.invoke()
alpha = it.get_tensor(out[0]["index"])[0, 0]   # [512,512] in [0,1]

Conversion

Converted with litert-torch (build_modnet.py): loads the trained MODNet weights, applies the two patches (SE 1×1-conv, SafeInstanceNorm), and exports.

License

Apache-2.0 (MODNet / ZHKKKe/MODNet).