mcap/modules/Segment
Musa Mahmood 8c8b2c6e38 [0.1.24] GPU memory: OCR runs largest first, gives its arenas' blocks back after each stage, and keeps its worker threads
OCR filled the GPU: 9 GB for the mobile model on gpu-bench's 30 images, past
16 GB for the server model, which Windows then paged into shared memory (1.5 to
11 s a capture). onnxruntime's logged arena extensions showed why: the arena
keeps every block it takes from CUDA and never joins or frees them, and OCR's
inputs grow from run to run (recognizer batches went narrowest first; images in
any order across calls), so each larger run left another block behind: 5.7 GB
of recognizer blocks where the widest batch needs 0.9 GB.

- Detection runs the largest images first, recognition the widest batches.
- The detector's and recognizer's arenas shrink once each stage of an
  ocr_process call is done, by a run on a small input with
  memory.enable_memory_arena_shrinkage. Shrinking after every run cost 33 to
  59% of the speed.
- The worker threads are made with the engine and kept, not made for every
  stage of every call (onnxruntime keeps GPU state for each thread that runs a
  session, microsoft/onnxruntime#33037); 6% faster.
- OCR's arenas grow by what is asked (kSameAsRequested).
- ort.log_level sets what the engines' onnxruntime environments log; INFO shows
  every arena extension.

Mobile OCR on gpu-bench (RTX 5080 laptop, 4 workers): peak 5.7 GB over the
baseline, 2 GB between calls, 142 ms a capture, against 9.0 GB and 126 ms. The
text read differs by a line in 2,986 (batches are grouped differently). The
server model is not yet measured again.
2026-10-17 21:12:57 -04:00
..
models.config [0.0.11] Segmentation phases 1-2: Mask2Former exported to ONNX, and modules/Segment to run it 2026-10-10 03:39:28 +01:00
module.jai [0.1.24] GPU memory: OCR runs largest first, gives its arenas' blocks back after each stage, and keeps its worker threads 2026-10-17 21:12:57 -04:00
README.md [0.0.11] Segmentation phases 1-2: Mask2Former exported to ONNX, and modules/Segment to run it 2026-10-10 03:39:28 +01: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. 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.