README.md

---
license: apache-2.0
library_name: litert
pipeline_tag: image-segmentation
base_model: xuebinqin/U-2-Net
tags:
  - litert
  - tflite
  - on-device
  - android
  - background-removal
  - salient-object-detection
  - image-matting
  - u2net
---

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

On-device [LiteRT](https://ai.google.dev/edge/litert) (`.tflite`) conversion of
**[U²-Net](https://github.com/xuebinqin/U-2-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.

![U²-Net — input, saliency mask, background removed (on-device LiteRT GPU)](samples/sample.png)

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`](https://github.com/google-ai-edge/ai-edge-torch) **with no custom
rewrites** (pure CNN).

## Files

| File | Size | Description |
|------|------|-------------|
| `u2net_fp16.tflite` | 88 MB | float16 weights, GPU-compatible |

## 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)**

```kotlin
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)**

```python
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](https://github.com/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](https://github.com/xuebinqin/U-2-Net) repo for dataset details.

## License & attribution

- License: **Apache-2.0** (© the U²-Net authors,
  [xuebinqin/U-2-Net](https://github.com/xuebinqin/U-2-Net/blob/master/LICENSE)).
- This is a format conversion of the official U²-Net weights (no architectural changes);
  all credit to the original authors.