MediaPipe selfie multiclass (256×256), ONNX

This is Google's MediaPipe selfie_multiclass_256x256 image segmenter, converted from TFLite to ONNX so it runs anywhere onnxruntime does, with no MediaPipe or TensorFlow at inference time. The weights are Google's and unchanged. Only the file format is different. See the original model card: MediaPipe image segmenter models.

It labels each pixel as one of six classes:

index class
0 background
1 hair
2 body skin
3 face skin
4 clothes
5 others (accessories)

Usage

  • Input input_29: float32 [1, 256, 256, 3], NHWC, RGB, values in [0, 1]. Stretch the image to 256×256; don't letterbox it.
  • Output Identity: float32 [1, 256, 256, 6]. These are logits. MediaPipe applies the softmax outside the model, so apply it yourself before treating them as probabilities.
import numpy as np, onnxruntime as ort
from PIL import Image

sess = ort.InferenceSession("onnx/model.onnx", providers=["CPUExecutionProvider"])
img = Image.open("photo.jpg").convert("RGB")
x = (np.asarray(img.resize((256, 256), Image.BILINEAR), np.float32) / 255.0)[None]
logits = sess.run(None, {"input_29": x})[0][0]
e = np.exp(logits - logits.max(-1, keepdims=True))
probs = e / e.sum(-1, keepdims=True)            # 256 x 256 x 6
hair = Image.fromarray((probs[..., 1] * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR)

This takes about 35 ms on a laptop CPU (Intel Core Ultra 9 185H, onnxruntime 1.30).

The masks are 256×256, so individual strands at the edge of the hair are not resolved. We multiply the upsampled probabilities by a high-resolution subject matte (BiRefNet), which removes background false positives and brings back the flyaway edge. MediaPipe's model card lists the model's intended use and limits (selfie framing, people within a few metres); read it before relying on it elsewhere.

Conversion

From the versioned TFLite file:

  • source: https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/1/selfie_multiclass_256x256.tflite
  • source sha256: c6748b1253a99067ef71f7e26ca71096cd449baefa8f101900ea23016507e0e0
pip install tf2onnx==1.17.0 tensorflow   # conversion only
python -m tf2onnx.convert --tflite selfie_multiclass_256x256.tflite --opset 17 --output model.onnx

The result was checked against the TFLite model run through LiteRT on three photographs: the largest absolute difference is 3.7e-5, and the per-pixel argmax agrees 100%.

sha256 of onnx/model.onnx: 35ec1ecd9ee7f85073c99c00020b7f6751b69506eeacf683bc8665f6117f85b0

This conversion is used by painter (painter.parts), for hair, face and clothes masks in its layered painter.

License

Apache-2.0, the same as the original model. The model and its weights are © Google LLC. This repository only changes the file format.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support