From cb21b7346043f559a4c3fb0379b0049de0b07e73 Mon Sep 17 00:00:00 2001 From: Musa Mahmood Date: Sat, 10 Oct 2026 03:39:28 +0100 Subject: [PATCH] [0.0.11] Segmentation phases 1-2: Mask2Former exported to ONNX, and modules/Segment to run it - 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. --- .gitignore | 2 + check_targets.jai | 2 +- first.jai | 13 +- modules/Segment/README.md | 75 +++++ modules/Segment/models.config | 18 + modules/Segment/module.jai | 573 ++++++++++++++++++++++++++++++++ modules/onnxruntime/session.jai | 36 +- segmentation-plan.html | 407 +++++++++++++++++++++++ tools/export_mask2former.py | 204 ++++++++++++ tools/segment_probe.jai | 236 +++++++++++++ 10 files changed, 1563 insertions(+), 3 deletions(-) create mode 100644 modules/Segment/README.md create mode 100644 modules/Segment/models.config create mode 100644 modules/Segment/module.jai create mode 100644 segmentation-plan.html create mode 100644 tools/export_mask2former.py create mode 100644 tools/segment_probe.jai diff --git a/.gitignore b/.gitignore index c6b80c5..11d64cf 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,5 @@ Thumbs.db /modules/OCR/models/ # Labels: downloaded, not committed (modules/Labels/fetch_models.jai) /modules/Labels/models/ +# Segmentation: exported by tools/export_mask2former.py, not committed +/modules/Segment/models/ diff --git a/check_targets.jai b/check_targets.jai index 43a6dee..6844b9f 100644 --- a/check_targets.jai +++ b/check_targets.jai @@ -29,7 +29,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; GIT_BRANCH :: "check"; GIT_HASH :: "-"; -VERSION :: "0.0.10"; +VERSION :: "0.0.11"; BUILD_DATE :: "-"; DEBUG :: true; COMPILE_TESTS :: true; diff --git a/first.jai b/first.jai index 181a02f..e237da5 100644 --- a/first.jai +++ b/first.jai @@ -8,7 +8,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; APP_BUNDLE_IDENTIFIER :: "com.example.mcap"; -VERSION :: "0.0.10"; +VERSION :: "0.0.11"; // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated // builds never write over the ones you run by hand. build_directory := "bin"; @@ -65,6 +65,8 @@ build_project :: () { if array_find(args, "ocr_bench") array_add(*pending, start_tool_build("ocr-bench", "tools/ocr_bench.jai", optimized = build_release)); // jai first.jai - image_probe: also bin/image-probe, what modules/Image_Decode reads from image files. if array_find(args, "image_probe") array_add(*pending, start_tool_build("image-probe", "tools/image_probe.jai", optimized = build_release)); + // jai first.jai - segment_probe: also bin/segment-probe, what modules/Segment makes of image files. + if array_find(args, "segment_probe") array_add(*pending, start_tool_build("segment-probe", "tools/segment_probe.jai", optimized = build_release)); wait_for_builds(pending); @@ -299,6 +301,15 @@ copy_assets :: () { } else { log("Labels: the models have not been fetched; run jai fetch_models.jai in modules/Labels for labels."); } + + // Segmentation: its model in bin/segment-models, the same way (made by tools/export_mask2former.py). + make_directory_if_it_does_not_exist(tprint("%/segment-models", build_directory)); + copy_if_changed("modules/Segment/models.config", tprint("%/segment-models/models.config", build_directory)); + if file_exists("modules/Segment/models") { + for file_list("modules/Segment/models") copy_if_changed(it, tprint("%/segment-models/%", build_directory, path_filename(it))); + } else { + log("Segmentation: the model has not been made; see modules/Segment/README.md."); + } } copy_if_changed :: (from: string, to: string) { diff --git a/modules/Segment/README.md b/modules/Segment/README.md new file mode 100644 index 0000000..a199c00 --- /dev/null +++ b/modules/Segment/README.md @@ -0,0 +1,75 @@ +# 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. diff --git a/modules/Segment/models.config b/modules/Segment/models.config new file mode 100644 index 0000000..4da4183 --- /dev/null +++ b/modules/Segment/models.config @@ -0,0 +1,18 @@ +@version 1 + +# Panoptic segmentation models (modules/Segment). Each file named here is in the models folder: made +# by tools/export_mask2former.py, which also checks the .onnx against PyTorch. + +[model mask2former-swin-base-coco-panoptic] +title: "Mask2Former Swin-B, COCO panoptic (133 classes)" +model: "mask2former-swin-base-coco-panoptic.onnx" +classes: "mask2former-swin-base-coco-panoptic.classes.txt" +image_size: 384 # squashed to 384 x 384, not keeping its shape, as the checkpoint's processor does +mean: 0.485 0.456 0.406 +std: 0.229 0.224 0.225 +batch_gpu: 4 +batch_cpu: 2 # the fastest per image on the CPU (tools/export_mask2former.py's check) +score_threshold: 0.5 # transformers 5.19's post_process_panoptic_segmentation defaults, which the check compares with +mask_threshold: 0.5 +overlap_threshold: 0.8 +map_max_side: 1024 # the stored segment map's longest side diff --git a/modules/Segment/module.jai b/modules/Segment/module.jai new file mode 100644 index 0000000..78793e0 --- /dev/null +++ b/modules/Segment/module.jai @@ -0,0 +1,573 @@ +// Segment: what is where in an image, from a panoptic segmentation model (Mask2Former): every pixel +// given a class ("person", "dog", "sky", "wall"), and countable things split into instances. +// +// engine, ok, error := segment_create(models_folder, "mask2former-swin-base-coco-panoptic"); +// results := segment_images(engine, pictures); // one Segment_Result per picture +// for results[0].segments print("% %\n", engine.classes[it.class].name, it.score); +// +// The model is listed in models_folder/models.config, made by tools/export_mask2former.py, and run +// with onnxruntime (modules/onnxruntime): on an NVIDIA GPU through CUDA when it can, else on the CPU. +// Its two outputs, a class and a mask for each of its 100 queries, are made into segments here, as +// transformers' post_process_panoptic_segmentation does (tools/export_mask2former.py checks that +// they match). + +Segment_Provider :: enum { CPU; CUDA; } + +// Tightly packed or not: stride is the bytes from one row to the next. +Segment_Image :: struct { + pixels: *u8; // RGBA + width, height, stride: s32; +} + +Segment_Model_Info :: struct { + name: string; + title: string; + model: string; // the .onnx, in the models folder + classes: string; // the class list, in the models folder + image_size: s32 = 384; // the input's width and height: images are resized to it, not keeping their shape + mean := float.[0.485, 0.456, 0.406]; // ImageNet's, per channel, after scaling to 0..1 + std := float.[0.229, 0.224, 0.225]; + batch_gpu: s32 = 4; + batch_cpu: s32 = 2; + score_threshold: float = 0.5; // a query's class score must be above this (transformers' default) + mask_threshold: float = 0.5; // a pixel is in a query's own mask at this (score times mask) + overlap_threshold: float = 0.8; // a segment keeps more than this of its own mask, or is dropped + map_max_side: s32 = 1024; // the segment map: the image scaled so its longest side is at most this +} + +Segment_Class :: struct { + name: string; // as people see and search it: "wall" + model_name: string;// the model's own: "wall-other-merged" + thing: bool; // countable (person, dog), not stuff (sky, wall) +} + +Segment_Engine :: struct { + model: Segment_Model_Info; + classes: [] Segment_Class; + classes_hash: u64; // of the class list, so results from another list are out of date + provider: Segment_Provider; + gpu_explanation: string; // why the GPU is not in use; "" when it is + env: *OrtEnv; + memory_info: *OrtMemoryInfo; + session: Ort_Session; + cancel: *bool; // when set, and true, segment_images stops between batches: the rest fail as "cancelled" + stats: Segment_Stats; // since it was made +} + +// Time spent in each step of segment_images, summed. +Segment_Stats :: struct { + images: s64; + prepare_us: s64; // resizing and normalizing + model_us: s64; // the model's run, on the GPU or the CPU + post_us: s64; // making segments of its outputs +} + +// One segment: a thing, or a stretch of stuff. Its pixels are those of map with its id. +Segment :: struct { + id: u8; // 1 up, in the order made + class: s32; // into engine.classes + score: float; // the model's confidence in its class, 0 to 1 + pixels: s64; // its area in the map + x0, y0, x1, y1: s32; // its bounds in the map; x1 and y1 are past its last pixel +} + +Segment_Result :: struct { + failed: bool; + error: string; + width, height: s32; // the map's: the image's shape, longest side at most map_max_side + map: [] u8; // a segment id per pixel, row by row; 0: none + segments: [] Segment; +} + +SEGMENT_MODELS_CONFIG :: "models.config"; +SEGMENT_CANCELLED :: "cancelled"; + +// The models models.config lists, in its order; from the context allocator. +segment_list_models :: (models_folder: string) -> [] Segment_Model_Info, error: string { + models: [..] Segment_Model_Info; + path := tprint("%/%", models_folder, SEGMENT_MODELS_CONFIG); + text, ok := read_entire_file(path, log_errors = false,, temp); + if !ok return models, tprint("Cannot read %.", path); + + model: *Segment_Model_Info; + version_seen := false; + line_number := 0; + for raw: split(text, "\n",, temp) { + line_number += 1; + line := trim(segment_strip_comment(raw)); + if !line.count continue; + if begins_with(line, "@version") { version_seen = true; continue; } + if !version_seen return models, tprint("% does not start with @version 1.", path); + if begins_with(line, "[model ") && ends_with(line, "]") { + model = array_add(*models); + model.* = .{}; + model.name = copy_string(trim(slice(line, 7, line.count - 8))); + continue; + } + colon := find_index_from_left(line, #char ":"); + if colon < 0 || !model { log_error("[Segment] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; } + key := trim(slice(line, 0, colon)); + value := trim(advance(line, colon + 1)); + if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\"" value = slice(value, 1, value.count - 2); + number, number_ok := string_to_float(value); + if key == { + case "title"; model.title = copy_string(value); + case "model"; model.model = copy_string(value); + case "classes"; model.classes = copy_string(value); + case "image_size"; if number_ok model.image_size = cast(s32) number; + case "mean"; segment_parse_three(value, *model.mean, path, line_number); + case "std"; segment_parse_three(value, *model.std, path, line_number); + case "batch_gpu"; if number_ok model.batch_gpu = cast(s32) number; + case "batch_cpu"; if number_ok model.batch_cpu = cast(s32) number; + case "score_threshold"; if number_ok model.score_threshold = number; + case "mask_threshold"; if number_ok model.mask_threshold = number; + case "overlap_threshold"; if number_ok model.overlap_threshold = number; + case "map_max_side"; if number_ok model.map_max_side = cast(s32) number; + case; log_error("[Segment] %:%: unknown key %. Ignored.", path, line_number, key); + } + } + return models, ""; +} + +// An engine for the model named in models_folder/models.config: on CUDA with prefer_gpu when every +// library that needs is present, else on the CPU. +segment_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Segment_Engine, ok: bool, error: string { + push_allocator(context.default_allocator); + loaded, load_error := ort_load(); + if !loaded return null, false, load_error; + + models, models_error := segment_list_models(models_folder); + if models_error return null, false, models_error; + model: *Segment_Model_Info; + for * models if it.name == model_name { model = it; break; } + if !model return null, false, tprint("% has no [model %].", SEGMENT_MODELS_CONFIG, model_name); + + engine := New(Segment_Engine); + engine.model = model.*; + ok, error := segment_load_classes(engine, tprint("%/%", models_folder, model.classes)); + if !ok { segment_destroy(engine); return null, false, error; } + + api := ort.api; + ok, error = ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "segment", *engine.env)); + if !ok { segment_destroy(engine); return null, false, error; } + ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info)); + + HINT :: "make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild"; + path := tprint("%/%", models_folder, model.model); + use_gpu := false; + if prefer_gpu { + available, explanation := ort_cuda_check("Segmentation"); + use_gpu = available; + if !available engine.gpu_explanation = copy_string(explanation); + } else { + engine.gpu_explanation = copy_string("Segmentation runs on the CPU, as it was asked to."); + } + if use_gpu { + ort_gpu_lock(); // one CUDA engine made or released at a time (modules/onnxruntime/session.jai) + engine.session, ok, error = ort_create_session(engine.env, path, true, 2, missing_hint = HINT); + if ok ok, error = segment_warm_up(engine); + if !ok ort_release_session(*engine.session); + ort_gpu_unlock(); + if ok { + engine.provider = .CUDA; + } else { + engine.gpu_explanation = sprint("Segmentation runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", error); + use_gpu = false; + } + } + if !use_gpu { + cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1); + engine.session, ok, error = ort_create_session(engine.env, path, false, cores, missing_hint = HINT); + if !ok { segment_destroy(engine); return null, false, error; } + engine.provider = .CPU; + } + return engine, true, ""; +} + +segment_destroy :: (engine: *Segment_Engine) { + if !engine return; + push_allocator(context.default_allocator); + if engine.provider == .CUDA ort_gpu_lock(); + ort_release_session(*engine.session); + if engine.provider == .CUDA ort_gpu_unlock(); + if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info); + if engine.env ort.api.ReleaseEnv(engine.env); + for engine.classes { free(it.name); free(it.model_name); } + array_free(engine.classes); + free(engine.gpu_explanation); + free(engine); +} + +// The map size for an image of width x height: its shape, the longest side at most max_side. +segment_map_size :: (width: s32, height: s32, max_side: s32) -> s32, s32 { + scale := min(cast(float64) 1, cast(float64) max_side / max(width, height)); + return max(cast(s32) 1, cast(s32) (width * scale + 0.5)), max(cast(s32) 1, cast(s32) (height * scale + 0.5)); +} + +// A segmentation of each image, in batches; from allocator. An image that fails says why in its +// result; with engine.cancel set and true between batches, every image left fails as SEGMENT_CANCELLED. +segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator := context.allocator) -> [] Segment_Result { + results := NewArray(images.count, Segment_Result,, allocator); + size := engine.model.image_size; + plane := cast(s64) size * size; + batch := cast(s64) max(ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu, 1); + input := NewArray(plane * 3 * batch, float,, context.default_allocator); + defer array_free(input,, context.default_allocator); + + first := 0; + while first < images.count { + count := min(batch, images.count - first); + if engine.cancel && engine.cancel.* { + for first..images.count - 1 results[it] = .{failed = true, error = copy_string(SEGMENT_CANCELLED,, allocator)}; + break; + } + start := current_time_monotonic(); + for 0..count - 1 segment_prepare_image(engine.model, images[first + it], input.data + plane * 3 * it); + shape := s64.[count, 3, size, size]; + outs: [2] Ort_Output; + ran := current_time_monotonic(); + engine.stats.prepare_us += to_microseconds(ran - start); + ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, array_view(input, 0, plane * 3 * count), shape, .["class_queries_logits", "masks_queries_logits"], outs); + posted := current_time_monotonic(); + engine.stats.model_us += to_microseconds(posted - ran); + engine.stats.images += count; + defer engine.stats.post_us += to_microseconds(current_time_monotonic() - posted); + if !ok { + for first..first + count - 1 results[it] = .{failed = true, error = copy_string(error,, allocator)}; + } else { + classes, masks := outs[0], outs[1]; + queries := classes.dims[1]; + labels := classes.dims[2]; // the classes and "no object", last + mask_h, mask_w := masks.dims[2], masks.dims[3]; + if classes.dim_count != 3 || masks.dim_count != 4 || classes.dims[0] != count || masks.dims[0] != count || masks.dims[1] != queries || labels != engine.classes.count + 1 { + message := tprint("The model gave outputs of shape % and %, not [%, Q, %] and [%, Q, H, W].", array_view(classes.dims, 0, classes.dim_count), array_view(masks.dims, 0, masks.dim_count), count, engine.classes.count + 1, count); + for first..first + count - 1 results[it] = .{failed = true, error = copy_string(message,, allocator)}; + } else { + for n: 0..count - 1 { + image := images[first + n]; + results[first + n] = segment_panoptic(engine, classes.data + n * queries * labels, masks.data + n * queries * mask_h * mask_w, + cast(s32) queries, cast(s32) labels, cast(s32) mask_w, cast(s32) mask_h, image.width, image.height, allocator); + } + } + ort_release_output(*outs[0]); + ort_release_output(*outs[1]); + } + first += count; + } + return results; +} + +segment_result_free :: (result: *Segment_Result, allocator := context.allocator) { + free(result.error,, allocator); + array_free(result.map,, allocator); + array_free(result.segments,, allocator); + result.* = .{}; +} + +// A class's colour, the same everywhere: golden-angle hues, so neighbouring class numbers differ. +segment_colour :: (class: s32) -> [3] u8 { + degrees := cast(float64) class * 137.508; + hue := cast(float) (degrees - floor(degrees / 360) * 360) / 60.0; + sector := cast(s32) hue; + f := hue - sector; + v, p, q, t := 230.0, 60.0, 230.0 - 170.0 * f, 60.0 + 170.0 * f; + rgb: [3] float; + if sector == { + case 0; rgb = .[v, t, p]; + case 1; rgb = .[q, v, p]; + case 2; rgb = .[p, v, t]; + case 3; rgb = .[p, q, v]; + case 4; rgb = .[t, p, v]; + case; rgb = .[v, p, q]; + } + return .[cast(u8) rgb[0], cast(u8) rgb[1], cast(u8) rgb[2]]; +} + +// The model's input for an image (3 x image_size x image_size floats at dest), as segment_images makes +// it: for checking against the checkpoint's own processor (segment-probe -dump-input). +segment_prepare :: (engine: *Segment_Engine, image: Segment_Image, dest: *float) { + segment_prepare_image(engine.model, image, dest); +} + +#scope_module + +// Transformers' post_process_panoptic_segmentation, for one image (no stuff fused): queries whose +// best class is not "no object" and scores above score_threshold are kept; their masks are scaled +// up to image_size, then to the map, through their sigmoid, and weighted by their scores; each +// pixel goes to the query weighing most there; and a query keeps its segment only when that is more +// than overlap_threshold of its own mask (its pixels weighing at least mask_threshold). +segment_panoptic :: (engine: *Segment_Engine, class_logits: *float, mask_logits: *float, queries: s32, labels: s32, mask_w: s32, mask_h: s32, + image_w: s32, image_h: s32, allocator: Allocator) -> Segment_Result { + model := engine.model; + result: Segment_Result; + result.width, result.height = segment_map_size(image_w, image_h, model.map_max_side); + map_w, map_h := result.width, result.height; + result.map = NewArray(cast(s64) map_w * map_h, u8, initialized = true,, allocator); + + // The queries kept: softmax over the classes, then the best. + Kept :: struct { query: s32; class: s32; score: float; } + kept: [..] Kept; + kept.allocator = temp; + for q: 0..queries - 1 { + row := class_logits + cast(s64) q * labels; + top := row[0]; + for 1..labels - 1 top = max(top, row[it]); + total: float64; + for 0..labels - 1 total += exp(cast(float64) (row[it] - top)); + best: s64; + best_score: float64 = -1; + for 0..labels - 1 { + p := exp(cast(float64) (row[it] - top)) / total; + if p > best_score { best, best_score = it, p; } // the first of equals, as torch's max + } + if best == labels - 1 || best_score <= model.score_threshold continue; + array_add(*kept, .{q, cast(s32) best, cast(float) best_score}); + } + if !kept.count || kept.count > 255 return result; + + // Each kept query's mask at image_size: its logits scaled up (bilinear, corners not aligned), then + // through the sigmoid. + size := model.image_size; + plane := cast(s64) size * size; + probs := NewArray(kept.count * plane, float, initialized = false,, temp); + across_in := segment_linear_taps(mask_w, size,, temp); + down_in := segment_linear_taps(mask_h, size,, temp); + for k, n: kept { + source := mask_logits + cast(s64) k.query * mask_w * mask_h; + dest := probs.data + n * plane; + for y: 0..size - 1 { + ty := down_in[y]; + row0 := source + cast(s64) ty.i0 * mask_w; + row1 := source + cast(s64) ty.i1 * mask_w; + for x: 0..size - 1 { + tx := across_in[x]; + top := row0[tx.i0] * (1 - tx.t) + row0[tx.i1] * tx.t; + bottom := row1[tx.i0] * (1 - tx.t) + row1[tx.i1] * tx.t; + logit := top * (1 - ty.t) + bottom * ty.t; + dest[cast(s64) y * size + x] = 1.0 / (1.0 + exp(-logit)); + } + } + } + + // Then at the map's size, weighted by the scores: each pixel to the query weighing most, and each + // query's own mask counted. + across := segment_linear_taps(size, map_w,, temp); + down := segment_linear_taps(size, map_h,, temp); + best_of := NewArray(cast(s64) map_w * map_h, u8, initialized = false,, temp); // index into kept + won := NewArray(kept.count, s64, initialized = true,, temp); + own := NewArray(kept.count, s64, initialized = true,, temp); + for y: 0..map_h - 1 { + ty := down[y]; + for x: 0..map_w - 1 { + tx := across[x]; + best, best_value := 0, -1.0; + for k, n: kept { + p := probs.data + n * plane; + row0 := p + cast(s64) ty.i0 * size; + row1 := p + cast(s64) ty.i1 * size; + top := row0[tx.i0] * (1 - tx.t) + row0[tx.i1] * tx.t; + bottom := row1[tx.i0] * (1 - tx.t) + row1[tx.i1] * tx.t; + value := (top * (1 - ty.t) + bottom * ty.t) * k.score; + if value >= model.mask_threshold own[n] += 1; + if value > best_value { best, best_value = n, value; } // the first of equals, as torch's argmax + } + best_of[cast(s64) y * map_w + x] = cast(u8) best; + won[best] += 1; + } + } + + // The segments: the queries, in order, that keep enough of their own mask. + id_of := NewArray(kept.count, u8, initialized = true,, temp); // 0: dropped + segments: [..] Segment; + segments.allocator = allocator; + for k, n: kept { + if !won[n] || !own[n] || !(cast(float64) won[n] / own[n] > model.overlap_threshold) continue; + id := cast(u8) (segments.count + 1); + id_of[n] = id; + array_add(*segments, .{id = id, class = k.class, score = k.score, x0 = map_w, y0 = map_h}); + } + for y: 0..map_h - 1 for x: 0..map_w - 1 { + at := cast(s64) y * map_w + x; + id := id_of[best_of[at]]; + result.map[at] = id; + if !id continue; + s := *segments[id - 1]; + s.pixels += 1; + s.x0 = min(s.x0, cast(s32) x); s.x1 = max(s.x1, cast(s32) x + 1); + s.y0 = min(s.y0, cast(s32) y); s.y1 = max(s.y1, cast(s32) y + 1); + } + result.segments = segments; + return result; +} + +// For each of out positions, the two in positions and the weight of the second, for torch's bilinear +// interpolate with align_corners = False: no averaging when shrinking. Temp. +Linear_Tap :: struct { i0, i1: s32; t: float; } + +segment_linear_taps :: (in: s32, out: s32) -> [] Linear_Tap { + taps := NewArray(out, Linear_Tap); + scale := cast(float) in / out; + for i: 0..out - 1 { + source := max((i + 0.5) * scale - 0.5, 0.0); + i0 := min(cast(s32) source, in - 1); + i1 := min(i0 + 1, in - 1); + taps[i] = .{i0, i1, source - i0}; + } + return taps; +} + +// The image resized to image_size x image_size (not keeping its shape, as the checkpoint's processor +// does) with PIL's bilinear filter, each pass rounded to 8 bits as PIL's is; then as planar RGB, +// scaled to 0..1 and normalized by the model's mean and std, at dest. +segment_prepare_image :: (model: Segment_Model_Info, image: Segment_Image, dest: *float) { + size := model.image_size; + across := segment_resample_weights(image.width, size); + down := segment_resample_weights(image.height, size); + rows := NewArray(cast(s64) image.height * size * 3, u8, initialized = false,, temp); + for y: 0..image.height - 1 { + row := image.pixels + cast(s64) y * image.stride; + for x: 0..size - 1 { + w := *across[x]; + r, g, b: float; + for k: 0..w.weights.count - 1 { + p := row + (w.first + k) * 4; + r += p[0] * w.weights[k]; g += p[1] * w.weights[k]; b += p[2] * w.weights[k]; + } + at := (cast(s64) y * size + x) * 3; + rows[at] = segment_byte(r); rows[at + 1] = segment_byte(g); rows[at + 2] = segment_byte(b); + } + } + plane := cast(s64) size * size; + scale: [3] float; + offset: [3] float; + for c: 0..2 { scale[c] = 1 / (255 * model.std[c]); offset[c] = model.mean[c] / model.std[c]; } + for y: 0..size - 1 { + w := *down[y]; + for x: 0..size - 1 { + r, g, b: float; + for k: 0..w.weights.count - 1 { + at := (cast(s64) (w.first + k) * size + x) * 3; + r += rows[at] * w.weights[k]; g += rows[at + 1] * w.weights[k]; b += rows[at + 2] * w.weights[k]; + } + i := cast(s64) y * size + x; + dest[i] = segment_byte(r) * scale[0] - offset[0]; + dest[i + plane] = segment_byte(g) * scale[1] - offset[1]; + dest[i + plane * 2] = segment_byte(b) * scale[2] - offset[2]; + } + } +} + +segment_byte :: (v: float) -> u8 { + return cast(u8) clamp(cast(s32) (v + 0.5), 0, 255); +} + +Resample_Weights :: struct { first: s32; weights: [] float; } + +// For each of out positions, the input positions it takes and their weights: a triangle filter +// widened by the scale when shrinking, as PIL's bilinear resize (and modules/Labels). Temp. +segment_resample_weights :: (in: s32, out: s32) -> [] Resample_Weights { + result := NewArray(out, Resample_Weights,, temp); + scale := cast(float) in / out; + filter_scale := max(scale, 1.0); + support := filter_scale; + for i: 0..out - 1 { + centre := (i + 0.5) * scale; + first := max(cast(s32) (centre - support + 0.5), 0); + last := min(cast(s32) (centre + support + 0.5), in); + weights := NewArray(max(last - first, 1), float,, temp); + total: float; + for k: 0..last - first - 1 { + x := (first + k - centre + 0.5) / filter_scale; + w := max(0.0, 1 - abs(x)); + weights[k] = w; + total += w; + } + if total > 0 for * weights it.* /= total; + else { first = clamp(cast(s32) centre, 0, in - 1); weights[0] = 1; } + result[i] = .{first, weights}; + } + return result; +} + +// Runs the model once, on a blank image, when it is made on the GPU: see labels_warm_up. +segment_warm_up :: (engine: *Segment_Engine) -> ok: bool, error: string { + size := engine.model.image_size; + input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator); + defer array_free(input,, context.default_allocator); + outs: [2] Ort_Output; + ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, s64.[1, 3, size, size], .["class_queries_logits", "masks_queries_logits"], outs); + if !ok return false, tprint("The first run on the GPU failed: %", error); + ort_release_output(*outs[0]); + ort_release_output(*outs[1]); + return true, ""; +} + +// The class list: `class: INDEX thing|stuff "model's name" "name"`, a line each, in index order. +segment_load_classes :: (engine: *Segment_Engine, path: string) -> ok: bool, error: string { + text, ok := read_entire_file(path, log_errors = false,, temp); + if !ok return false, tprint("Cannot read %; make it with tools/export_mask2former.py.", path); + classes: [..] Segment_Class; + for raw, line_index: split(text, "\n",, temp) { + line := trim(raw); + if !begins_with(line, "class:") continue; + rest := trim(advance(line, 6)); + space := find_index_from_left(rest, #char " "); + index, index_ok := string_to_int(slice(rest, 0, max(space, 0))); + if space < 0 || !index_ok || index != classes.count return false, tprint("%:%: expected class % next.", path, line_index + 1, classes.count); + rest = trim(advance(rest, space + 1)); + thing := begins_with(rest, "thing"); + if !thing && !begins_with(rest, "stuff") return false, tprint("%:%: expected thing or stuff.", path, line_index + 1); + quoted: [2] string; + for 0..1 { + open := find_index_from_left(rest, #char "\""); + if open < 0 return false, tprint("%:%: expected two quoted names.", path, line_index + 1); + close := find_index_from_left(rest, #char "\"", open + 1); + if close < 0 return false, tprint("%:%: a name is not closed.", path, line_index + 1); + quoted[it] = slice(rest, open + 1, close - open - 1); + rest = advance(rest, close + 1); + } + array_add(*classes, .{name = copy_string(quoted[1]), model_name = copy_string(quoted[0]), thing = thing}); + } + if !classes.count return false, tprint("% lists no classes.", path); + engine.classes = classes; + + // FNV-1a over each class's index, kind and name, as a line. + hash := FNV_64_OFFSET_BIAS; + for engine.classes { + line := tprint("%\t%\t%\n", it_index, ifx it.thing then "thing" else "stuff", it.model_name); + hash = fnv1a_hash(line.data, line.count, hash); + } + engine.classes_hash = hash; + return true, ""; +} + +segment_parse_three :: (value: string, dest: *[3] float, path: string, line_number: s64) { + parts := split(value, " ",, temp); + got := 0; + for parts { + if !it.count continue; + number, ok := string_to_float(it); + if !ok || got >= 3 { got = -1; break; } + dest.*[got] = number; + got += 1; + } + if got != 3 log_error("[Segment] %:%: expected three numbers.", path, line_number); +} + +segment_strip_comment :: (line: string) -> string { + in_quotes := false; + for i: 0..line.count - 1 { + c := line[i]; + if c == #char "\"" in_quotes = !in_quotes; + else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t") return slice(line, 0, i); + } + return line; +} + +#import "Basic"; +#import "String"; +#import "File"; +#import "Math"; +#import "System"; +#import "Hash"; +#import "onnxruntime"; diff --git a/modules/onnxruntime/session.jai b/modules/onnxruntime/session.jai index 045c213..3f5497e 100644 --- a/modules/onnxruntime/session.jai +++ b/modules/onnxruntime/session.jai @@ -1,5 +1,5 @@ // Sessions: one model file each, run on an NVIDIA GPU through CUDA or on the CPU, and running them on -// a tensor. Shared by the modules that run models (OCR, Labels). +// a tensor. Shared by the modules that run models (OCR, Labels, Segment). // CUDA sessions are made, first run and released one at a time, process-wide: hold ort_gpu_lock // around it. Done on two threads at once (onnxruntime 1.30.0, CUDA 13, cuBLASLt 13.8, RTX 5080 @@ -151,6 +151,40 @@ ort_run_float :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, return out, ok, error; } +// As ort_run_float, for a model with several outputs: those named, into outs (one each, in order), +// each released with ort_release_output. On failure none is set. +ort_run_float_outputs :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64, names: [] string, outs: [] Ort_Output) -> ok: bool, error: string { + assert(names.count == outs.count); + api := ort.api; + input: *OrtValue; + ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data.data, xx (data.count * size_of(float)), shape.data, xx shape.count, .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, *input)); + if !ok return false, error; + defer api.ReleaseValue(input); + + output_names := NewArray(names.count, *u8,, temp); + for names output_names[it_index] = temp_c_string(it); + values := NewArray(names.count, *OrtValue, initialized = true,, temp); + ok, error = ort_check(api.Run(s.session, null, *s.input_name, *input, 1, output_names.data, xx names.count, values.data)); + if !ok { + for values if it api.ReleaseValue(it); + return false, error; + } + for * out: outs { + out.* = .{value = values[it_index]}; + raw: *void; + ort_check(api.GetTensorMutableData(out.value, *raw)); + out.data = cast(*float) raw; + info: *OrtTensorTypeAndShapeInfo; + ort_check(api.GetTensorTypeAndShape(out.value, *info)); + count: u64; + ort_check(api.GetDimensionsCount(info, *count)); + out.dim_count = min(cast(s64) count, out.dims.count); + ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count)); + api.ReleaseTensorTypeAndShapeInfo(info); + } + return true, ""; +} + ort_release_output :: (out: *Ort_Output) { if out.value ort.api.ReleaseValue(out.value); out.* = .{}; diff --git a/segmentation-plan.html b/segmentation-plan.html new file mode 100644 index 0000000..7b94f83 --- /dev/null +++ b/segmentation-plan.html @@ -0,0 +1,407 @@ + + + + + +Segmentation Plan + + + +
+ +
+

