---
license: mit
library_name: litert
pipeline_tag: image-segmentation
tags:
  - litert
  - tflite
  - android
  - on-device
  - gpu
  - clothing-segmentation
  - fashion
  - virtual-try-on
  - u2net
---

# Cloth Segmentation (U²-Net) — LiteRT GPU

On-device **clothing segmentation** running **fully on the LiteRT `CompiledModel` GPU**
delegate (no CPU fallback). [cloth-segmentation](https://github.com/levindabhi/cloth-segmentation)
is a U²-Net trained on iMaterialist-Fashion to segment **upper-body / lower-body /
full-body clothing** — the building block for virtual try-on and fashion apps. ~88 ms/frame
on a Pixel 8a.

- **Architecture:** U²-Net (RSU nested residual blocks), 4-class head — pure CNN.
- **Weights:** [levindabhi/cloth-segmentation](https://github.com/levindabhi/cloth-segmentation) (iMaterialist-Fashion) · MIT.
- **Size:** 176 MB.

![Cloth segmentation](hero.png)

*Upper-body clothing (cyan) + lower-body (orange). Photo: Unsplash (free license).*

## I/O

- **Input:** `[1, 3, 768, 768]` NCHW, RGB, `(x/255 - 0.5)/0.5` (i.e. [-1, 1]).
- **Output:** `[1, 4, 768, 768]` logits — `argmax` over the 4 classes: 0 = background,
  1 = upper body, 2 = lower body, 3 = full body (dress).

## GPU conversion

U²-Net is a pure CNN → fully GPU-compatible (**254/254 nodes on the delegate, 1
partition**; device corr 0.999798, ~88 ms) with **one defensive patch**: `align_corners=True`
→ `False` on the bilinear upsamples (the GPU delegate rejects `align_corners=True`).
CPU-exact vs PyTorch (corr 1.0).

## Minimal usage

### Kotlin (Android, LiteRT CompiledModel GPU)

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

inBufs[0].writeFloat(inputNCHW)          // [1,3,768,768] RGB, (x/255-0.5)/0.5
model.run(inBufs, outBufs)
val out = outBufs[0].readFloat()         // [4*768*768]; per pixel p argmax over the 4 class planes
// class 0 bg, 1 upper, 2 lower, 3 full-body
```

### Python (LiteRT / ai-edge-litert)

```python
import numpy as np
from ai_edge_litert.interpreter import Interpreter

it = Interpreter(model_path="clothseg.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
it.set_tensor(inp[0]["index"], x)        # [1,3,768,768] float32, RGB, (x/255-0.5)/0.5
it.invoke()
seg = it.get_tensor(out[0]["index"])[0].argmax(0)   # [768,768] 0=bg 1=upper 2=lower 3=full
```

## Conversion

Converted with **litert-torch** (`build_clothseg.py`): loads the MIT U²-Net cloth weights
and exports the 4-class graph.

## License

MIT (cloth-segmentation / levindabhi). Trained on iMaterialist-Fashion-2019.