litert-community/Cloth-Segmentation-U2Net-LiteRT

Cloth Segmentation (U²-Net) — LiteRT GPU

On-device clothing segmentation running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). 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 (iMaterialist-Fashion) · MIT.
• Size: 176 MB.

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)

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)

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.