Semantic segmentation in mcap

+

Planning document · 2026-10-10 · for mcap after 0.0.10

+
+ +
+ In short +
    +
  • Decided Mask2Former, exported to ONNX and run with the onnxruntime mcap already ships. It is swappable later, because results are keyed by model name.
    • +
  • Decided Images only, and only in the library folders you opt in. The setting is per folder, kept in a file in the library root, so it travels with the library.
    • +
  • Decided No .mcap version bump. Results go in two new chunk kinds: a small segs summary at the front for search, and a mask at the back for drawing.
    • +
  • It is built in six phases. The first proves the ONNX export matches PyTorch before anything else depends on it.
    • +
  • Progress Phases 1 and 2 are done (2026-10-10). The model is exported and checked, and modules/Segment runs it in mcap at about 57 ms per image on the RTX 5080. Next is phase 3: storing results in .mcap files.
    • +
  • Later Video, by treating sampled frames as images, is shelved.
    • +
+
+ + + +

1. Scope

+

Each opted-in image gets a panoptic segmentation. Every pixel gets a class ("person", "dog", "sky", "wall"), and countable things are split into instances: two people are two segments. The results serve three purposes:

+ +

What counts as an image. Imported photos and screenshots, Snagit captures and animated GIFs (their first frame) qualify. Videos are skipped even in opted-in folders: a video's poster frame is an image technically, but frames are the later project. A capture is a video when its catalog entry's media is VIDEO, which comes from VideoCodec in its import.json.

