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. |
||
|---|---|---|
| .. | ||
| models.config | ||
| module.jai | ||
| README.md | ||
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_segmentationmakes them (version 5.19 defaults; stuff of one class is not merged):- Queries whose best class is "no object", or scores 0.5 or less, are dropped.
- Each mask is scaled up to 384 x 384 through its sigmoid, then to the map, and weighted by its query's score.
- Each pixel goes to the query weighing most there.
- 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. EachSegmenthas its class, score, area and bounds in the map. - Classes:
engine.classeshas a name for people ("wall") and the model's own ("wall-other-merged"), and whether it is a countable thing.engine.classes_hashidentifies the list, so results from another list can be found out of date. - Cancelling: set
engine.cancelto a flag; when it is true between batches, every image left fails asSEGMENT_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.