Musa Mahmood
cb21b73460
- segmentation-plan.html: the plan for panoptic segmentation of images, per library folder, stored
in .mcap chunks (no file version bump), in six phases.
- tools/export_mask2former.py: exports facebook/mask2former-swin-base-coco-panoptic to ONNX (opset 17,
dynamic batch, 384 x 384 input) and checks it against PyTorch: on 25 images, all 180 segments
identical (IoU 1.000). Writes the class list with names for people ("sky-other-merged" is "sky"),
and PyTorch's own segmentations for segment-probe -check.
- modules/Segment: runs the model with onnxruntime (CUDA, else the CPU) and makes segments of its
outputs as transformers' post_process_panoptic_segmentation does: a segment id map (longest side at
most 1024) and each segment's class, score, area and bounds. Matches PyTorch on 99.8% of pixels,
every segment scored 0.6 or more found. 57 ms an image at batch 4 on an RTX 5080 laptop.
- tools/segment_probe.jai (jai first.jai - segment_probe): segments image files; -out draws them,
-check compares with PyTorch, -dump-input writes the model's input.
- modules/onnxruntime: ort_run_float_outputs, for a model with several outputs.
- first.jai copies the model to bin/segment-models. The model itself is not in git
(modules/Segment/models, 462 MB); tools/export_mask2former.py makes it.
76 lines
4.0 KiB
Markdown
76 lines
4.0 KiB
Markdown
# Segment
|
|
|
|
What is where in an image: panoptic segmentation with Mask2Former (Swin-B, COCO panoptic: 80 kinds of
|
|
thing and 53 of stuff). Every pixel gets a class ("person", "dog", "sky", "wall"), and countable things
|
|
are split into instances. It runs the model through onnxruntime ([`modules/onnxruntime`](../onnxruntime)),
|
|
sharing its session and CUDA code with [`modules/OCR`](../OCR) and [`modules/Labels`](../Labels). The
|
|
plan for using it in mcap is in [`segmentation-plan.html`](../../segmentation-plan.html).
|
|
|
|
## Setting up
|
|
|
|
There is no ONNX release of the model, so it is made from the Hugging Face checkpoint, once:
|
|
|
|
```
|
|
conda create -n mcap-seg python=3.11
|
|
conda run -n mcap-seg pip install torch --index-url https://download.pytorch.org/whl/cpu
|
|
conda run -n mcap-seg pip install transformers onnx onnxruntime pillow scipy timm
|
|
conda run -n mcap-seg python tools/export_mask2former.py --check FOLDER_OF_IMAGES
|
|
```
|
|
|
|
This writes `models/mask2former-swin-base-coco-panoptic.onnx` (462 MB) and its class list, then checks
|
|
the .onnx against PyTorch on the images. Building the app copies `models.config` and the models into
|
|
`bin/segment-models`. The models folder is not in git.
|
|
|
|
## Using it
|
|
|
|
```jai
|
|
#import "Segment";
|
|
|
|
engine, ok, error := segment_create(models_folder, "mask2former-swin-base-coco-panoptic"); // prefer_gpu = true
|
|
results := segment_images(engine, pictures); // [] Segment_Image (RGBA): one Segment_Result each
|
|
for results[0].segments print("% % %\n", engine.classes[it.class].name, it.score, it.pixels);
|
|
segment_destroy(engine);
|
|
```
|
|
|
|
- **Images** are squashed to 384 x 384 (not keeping their shape), with PIL's bilinear filter, each pass
|
|
rounded to 8 bits, then normalized by ImageNet's mean and standard deviation, as the checkpoint's own
|
|
processor does. Batches of `batch_gpu` (4) on CUDA, `batch_cpu` (2) on the CPU.
|
|
- **Segments** are made from the model's 100 queries as transformers' `post_process_panoptic_segmentation`
|
|
makes them (version 5.19 defaults; stuff of one class is not merged):
|
|
1. Queries whose best class is "no object", or scores 0.5 or less, are dropped.
|
|
2. Each mask is scaled up to 384 x 384 through its sigmoid, then to the map, and weighted by its
|
|
query's score.
|
|
3. Each pixel goes to the query weighing most there.
|
|
4. A query keeps its segment only when that is more than 80% of its own mask, which is its pixels
|
|
weighing 0.5 or more.
|
|
- **The map** (`Segment_Result.map`) is a segment id per pixel (0 for none), at the image's shape with
|
|
its longest side at most 1024. Each `Segment` has its class, score, area and bounds in the map.
|
|
- **Classes**: `engine.classes` has a name for people (`"wall"`) and the model's own
|
|
(`"wall-other-merged"`), and whether it is a countable thing. `engine.classes_hash` identifies the
|
|
list, so results from another list can be found out of date.
|
|
- **Cancelling**: set `engine.cancel` to a flag; when it is true between batches, every image left fails
|
|
as `SEGMENT_CANCELLED`.
|
|
|
|
## Checking it
|
|
|
|
`bin/segment-probe` (`jai first.jai - segment_probe`) runs it on image files and prints their segments.
|
|
`-out FOLDER` draws them over the images. `-check FOLDER` compares them with PyTorch's own segmentation,
|
|
made by `tools/export_mask2former.py --check IMAGES --reference FOLDER`.
|
|
|
|
On 25 images (16 COCO val2017 photos, 5 photos and 4 video frames), 2026-10-10:
|
|
|
|
- **Agreement:** 99.8% of pixels have the same class as PyTorch's, and every segment PyTorch scores 0.6
|
|
or more is found.
|
|
- **Differences:** two people at 0.56 and 0.58 are not found, and two segments are found that PyTorch
|
|
does not make. These sit at a threshold, tipped by mcap's input being a step or two (out of 255) off
|
|
PyTorch's, from stb_image decoding JPEGs a little differently from libjpeg.
|
|
- **Time per image, release build, RTX 5080 laptop:**
|
|
|
|
| Batch | Preparing | Model | Making segments | Peak GPU memory |
|
|
| --- | --- | --- | --- | --- |
|
|
| 1 | 3.7 ms | 38 ms | 31 ms | 1.4 GB |
|
|
| 4 | 3.2 ms | 21 ms | 30 ms | 2.4 GB |
|
|
| 8 | 4.0 ms | 22 ms | 31 ms | 2.4 GB |
|
|
|
|
On the CPU it is about 670 ms per image.
|