mcap/modules/Segment
2026-10-17 21:17:00 -04:00
..
models.config [0.1.32] macOS phase 5: Mask2Former segmentation on the GPU through native Core ML 2026-10-17 21:17:00 -04:00
module.jai [0.1.32] macOS phase 5: Mask2Former segmentation on the GPU through native Core ML 2026-10-17 21:17:00 -04:00
README.md [0.1.32] macOS phase 5: Mask2Former segmentation on the GPU through native Core ML 2026-10-17 21:17:00 -04:00

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), sharing its session and CUDA code with modules/OCR and modules/Labels. The plan for using it in mcap is in 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

#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.