Catrex Lite Image Segmentation

A compact neural network for image segmentation, trained from scratch.

Results

ONNX

import numpy as np, onnxruntime as ort
from PIL import Image
from huggingface_hub import hf_hub_download

sess = ort.InferenceSession(hf_hub_download("Catniti/catrex-lite-image-segmentation", "model.onnx"))

img = Image.open("photo.jpg").convert("RGB")
x = np.array(img.resize((512, 512))).astype(np.float32) / 255.
x = (x - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
x = x.transpose(2, 0, 1)[None].astype(np.float32)

mask = sess.run(None, {"input": x})[0][0, 0]
mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(img.size)

cutout = Image.fromarray(np.dstack([np.array(img), np.array(mask)]), "RGBA")
cutout.save("no_background.png")

onnxruntime-web:

<script src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js"></script>
<script type="module">
const url = "https://huggingface.co/Catniti/catrex-lite-image-segmentation/resolve/main/model.onnx";
const session = await ort.InferenceSession.create(url);
const { mask } = await session.run({ input: tensor });
</script>

Node.js:

npm i onnxruntime-node

Python without internet:

hf download Catniti/catrex-lite-image-segmentation model.onnx --local-dir .