mlboydaisuke's picture
Upload README.md with huggingface_hub
f2549f0 verified
|
Raw
History Blame Contribute Delete
2.76 kB
---
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.