nobg/DIS5K
Viewer • Updated • 5.47k • 1.26k • 3
A compact neural network for image segmentation, trained from scratch.
| F1 on DIS-VD | 0.6461 |
| MAE | 0.1178 |
| Resolution | 512×512 |
| Parameters | 4.63M |
| Training | 1× T4, 40 epohs |
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 .