| --- |
| license: apache-2.0 |
| library_name: pytorch |
| pipeline_tag: image-segmentation |
| tags: |
| - image-segmentation |
| - background-removal |
| - onnx |
| - catrex |
| datasets: |
| - nobg/DIS5K |
| metrics: |
| - f1 |
| --- |
| |
| # Catrex Lite Image Segmentation |
|
|
| Компактная сегментационная сеть, обученная **с нуля** (случайная инициализация, |
| без предобученных весов) на DIS5K для удаления фона. |
|
|
| ## Результаты |
|
|
| | | | |
| |---|---| |
| | F1 на DIS-VD | **0.6461** | |
| | MAE | 0.1178 | |
| | Разрешение | 512×512 | |
| | Параметров | 4.63M | |
| | Обучение | 1× T4, 40 эпох | |
|
|
| ## Файлы |
|
|
| | Файл | Зачем | |
| |---|---| |
| | `model.safetensors` | веса для PyTorch, безопасный формат | |
| | `model.onnx` | инференс без PyTorch, динамический размер входа | |
| | `config.json` | параметры архитектуры | |
|
|
| ## Использование (ONNX — проще всего) |
|
|
| ```python |
| 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") |
| ``` |
|
|
| ## Ограничения |
|
|
| Обучалась на одной T4 в разрешении 512px, поэтому уступает BiRefNet |
| и RMBG-2.0 на тонких деталях: волосы, мех, прозрачные и полупрозрачные объекты. |
| На простых сюжетах с чётко выраженным объектом работает уверенно. |
|
|
| Обучающие данные — DIS5K (Apache-2.0), веса под той же лицензией. |
|
|