+ +

2. The model

+

Which Mask2Former

+
+ + + + + + +
Checkpoint (Hugging Face)ClassesSizeUse
facebook/mask2former-swin-base-coco-panoptic133 COCO panoptic: 80 things, 53 stuff~107MDefault. Things are counted as instances, and the stuff classes (sky, wall, grass, road) cover the rest.
facebook/mask2former-swin-large-coco-panopticsame~216MIf base misses too much and the GPU can take it.
facebook/mask2former-swin-tiny-coco-panopticsame~47MThe CPU fallback, when CUDA is unavailable.
facebook/mask2former-swin-base-ade-semantic150 ADE20K, no instances~107MThe alternative, if indoor scenes and objects matter more than counting people. See open questions.
+
+

COCO panoptic is the default because a photo library's questions are mostly about things ("photos with the dog", "group shots"), and only the panoptic models separate instances.

+ +

Export to ONNX

+

There is no official ONNX release, so phase 1 exports one with a Python script kept in the repo (tools/export_mask2former.py). It runs once on a developer machine and is never shipped.

+ +
Phase 1 result (2026-10-10): passed. +
    +
  • Test set: 25 images: 16 COCO val2017 photos, 5 of your photos and 4 video frames.
    • +
  • Accuracy: 180 segments, every one with an IoU of 1.000 against PyTorch, none missing and none extra.
    • +
  • Batching: batches of 2, 4 and 8 give bit-identical results to one at a time.
    • +
  • Speed: CPU 348 ms per image at batch 1, 290 ms at batch 2.
    • +
  • Size: 462 MB at opset 17.
    • +
  • Load warning: the "missing" swin.layernorm that transformers 5.19 reports when loading is unused; randomizing it changes no output.
    • +
  • Not yet measured: GPU timing and memory. Those come in phase 2, with mcap's own onnxruntime build.
    • +
