litert-community / U-2-Net

U²-Net — LiteRT (TFLite) GPU, FP16

On-device LiteRT ( .tflite ) conversion of U²-Net for salient-object segmentation / background removal . U²-Net is a nested U-structure ("U-net of U-nets", a pure CNN) that predicts a single-channel saliency mask; the foreground is composited onto transparency to cut the subject out of its background.

The model runs fully on the LiteRT CompiledModel GPU accelerator (ML Drift): every op is GPU-native, no CPU fallback, no Flex ops. It converts with litert-torch with no custom rewrites (pure CNN).

Files

I/O

• Input : [1, 3, 320, 320] float32, NCHW , RGB. Preprocessing: resize to 320×320, divide by the per-image max, then ImageNet normalize ( mean = [0.485, 0.456, 0.406] , std = [0.229, 0.224, 0.225] ).
• Output : [1, 1, 320, 320] saliency mask in [0, 1] (sigmoid). Upscale to the input size and use as the foreground alpha.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "u2net_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(chw)            // [1,3,320,320] /max then ImageNet-norm, NCHW
model.run(inputs, outputs)
val mask = outputs[0].readFloat()    // [1,1,320,320] saliency in [0,1]

Python (desktop verification)

MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD  = np.array([0.229, 0.224, 0.225], np.float32)
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

orig = Image.open("photo.jpg").convert("RGB")
a = np.asarray(orig.resize((320, 320)), np.float32)
a = a / a.max()                                                   # per-image max, then ImageNet
x = ((a - MEAN) / STD).transpose(2, 0, 1)[None]                   # [1,3,320,320]

it = Interpreter(model_path="u2net_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
m = it.get_tensor(it.get_output_details()[0]["index"])[0, 0]      # [320,320], [0,1]
alpha = Image.fromarray((m * 255).astype(np.uint8)).resize(orig.size)
cutout = orig.copy(); cutout.putalpha(alpha)                      # foreground on transparency
cutout.save("cutout.png")

A complete Android sample (live camera + gallery background removal) is available in google-ai-edge/litert-samples .

Performance

• ~147 ms / frame on a Pixel 8a (Tensor G3, Mali) GPU.

Conversion notes

Converted with litert-torch (full U2NET, 44M params) and float16-quantized with ai-edge-quantizer . Verified: all ops GPU-native, output correlation = 1.0 vs the PyTorch reference (FP32), ~0.9999 for the FP16 build.

Training data & PII

This is a weights-exact format conversion of the public U²-Net salient-object-detection model; no new training was performed. U²-Net was trained on the DUTS-TR saliency dataset (web images with binary salient-object masks). Such web images may incidentally contain people and other PII; none was deliberately collected and this conversion adds none. The model outputs a saliency mask only and performs no identification. Apply your own content/PII filtering before deployment. See the original U²-Net repo for dataset details.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator — the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

Snapdragon NPU (Hexagon)

The NPU is 4.17x faster than the GPU (8.82 ms against 36.80 ms) and loads 11.10x faster (146 ms against 1619 ms).

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16) with LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.82, where 1.0 is the throttling threshold.

The NPU rows ran the published file unchanged. LiteRT compiled it for the Hexagon on the device at first load. That first compile took 5.6 s here. The load column above is the cached load every later run pays. Recipe and the runtime libraries it needs: NPU guide .

GPU wiring: GPU guide .

Raspberry Pi 5 (CPU)

Measured on a Raspberry Pi 5 Model B Rev 1.1 (8 GB, Raspberry Pi OS 64-bit) with the LiteRT benchmark_model tool from litert-cli-nightly 0.2.0.dev20260805: CPU inference (XNNPACK, 4 threads), 3 invocations per file of 10 warm-up plus 50 timed runs (the tool caps a phase at 150 s, so very slow graphs run fewer — the Runs column is the actual timed total). The latency is the median across invocations; the spread is the min–max over all timed runs. No thermal throttling occurred during these runs ( vcgencmd get_throttled stayed 0x0 ).

License & attribution

• License: Apache-2.0 (© the U²-Net authors, xuebinqin/U-2-Net ).
• This is a format conversion of the official U²-Net weights (no architectural changes); all credit to the original authors.