79 lines
4.3 KiB
Markdown
79 lines
4.3 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. On macOS, `tools/export_mask2former_coreml.py --check
|
|
FOLDER_OF_IMAGES` then makes `models/mask2former-swin-base-coco-panoptic.mlpackage`, which runs on the
|
|
GPU through Core ML (`modules/Core_ML`), one image a run, about four times as fast as the CPU with the
|
|
same segments. 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.
|