+Script: tools/export_mask2former.py.
+
Risk. The export is the one part of this plan that may fail outright. If grid_sample export misbehaves, fall back to the ADE20K semantic SegFormer, which exports trivially, to build phases 2 to 6. Then return to Mask2Former, since everything after phase 1 is keyed by model name.
+ +

Packaging, as Labels and OCR are

+ + +

3. Per-folder opt-in

+

A library folder is a top-level folder of the library, one per place imported from (SnagitLibraryWork, Phone, Photos 2014). Segmentation runs only in the folders you turn it on for. A new folder starts with it off, so adding a 40,000-photo archive never starts an hours-long GPU job by surprise.

+ +

Where the setting lives

+

It goes in a new file in the library root, .library-folders, next to .import-index and .import-sources. Not in %APPDATA%\mcap\settings.config, because the choice belongs to the library, not to the machine. It follows the library when it is opened elsewhere, or with -library FOLDER. The format follows the Ex4 style used in the rest of mcap, so later per-folder options (OCR off for a photos folder, say) go in the same file:

+
@version 1
+
+[folder "Photos 2014"]
+segment: true
+
+[folder "SnagitLibraryWork"]
+segment: false
+ + +

Where you set it

+ + +

4. The segmentation scan

+
+ + + + .library-folders + segment: true + + Items + images not yet done + + Decode, segment + batches on the scan thread + + mcap_store + segs front, mask back + + + + + catalog_load → viewer.catalog_updates → search and Stats + +
Only the scan's choice of captures is new. Batching, pause and cancel, storing and the catalog hand-off are the OCR scan's.
+
+ + + +

5. Storage in .mcap

+
No file version bump. Version 3 added tagged, model-keyed chunk lists, and the format README already names segmentation masks as what back chunks are for. A reader that does not know a chunk keeps it on rewrite. Each new chunk carries its own internal version, as labl does.
+ +

segs, front chunk, keyed by model name

+

This is what search and Stats need, small enough for the front that every library browse reads. It is text in the style of labl:

+
@version 1
+engine:  "onnxruntime 1.30.0, CUDA"
+date:    2026-10-12T09:30:00Z
+image:   rendered 4032 3024
+classes: 5c1e9a0d77b2f413
+segment: 1 0.981 0.2140 1210 380 900 2210 "person"
+segment: 2 0.955 0.1022 2300 900 700 1400 "person"
+segment: 3 0.912 0.0811 400 1900 800 700 "dog"
+segment: 4 0.993 0.3105 0 0 4032 1100 "sky"
+
+ + + + + + +
KeyValue
engine, dateAs in OCR text.
imageWhich image was read (rendered or original) and its size. Boxes are in its pixels.
classesA 64-bit FNV-1a hash of the model's class list, so a model with new classes is seen as needing a new scan.
segmentThe segment's id in the mask, its score, its area (as a fraction of the image), its box (x y w h), and its class, quoted. Largest first. Stuff classes have one segment per class.
+
+

A busy photo has 10 to 30 segments, about 1 to 2 KB.

+ +

mask, back chunk, keyed by model name

+

This is the segment-id map, for drawing only. It is binary:

+
u8   format        1: run-length encoded, row-major
+u16  width         the map's size: the image scaled so its longest side is at most 1024
+u16  height
+u8   id_bits       8 (up to 255 segments), or 16 later if needed
+then runs:         (id, length) pairs over the whole map, row by row, wrapping across rows
+                   id as u8; length as a LEB128 varint
+                   id 0: no segment (below threshold, or "no object")
+ + +

Code changes

+ + + + +
has: and label: answer different questions. Labels (SigLIP 2) describe the whole image against your own list ("Settings", "Installer", "Beach"). Segments say which COCO things are in it, and where. Both stay, and they can be combined: label:beach has:dog.
+ +

7. Capture window overlay

+ + +

8. Phases

+

Each phase ends with something to run and check, and is a commit (0.0.11 onwards).

+ +
+

Phase 1: Export and check the model Done

+

tools/export_mask2former.py exports Swin-B COCO panoptic. It is checked against transformers on 20 or more photos, with the IoU target above, and the timings and peak GPU memory are recorded at batch sizes 1, 2 and 4 on your GPU.

+

Done when: an .onnx that loads in onnxruntime 1.30 with CUDA and passes the IoU check, plus a note of the numbers.

+
+ +
+

Phase 2: modules/Segment and the probe Done

+

The module (create, segment, panoptic merge, cancel), models.config, the copy step in first.jai, and tools/segment_probe.jai. fetch_models.jai waits on where the model will be hosted (open question 2); until then tools/export_mask2former.py makes it.

+

Done when: segment-probe on the phase 1 photos writes masks that match phase 1's, and runs on both CPU and CUDA.

+
Result (2026-10-10). +
    +
  • Agreement: 99.8% of pixels have the same class as PyTorch's, on both CUDA and the CPU. Every segment PyTorch scores 0.6 or more is found. The test set is the same 25 images.
    • +
  • Differences: two people at 0.56 and 0.58 are missed, and there are two extra segments. These are borderline cases, tipped by the input being a step or two (out of 255) off PyTorch's, because stb_image decodes JPEGs slightly differently from libjpeg.
    • +
  • Release build, RTX 5080: 57 ms per image at batch 4: preparing 3 ms, the model 21 ms, making segments 30 ms. Peak GPU memory is 2.4 GB. Batch 8 is no faster; batch 1 is 38 ms for the model. On the CPU, about 670 ms per image.
    • +
+

Things learned on the way:

+
    +
  • The score threshold. transformers 5.19's default is 0.5, not the 0.8 this plan assumed, and the check compares with that default. It's score_threshold in models.config, to raise if weak segments prove noisy in search.
    • +
  • Stuff isn't merged per class by default, so two patches of sky can be two segments.
    • +
  • Making segments now takes longer than the model. Spreading it across threads can roughly double throughput if the scan needs it.
    • +
  • Overlay colours. The overlay colours by class, so neighbouring people share a colour. Phase 6 should vary the shade per instance.
    • +
+
+ +
+

Phase 3: Chunks and storing

+

segment_chunks.jai (text, RLE), mcap_store with back chunks, the README section, and tests:

+
    +
  • Round trip.
    • +
  • RLE edge cases: one segment covering everything, a map with no segments, runs wrapping across rows, the largest id.
    • +
  • Storing segs and mask together leaves the original and every other chunk as they were.
    • +
+

Done when: the tests pass and a probe run can store its results into a scratch library's files.

+
+ +
+

Phase 4: Per-folder opt-in and the scan

+

The pieces:

+
    +
  • .library-folders (load, save, keep unknown keys).
    • +
  • The Segmentation window's folder table.
    • +
  • The Import window's checkbox.
    • +
  • The SEGMENT task and its item selection.
    • +
  • The import follow-up scan.
    • +
  • -segment and -erase segments.
    • +
  • The shutdown status line.
    • +
+

Done when: on a scratch library with two folders, only the opted-in one is segmented. Videos in it are skipped. Cancelling mid-batch stores nothing. Erase per folder works.

+
+ +
+

Phase 5: Catalog, has: and Stats

+

Done when: has:person finds the right captures in the scratch library, catalog tests cover the new field, and Stats shows the top classes, each clickable.

+
+ +
+

Phase 6: The overlay

+

The mask texture on a worker, the legend, hiding classes, opacity and hover.

+

Done when: masks line up with the image at every fit and zoom, including on Snagit captures (rendered image) and on photos turned by their EXIF orientation.

+
+ +

9. Open questions

+
    +
  1. COCO panoptic or ADE20K? COCO counts things (people, pets, cars); ADE20K knows more indoor objects and scene parts, but has no instances. The default is COCO. Revisit after phase 1 by trying both on a sample of your photo libraries.
    2. +
  2. Where to host the exported .onnx for fetch_models.jai: a private Hugging Face repo, a release asset, or committing it as the FFmpeg DLLs are. At about 400 MB for Swin-B, committing it is probably too much.
    3. +
  3. The mask resolution cap. 1024 on the long side is a guess. Phase 3 should measure chunk sizes on real photos at 512, 1024 and 2048.
    4. +
  4. Snagit and screenshot folders. They stay off by default, and COCO classes mean little on UI. Turning one on is allowed, but the Segmentation window could warn that most of the folder is screenshots, judged from Format and AppName.
    5. +
  5. Re-scanning on a model change. A new model name means a new key, so old results stay until erased. Should the scan offer to erase other models' masks to save space? Probably yes, as an explicit button rather than automatically.
    6. +
+ +

10. Later: video frames Shelved

+

Treat strategically sampled frames of a video as images, and run the same analysis (segmentation, and OCR and labels too) on them. Notes for when it comes back:

+ + +
+ + diff --git a/tools/export_mask2former.py b/tools/export_mask2former.py new file mode 100644 index 0000000..19f32d7 --- /dev/null +++ b/tools/export_mask2former.py @@ -0,0 +1,204 @@ +# Exports a Mask2Former checkpoint from Hugging Face to ONNX for modules/Segment, then checks the ONNX +# model against PyTorch on a folder of images: the raw logits, and the panoptic segmentation made from +# them, segment by segment. Run once, on a developer machine; mcap only ships the .onnx it writes. +# +# 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 [--reference OUT] +# +# --reference also writes PyTorch's own segmentation of each image, which bin/segment-probe -check +# compares modules/Segment's against. +# +# The model takes pixel_values [batch, 3, 384, 384] (RGB, resized to 384 x 384 without keeping the +# aspect ratio, as the checkpoint's processor does; scaled to 0..1, then ImageNet mean and std) and +# gives class_queries_logits [batch, 100, 134] (133 classes and "no object") and masks_queries_logits +# [batch, 100, 96, 96]. Making segments of those (the panoptic merge) is mcap's, in Jai. + +import argparse, json, os, sys, time +import numpy as np +import torch +from PIL import Image +from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation +from transformers.utils import logging + +logging.set_verbosity_error() + + +class Logits(torch.nn.Module): + """The model, from pixel values to the two logit tensors, as the ONNX graph has it.""" + + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, pixel_values): + out = self.model(pixel_values=pixel_values) + return out.class_queries_logits, out.masks_queries_logits + + +def export(model, path, size, opset): + wrapper = Logits(model).eval() + dummy = torch.randn(1, 3, size, size) + start = time.time() + torch.onnx.export( + wrapper, (dummy,), path, + input_names=["pixel_values"], + output_names=["class_queries_logits", "masks_queries_logits"], + dynamic_axes={"pixel_values": {0: "batch"}, "class_queries_logits": {0: "batch"}, "masks_queries_logits": {0: "batch"}}, + opset_version=opset, + do_constant_folding=True, + dynamo=False, + ) + print(f"exported {path} ({os.path.getsize(path) / 1e6:.0f} MB) in {time.time() - start:.0f} s") + + +class Outputs: + """What post_process_panoptic_segmentation reads, from ONNX's arrays.""" + + def __init__(self, class_logits, mask_logits): + self.class_queries_logits = torch.from_numpy(class_logits) + self.masks_queries_logits = torch.from_numpy(mask_logits) + + +def segments(processor, outputs, size): + """[(label, mask)] of the panoptic segmentation, at the image's size.""" + result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[size])[0] + seg = result["segmentation"] + return [(s["label_id"], (seg == s["id"]).numpy(), s["score"]) for s in result["segments_info"]] + + +def iou(a, b): + union = np.logical_or(a, b).sum() + return 1.0 if union == 0 else np.logical_and(a, b).sum() / union + + +def check(model, processor, path, folder, names, providers): + import onnxruntime as ort + session = ort.InferenceSession(path, providers=providers) + print(f"onnxruntime {ort.__version__} on {session.get_providers()[0]}") + files = sorted(f for f in os.listdir(folder) if f.lower().endswith((".jpg", ".jpeg", ".png"))) + worst_logit, ious, unmatched, torch_s, onnx_s = 0.0, [], 0, 0.0, 0.0 + for name in files: + image = Image.open(os.path.join(folder, name)).convert("RGB") + inputs = processor(images=image, return_tensors="pt") + start = time.time() + with torch.no_grad(): + reference = model(pixel_values=inputs["pixel_values"]) + torch_s += time.time() - start + start = time.time() + class_logits, mask_logits = session.run(None, {"pixel_values": inputs["pixel_values"].numpy()}) + onnx_s += time.time() - start + + diff = max(np.abs(class_logits - reference.class_queries_logits.numpy()).max(), + np.abs(mask_logits - reference.masks_queries_logits.numpy()).max()) + worst_logit = max(worst_logit, diff) + + size = image.size[::-1] + ref = segments(processor, reference, size) + got = segments(processor, Outputs(class_logits, mask_logits), size) + # Each reference segment against the best ONNX segment of the same class. + line = [] + for label, mask, score in ref: + best = max((iou(mask, m) for l, m, _ in got if l == label), default=0.0) + ious.append(best) + if best < 0.5: + unmatched += 1 + line.append(f"{names[label]} {best:.3f}") + extra = len(got) - len(ref) + print(f" {name}: logits max diff {diff:.1e}; {len(ref)} segments{f' ({extra:+d} in ONNX)' if extra else ''}: " + ", ".join(line[:6]) + (" ..." if len(line) > 6 else "")) + + ious = np.array(ious) + print(f"\n{len(files)} images, {len(ious)} segments") + print(f" logits: worst max abs diff {worst_logit:.2e}") + print(f" segment IoU, ONNX vs PyTorch: min {ious.min():.4f}, mean {ious.mean():.4f}, below 0.98: {(ious < 0.98).sum()}, unmatched: {unmatched}") + print(f" time per image: PyTorch {torch_s / len(files) * 1000:.0f} ms, ONNX {onnx_s / len(files) * 1000:.0f} ms (CPU, batch 1)") + return ious.min() >= 0.98 and unmatched == 0 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="facebook/mask2former-swin-base-coco-panoptic") + parser.add_argument("--out", default="modules/Segment/models") + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--check", help="a folder of images to check the export on") + parser.add_argument("--skip-export", action="store_true", help="check an .onnx already written") + parser.add_argument("--reference", help="write PyTorch's segmentation of the --check images here, for segment-probe -check") + args = parser.parse_args() + + processor = AutoImageProcessor.from_pretrained(args.model) + model = Mask2FormerForUniversalSegmentation.from_pretrained(args.model).eval() + size = processor.size["height"] + assert processor.size["width"] == size, "the export assumes a square input" + names = model.config.id2label + + os.makedirs(args.out, exist_ok=True) + base = args.model.split("/")[-1] + path = os.path.join(args.out, base + ".onnx") + if not args.skip_export: + export(model, path, size, args.opset) + + # The class list, for models.config: index, whether it is a countable thing, the model's name, and + # the name people see and search (has:): COCO's "wall-other-merged" is "wall". + with open(os.path.join(args.out, base + ".classes.txt"), "w", encoding="utf-8") as f: + f.write("@version 1\n") + for i in range(len(names)): + thing = "thing" if i < 80 else "stuff" # COCO panoptic: 80 things, then 53 stuff + f.write(f'class: {i} {thing} "{names[i]}" "{display_name(names[i])}"\n') + print(f"wrote {base}.classes.txt ({len(names)} classes)") + + if args.reference: + if not args.check: + sys.exit("--reference needs --check FOLDER, the images to make it of") + write_reference(model, processor, args.check, args.reference) + + if args.check: + ok = check(model, processor, path, args.check, names, ["CPUExecutionProvider"]) + print("PASS" if ok else "FAIL") + sys.exit(0 if ok else 1) + + +def display_name(name): + """COCO panoptic's class names for people: no -merged, -other or -stuff, spaces for dashes.""" + for suffix in ("-merged", "-other", "-stuff"): + name = name.replace(suffix, "") + return name.replace("-", " ").strip() + + +MASK_MAX_SIDE = 1024 # mcap's mask map: the image scaled so its longest side is at most this + + +def map_size(width, height): + """The mask map's size for an image, as modules/Segment works it out (rounding half up).""" + scale = min(1.0, MASK_MAX_SIDE / max(width, height)) + return max(1, int(width * scale + 0.5)), max(1, int(height * scale + 0.5)) + + +def write_reference(model, processor, folder, out): + """PyTorch's panoptic segmentation of each image, at the mask map's size, for segment-probe -check: + NAME.ref (b"M2FR", u16 width, u16 height, then a u8 segment id per pixel, row by row) and + NAME.ref.txt (a line per segment: id, class index, score).""" + import struct + from PIL import ImageOps + os.makedirs(out, exist_ok=True) + files = sorted(f for f in os.listdir(folder) if f.lower().endswith((".jpg", ".jpeg", ".png"))) + for name in files: + image = ImageOps.exif_transpose(Image.open(os.path.join(folder, name))).convert("RGB") # turned, as mcap shows it + w, h = map_size(*image.size) + inputs = processor(images=image, return_tensors="pt") + with torch.no_grad(): + outputs = model(pixel_values=inputs["pixel_values"]) + result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[(h, w)])[0] + seg = result["segmentation"].numpy().astype(np.int64) + seg[seg < 0] = 0 # no segment + assert seg.max() < 256 + with open(os.path.join(out, name + ".ref"), "wb") as f: + f.write(b"M2FR" + struct.pack("= 0.5, and ours left over + missed_confident: s64; // PyTorch's segments scoring CONFIDENT or more that ours do not match + worst: float64 = 1; // the lowest IoU of a PyTorch segment with its best match +} + +// Ours against PyTorch's: classes pixel by pixel, and each of its segments against our best of the same class. +compare :: (engine: *Segment_Engine, result: Segment_Result, ref_path: string, info_path: string) -> Comparison { + c: Comparison; + data, ok := read_entire_file(ref_path, log_errors = false,, temp); + if !ok { c.error = tprint("no reference % (tools/export_mask2former.py --reference)", ref_path); return c; } + if data.count < 8 || slice(data, 0, 4) != "M2FR" { c.error = tprint("% is not a reference", ref_path); return c; } + w := (cast(s32) data[4]) + ((cast(s32) data[5]) << 8); + h := (cast(s32) data[6]) + ((cast(s32) data[7]) << 8); + if w != result.width || h != result.height { c.error = tprint("the reference is % x %, ours % x %", w, h, result.width, result.height); return c; } + ref_map := cast([] u8) advance(data, 8); + if ref_map.count != cast(s64) w * h { c.error = "the reference is cut short"; return c; } + + // Its segments' classes, by id. + info, info_ok := read_entire_file(info_path, log_errors = false,, temp); + if !info_ok { c.error = tprint("no %", info_path); return c; } + ref_class: [256] s32; + for * ref_class it.* = -1; + ref_score: [256] float64; + ref_ids: [..] u8; + ref_ids.allocator = temp; + for line: split(info, "\n",, temp) { + parts := split(trim(line), " ",, temp); + if parts.count < 2 continue; + id := string_to_int(parts[0]); + ref_class[id] = cast(s32) string_to_int(parts[1]); + if parts.count >= 3 ref_score[id] = string_to_float64(parts[2]); + array_add(*ref_ids, cast(u8) id); + } + our_class: [256] s32; + for * our_class it.* = -1; + for result.segments our_class[it.id] = it.class; + + c.pixels = ref_map.count; + for 0..ref_map.count - 1 if ref_class[ref_map[it]] == our_class[result.map[it]] c.agree += 1; + + used: [256] bool; + for rid: ref_ids { + c.checked += 1; + best: float64; + best_id: u8; + for result.segments { + if it.class != ref_class[rid] continue; + both, either: s64; + for p: 0..ref_map.count - 1 { + a := ref_map[p] == rid; + b := result.map[p] == it.id; + if a && b both += 1; + if a || b either += 1; + } + iou := ifx either then cast(float64) both / either else 1.0; + if iou > best { best, best_id = iou, it.id; } + } + if best >= 0.5 { c.matched += 1; used[best_id] = true; } + else if ref_score[rid] >= CONFIDENT c.missed_confident += 1; + if best < 0.5 print(" check: PyTorch's % (score %) has no match here (best IoU %)\n", engine.classes[ref_class[rid]].name, formatFloat(ref_score[rid], trailing_width = 3), formatFloat(best, trailing_width = 3)); + c.worst = min(c.worst, best); + } + for result.segments if !used[it.id] c.extra += 1; + c.ok = true; + return c; +} + +// The segments in their classes' colours over the image, at the map's size, as a PNG. +write_overlay :: (engine: *Segment_Engine, image: Image, result: Segment_Result, path: string) { + w, h := result.width, result.height; + pixels := NewArray(cast(s64) w * h * 4, u8,, temp); + for y: 0..h - 1 for x: 0..w - 1 { + sx := min(cast(s32) ((x + 0.5) * image.width / w), image.width - 1); + sy := min(cast(s32) ((y + 0.5) * image.height / h), image.height - 1); + source := image.pixels + (cast(s64) sy * image.width + sx) * 4; + dest := pixels.data + (cast(s64) y * w + x) * 4; + id := result.map[cast(s64) y * w + x]; + if id { + colour := segment_colour(result.segments[id - 1].class); + for c: 0..2 dest[c] = cast(u8) ((cast(s32) source[c] + colour[c]) / 2); + } else { + for c: 0..2 dest[c] = source[c] / 3; // no segment: darkened + } + dest[3] = 255; + } + stbi_write_png(temp_c_string(path), w, h, 4, pixels.data, w * 4); +} + +ms :: (since: Apollo_Time) -> float64 { + return to_float64_seconds(current_time_monotonic() - since) * 1000; +} + +percent :: (fraction: float64) -> string { + return join(tprint("%", formatFloat(fraction * 100, trailing_width = 1)), "%",, temp); +} + +pad :: (text: string, width: s64) -> string { + if text.count >= width return text; + builder: String_Builder; + builder.allocator = temp; + append(*builder, text); + for 1..width - text.count append(*builder, " "); + return builder_to_string(*builder,, temp); +} + +#import "Basic"; +#import "String"; +#import "File"; +#import "File_Utilities"; +#import "System"; +#import "Image_Decode"; +#import "Segment"; +#import "stb_image_write";