[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.
This commit is contained in:
parent
775c2873e4
commit
cb21b73460
2
.gitignore
vendored
2
.gitignore
vendored
@ -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/
|
||||
|
||||
@ -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;
|
||||
|
||||
13
first.jai
13
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) {
|
||||
|
||||
75
modules/Segment/README.md
Normal file
75
modules/Segment/README.md
Normal file
@ -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.
|
||||
18
modules/Segment/models.config
Normal file
18
modules/Segment/models.config
Normal file
@ -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
|
||||
573
modules/Segment/module.jai
Normal file
573
modules/Segment/module.jai
Normal file
@ -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";
|
||||
@ -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.* = .{};
|
||||
|
||||
407
segmentation-plan.html
Normal file
407
segmentation-plan.html
Normal file
@ -0,0 +1,407 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Segmentation Plan</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #fbfaf8; --surface: #ffffff; --text: #1d1d1f; --dim: #5d5d66; --line: #e3e1dc;
|
||||
--accent: #3a5bd9; --accent-soft: #e9edfb; --warn: #9a5b00; --warn-soft: #fbf1e1;
|
||||
--code-bg: #f2f0ec; --ok: #2e7d4f; --ok-soft: #e5f3ea;
|
||||
}
|
||||
@media (prefers-color-scheme: dark) {
|
||||
:root:not([data-theme="light"]) {
|
||||
--bg: #16171a; --surface: #1d1f23; --text: #e7e7ea; --dim: #a2a3ab; --line: #30333a;
|
||||
--accent: #8aa2ff; --accent-soft: #232a45; --warn: #f0b45a; --warn-soft: #33281a;
|
||||
--code-bg: #25272c; --ok: #6fcf97; --ok-soft: #1c2e24;
|
||||
}
|
||||
}
|
||||
:root[data-theme="dark"] {
|
||||
--bg: #16171a; --surface: #1d1f23; --text: #e7e7ea; --dim: #a2a3ab; --line: #30333a;
|
||||
--accent: #8aa2ff; --accent-soft: #232a45; --warn: #f0b45a; --warn-soft: #33281a;
|
||||
--code-bg: #25272c; --ok: #6fcf97; --ok-soft: #1c2e24;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0; background: var(--bg); color: var(--text);
|
||||
font: 16px/1.6 system-ui, -apple-system, "Segoe UI", Roboto, sans-serif;
|
||||
}
|
||||
main { max-width: 860px; margin: 0 auto; padding: 40px 16px 80px; }
|
||||
header p { color: var(--dim); margin-top: 4px; }
|
||||
h1 { font-size: 2rem; line-height: 1.2; margin: 0; }
|
||||
h2 { font-size: 1.35rem; margin: 48px 0 8px; padding-top: 16px; border-top: 1px solid var(--line); }
|
||||
h3 { font-size: 1.05rem; margin: 28px 0 6px; }
|
||||
p, li { max-width: 72ch; }
|
||||
a { color: var(--accent); }
|
||||
code, pre { font-family: ui-monospace, "Cascadia Code", Consolas, monospace; font-size: 0.88em; }
|
||||
code { background: var(--code-bg); padding: 1px 5px; border-radius: 4px; }
|
||||
pre { background: var(--code-bg); padding: 12px 14px; border-radius: 8px; overflow-x: auto; line-height: 1.45; }
|
||||
pre code { background: none; padding: 0; }
|
||||
table { border-collapse: collapse; width: 100%; margin: 12px 0; font-size: 0.94rem; }
|
||||
th, td { text-align: left; padding: 7px 10px; border-bottom: 1px solid var(--line); vertical-align: top; }
|
||||
th { color: var(--dim); font-weight: 600; }
|
||||
.table-wrap { overflow-x: auto; }
|
||||
.summary { background: var(--surface); border: 1px solid var(--line); border-radius: 10px; padding: 16px 20px; margin-top: 24px; }
|
||||
.summary ul { margin: 6px 0 0; padding-left: 20px; }
|
||||
.note { background: var(--accent-soft); border-radius: 8px; padding: 10px 14px; margin: 14px 0; }
|
||||
.risk { background: var(--warn-soft); border-radius: 8px; padding: 10px 14px; margin: 14px 0; }
|
||||
.risk strong { color: var(--warn); }
|
||||
.tag { display: inline-block; font-size: 0.78rem; font-weight: 600; padding: 1px 8px; border-radius: 99px; background: var(--accent-soft); color: var(--accent); }
|
||||
.tag.later { background: var(--code-bg); color: var(--dim); }
|
||||
.tag.decided { background: var(--ok-soft); color: var(--ok); }
|
||||
nav ol { columns: 2; column-gap: 32px; padding-left: 20px; margin: 8px 0 0; }
|
||||
@media (max-width: 600px) { nav ol { columns: 1; } h1 { font-size: 1.6rem; } }
|
||||
.phase { border-left: 3px solid var(--accent); padding: 2px 0 2px 14px; margin: 16px 0; }
|
||||
.phase h3 { margin-top: 0; }
|
||||
figure { margin: 18px 0; }
|
||||
figcaption { color: var(--dim); font-size: 0.88rem; margin-top: 6px; }
|
||||
svg text { fill: var(--text); font: 13px system-ui, sans-serif; }
|
||||
svg .dim { fill: var(--dim); }
|
||||
svg .box { fill: var(--surface); stroke: var(--line); }
|
||||
svg .hl { fill: var(--accent-soft); stroke: var(--accent); }
|
||||
svg .arrow { stroke: var(--dim); fill: none; marker-end: url(#a); }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main>
|
||||
|
||||
<header>
|
||||
<h1>Semantic segmentation in mcap</h1>
|
||||
<p>Planning document · 2026-10-10 · for mcap after 0.0.10</p>
|
||||
</header>
|
||||
|
||||
<div class="summary">
|
||||
<strong>In short</strong>
|
||||
<ul>
|
||||
<li><span class="tag decided">Decided</span> <strong>Mask2Former</strong>, exported to ONNX and run with the onnxruntime mcap already ships. It is swappable later, because results are keyed by model name.</li>
|
||||
<li><span class="tag decided">Decided</span> <strong>Images only</strong>, and only in the <strong>library folders you opt in</strong>. The setting is per folder, kept in a file in the library root, so it travels with the library.</li>
|
||||
<li><span class="tag decided">Decided</span> <strong>No .mcap version bump.</strong> Results go in two new chunk kinds: a small <code>segs</code> summary at the front for search, and a <code>mask</code> at the back for drawing.</li>
|
||||
<li>It is built in six phases. The first proves the ONNX export matches PyTorch before anything else depends on it.</li>
|
||||
<li><span class="tag decided">Progress</span> <strong>Phases 1 and 2 are done</strong> (2026-10-10). The model is exported and checked, and <code>modules/Segment</code> runs it in mcap at about 57 ms per image on the RTX 5080. Next is phase 3: storing results in .mcap files.</li>
|
||||
<li><span class="tag later">Later</span> Video, by treating sampled frames as images, is shelved.</li>
|
||||
</ul>
|
||||
</div>
|
||||
|
||||
<nav>
|
||||
<ol>
|
||||
<li><a href="#scope">Scope</a></li>
|
||||
<li><a href="#model">The model</a></li>
|
||||
<li><a href="#folders">Per-folder opt-in</a></li>
|
||||
<li><a href="#scan">The segmentation scan</a></li>
|
||||
<li><a href="#storage">Storage in .mcap</a></li>
|
||||
<li><a href="#search">Catalog, search, Stats</a></li>
|
||||
<li><a href="#overlay">Capture window overlay</a></li>
|
||||
<li><a href="#phases">Phases</a></li>
|
||||
<li><a href="#questions">Open questions</a></li>
|
||||
<li><a href="#later">Later: video frames</a></li>
|
||||
</ol>
|
||||
</nav>
|
||||
|
||||
<h2 id="scope">1. Scope</h2>
|
||||
<p>Each opted-in image gets a <em>panoptic</em> 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:</p>
|
||||
<ul>
|
||||
<li><strong>Search:</strong> <code>has:dog</code>, <code>has:person</code>, and combined with the other search terms.</li>
|
||||
<li><strong>Stats:</strong> which things the shown captures hold most.</li>
|
||||
<li><strong>Looking:</strong> coloured masks over the image in the Capture window, with the class under the mouse.</li>
|
||||
</ul>
|
||||
<p><strong>What counts as an image.</strong> 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 <a href="#later">later</a> project. A capture is a video when its catalog entry's <code>media</code> is <code>VIDEO</code>, which comes from <code>VideoCodec</code> in its <code>import.json</code>.</p>
|
||||
|
||||
<h2 id="model">2. The model</h2>
|
||||
<h3>Which Mask2Former</h3>
|
||||
<div class="table-wrap">
|
||||
<table>
|
||||
<tr><th>Checkpoint (Hugging Face)</th><th>Classes</th><th>Size</th><th>Use</th></tr>
|
||||
<tr><td><code>facebook/mask2former-swin-base-coco-panoptic</code></td><td>133 COCO panoptic: 80 things, 53 stuff</td><td>~107M</td><td><strong>Default.</strong> Things are counted as instances, and the stuff classes (sky, wall, grass, road) cover the rest.</td></tr>
|
||||
<tr><td><code>facebook/mask2former-swin-large-coco-panoptic</code></td><td>same</td><td>~216M</td><td>If base misses too much and the GPU can take it.</td></tr>
|
||||
<tr><td><code>facebook/mask2former-swin-tiny-coco-panoptic</code></td><td>same</td><td>~47M</td><td>The CPU fallback, when CUDA is unavailable.</td></tr>
|
||||
<tr><td><code>facebook/mask2former-swin-base-ade-semantic</code></td><td>150 ADE20K, no instances</td><td>~107M</td><td>The alternative, if indoor scenes and objects matter more than counting people. See <a href="#questions">open questions</a>.</td></tr>
|
||||
</table>
|
||||
</div>
|
||||
<p>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.</p>
|
||||
|
||||
<h3>Export to ONNX</h3>
|
||||
<p>There is no official ONNX release, so phase 1 exports one with a Python script kept in the repo (<code>tools/export_mask2former.py</code>). It runs once on a developer machine and is never shipped.</p>
|
||||
<ul>
|
||||
<li><strong>Deformable attention.</strong> The pixel decoder uses multi-scale deformable attention. Transformers has a pure-PyTorch version built on <code>grid_sample</code>, which exports at opset 17 or later. The script forces that path.</li>
|
||||
<li><strong>Outputs.</strong> The script exports up to <code>class_queries_logits</code> [B, 100, 134] and <code>masks_queries_logits</code> [B, 100, H/4, W/4]. The panoptic merge (below) is done in Jai, as it is short, and keeping it out of the graph keeps the export simple.</li>
|
||||
<li><strong>Input.</strong> <code>pixel_values</code> [batch, 3, 384, 384]: RGB <em>squashed</em> to 384 × 384, ignoring the aspect ratio, as this checkpoint's processor does (bilinear), scaled to 0..1, then ImageNet mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225). The batch size is dynamic; the size is fixed.</li>
|
||||
<li><strong>Check.</strong> On 20 or more images, compare the ONNX output, run through mcap's post-processing, with transformers' <code>post_process_panoptic_segmentation</code>. Per-segment IoU should be at least 0.98.</li>
|
||||
</ul>
|
||||
<div class="note"><strong>Phase 1 result (2026-10-10): passed.</strong>
|
||||
<ul>
|
||||
<li><strong>Test set:</strong> 25 images: 16 COCO val2017 photos, 5 of your photos and 4 video frames.</li>
|
||||
<li><strong>Accuracy:</strong> 180 segments, every one with an IoU of 1.000 against PyTorch, none missing and none extra.</li>
|
||||
<li><strong>Batching:</strong> batches of 2, 4 and 8 give bit-identical results to one at a time.</li>
|
||||
<li><strong>Speed:</strong> CPU 348 ms per image at batch 1, 290 ms at batch 2.</li>
|
||||
<li><strong>Size:</strong> 462 MB at opset 17.</li>
|
||||
<li><strong>Load warning:</strong> the "missing" <code>swin.layernorm</code> that transformers 5.19 reports when loading is unused; randomizing it changes no output.</li>
|
||||
<li><strong>Not yet measured:</strong> GPU timing and memory. Those come in phase 2, with mcap's own onnxruntime build.</li>
|
||||
</ul>
|
||||
Script: <code>tools/export_mask2former.py</code>.</div>
|
||||
<div class="risk"><strong>Risk.</strong> The export is the one part of this plan that may fail outright. If <code>grid_sample</code> 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.</div>
|
||||
|
||||
<h3>Packaging, as Labels and OCR are</h3>
|
||||
<ul>
|
||||
<li>A new module, <code>modules/Segment</code>, laid out like <code>modules/Labels</code>.
|
||||
<ul>
|
||||
<li><code>models.config</code>: <code>@version 1</code>, then <code>[model mask2former-swin-base-coco-panoptic]</code> with <code>title</code>, <code>model</code> (the .onnx file name), <code>classes</code> (a text file with one class name per line, plus whether each is a "thing"), <code>input_size</code>, <code>batch_gpu</code>, <code>batch_cpu</code>, <code>score_threshold</code> (0.8) and <code>overlap_threshold</code> (0.8).</li>
|
||||
<li><code>fetch_models.jai</code>: the files, their SHA-256 hashes, and where they come from. The exported .onnx has to be hosted somewhere you control, such as a Hugging Face repo or a release asset. See <a href="#questions">open questions</a>.</li>
|
||||
<li><code>first.jai</code> copies the model folder to <code>bin/segment-models</code>, as it does <code>bin/label-models</code>.</li>
|
||||
</ul>
|
||||
</li>
|
||||
<li><code>segment_create</code> and <code>segment_destroy</code> follow <code>labels_create</code>. That means <code>ort_cuda_check</code>, a GPU session made under <code>ort_gpu_lock</code>, and a CPU fallback with the tiny model.</li>
|
||||
<li><code>segment_images(engine, images, cancel) -> [] Segment_Result</code>, given RGBA in, as <code>Labels_Image</code> and <code>Ocr_Image</code> are:
|
||||
<ol>
|
||||
<li>Squash to 384 × 384 and normalise, as above.</li>
|
||||
<li>Run in batches.</li>
|
||||
<li>The panoptic merge, as transformers 5.19 does it:
|
||||
<ol>
|
||||
<li>Softmax over the classes, and sigmoid over the masks scaled up to 384, then to the map.</li>
|
||||
<li>Drop queries scoring 0.5 or under, or whose class is "no object".</li>
|
||||
<li>Weight each mask by its query's score, and assign each pixel to its best query.</li>
|
||||
<li>Drop a segment if it keeps 80% or less of its own mask.</li>
|
||||
</ol>
|
||||
Stuff of one class is not merged.</li>
|
||||
<li>Result: a segment-id map at the mask resolution, plus a list of segments (id, class, score, area, box).</li>
|
||||
</ol>
|
||||
It checks <code>cancel</code> between batches, as <code>ocr_process</code> now does.</li>
|
||||
<li><code>tools/segment_probe.jai</code>, like <code>image-probe</code>: segment some files, print their segments, and write the masks out as PNGs to look at.</li>
|
||||
</ul>
|
||||
|
||||
<h2 id="folders">3. Per-folder opt-in</h2>
|
||||
<p>A library folder is a top-level folder of the library, one per place imported from (<code>SnagitLibraryWork</code>, <code>Phone</code>, <code>Photos 2014</code>). 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.</p>
|
||||
|
||||
<h3>Where the setting lives</h3>
|
||||
<p>It goes in a new file in the library root, <code>.library-folders</code>, next to <code>.import-index</code> and <code>.import-sources</code>. Not in <code>%APPDATA%\mcap\settings.config</code>, because the choice belongs to the library, not to the machine. It follows the library when it is opened elsewhere, or with <code>-library FOLDER</code>. 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:</p>
|
||||
<pre><code>@version 1
|
||||
|
||||
[folder "Photos 2014"]
|
||||
segment: true
|
||||
|
||||
[folder "SnagitLibraryWork"]
|
||||
segment: false</code></pre>
|
||||
<ul>
|
||||
<li><strong>Missing file or folder:</strong> every option is off. Unknown keys are kept on rewrite, as unknown chunks are.</li>
|
||||
<li><strong>Who reads and writes it:</strong> <code>library_folders_load</code> and <code>library_folders_save</code> in a new <code>src/app/library_folders.jai</code>, on the main thread, written beside itself and swapped in, as the import index is. It lives in a table keyed by folder name.</li>
|
||||
<li><strong>A renamed folder</strong> loses its setting. The Segmentation window lists settings for folders that no longer exist, with a button to remove them.</li>
|
||||
</ul>
|
||||
|
||||
<h3>Where you set it</h3>
|
||||
<ul>
|
||||
<li><strong>The Segmentation window</strong> (a new button next to Labels) has a table of library folders with these columns:
|
||||
<ul>
|
||||
<li><strong>Segment</strong>: the checkbox.</li>
|
||||
<li><strong>Images</strong>: the count, excluding videos.</li>
|
||||
<li><strong>Segmented</strong>: how many are done with the current model.</li>
|
||||
<li><strong>Erase</strong>: a per-folder button.</li>
|
||||
</ul>
|
||||
Turning a folder off stops future scans but erases nothing.</li>
|
||||
<li><strong>The Import window</strong> adds a checkbox, "Segment the images in this folder". It shows and sets the destination folder's option. For a new folder it starts off; for an existing folder it shows that folder's setting.</li>
|
||||
<li><strong>The folder combo above the grid</strong> marks opted-in folders, for example with a small "seg" after the name. This is a nicety; it can come last.</li>
|
||||
</ul>
|
||||
|
||||
<h2 id="scan">4. The segmentation scan</h2>
|
||||
<figure>
|
||||
<svg viewBox="0 0 820 190" width="100%" role="img" aria-label="Scan flow: opted-in folders give image items, which are decoded once, segmented in batches, then stored and handed to the catalog">
|
||||
<defs><marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0,0 L10,5 L0,10 z" fill="currentColor" style="color: var(--dim)"/></marker></defs>
|
||||
<rect class="box" x="10" y="30" width="170" height="70" rx="8"/>
|
||||
<text x="95" y="60" text-anchor="middle">.library-folders</text>
|
||||
<text x="95" y="80" text-anchor="middle" class="dim">segment: true</text>
|
||||
<rect class="box" x="220" y="30" width="170" height="70" rx="8"/>
|
||||
<text x="305" y="58" text-anchor="middle">Items</text>
|
||||
<text x="305" y="78" text-anchor="middle" class="dim">images not yet done</text>
|
||||
<rect class="hl" x="430" y="30" width="170" height="70" rx="8"/>
|
||||
<text x="515" y="58" text-anchor="middle">Decode, segment</text>
|
||||
<text x="515" y="78" text-anchor="middle" class="dim">batches on the scan thread</text>
|
||||
<rect class="box" x="640" y="30" width="170" height="70" rx="8"/>
|
||||
<text x="725" y="58" text-anchor="middle">mcap_store</text>
|
||||
<text x="725" y="78" text-anchor="middle" class="dim">segs front, mask back</text>
|
||||
<path class="arrow" d="M180 65 H216"/>
|
||||
<path class="arrow" d="M390 65 H426"/>
|
||||
<path class="arrow" d="M600 65 H636"/>
|
||||
<path class="arrow" d="M725 100 V140 H520"/>
|
||||
<text x="520" y="160" text-anchor="middle" class="dim">catalog_load → viewer.catalog_updates → search and Stats</text>
|
||||
</svg>
|
||||
<figcaption>Only the scan's choice of captures is new. Batching, pause and cancel, storing and the catalog hand-off are the OCR scan's.</figcaption>
|
||||
</figure>
|
||||
|
||||
<ul>
|
||||
<li><strong>The thread.</strong> It runs on the existing scan thread (the "ocr" thread in <code>ocr_ui.jai</code>), as a new task kind <code>SEGMENT</code>. One thread doing GPU work at a time is already the rule there, and it brings pause, cancel, the progress bar and <code>-scan</code>-style headless runs for free.</li>
|
||||
<li><strong>Which captures.</strong> Every capture in an opted-in folder whose media is not <code>VIDEO</code> and has no <code>segs</code> chunk from the current model. Unlike the OCR scan, it ignores the grid's folder filter and search: the scope is the folder settings. A "Segment the captures shown" button can be added later if wanted.</li>
|
||||
<li><strong>Batches.</strong> Batches of <code>SEGMENT_SCAN_BATCH</code> (4 to start, since Mask2Former is heavier per image than OCR). Each image is decoded once with <code>mcap_read_shown_image</code>, so a Snagit capture uses its rendered image. Cancel is checked between batches and inside <code>segment_images</code>; a cut-short batch stores nothing, as for OCR.</li>
|
||||
<li><strong>After an import.</strong> When "Read text and labels" is on and the destination folder segments, the import's follow-up scan becomes OCR and labels, then segmentation of the new images. These are two tasks queued one after the other, not one combined pass. They write different chunks, and keeping them apart keeps each simple.</li>
|
||||
<li><strong>Headless.</strong> <code>-segment</code> segments every opted-in folder and exits, 0 or 1, as <code>-scan</code> does.</li>
|
||||
<li><strong>Erase.</strong> <code>-erase segments</code>, and <code>-erase all</code> now includes it. Add <code>Erase_What.SEGMENTS</code>, its tags (<code>segs</code>, <code>mask</code>) and its "has" test. Per-folder erase in the Segmentation window uses the same task, limited to that folder's captures.</li>
|
||||
<li><strong>Shutdown.</strong> It is the same thread, so the 0.0.10 shutdown status covers it. Add a "Stopping the segmentation scan" line.</li>
|
||||
</ul>
|
||||
|
||||
<h2 id="storage">5. Storage in .mcap</h2>
|
||||
<div class="note">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 <code>labl</code> does.</div>
|
||||
|
||||
<h3><code>segs</code>, front chunk, keyed by model name</h3>
|
||||
<p>This is what search and Stats need, small enough for the front that every library browse reads. It is text in the style of <code>labl</code>:</p>
|
||||
<pre><code>@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"</code></pre>
|
||||
<div class="table-wrap">
|
||||
<table>
|
||||
<tr><th>Key</th><th>Value</th></tr>
|
||||
<tr><td><code>engine</code>, <code>date</code></td><td>As in OCR text.</td></tr>
|
||||
<tr><td><code>image</code></td><td>Which image was read (rendered or original) and its size. Boxes are in its pixels.</td></tr>
|
||||
<tr><td><code>classes</code></td><td>A 64-bit FNV-1a hash of the model's class list, so a model with new classes is seen as needing a new scan.</td></tr>
|
||||
<tr><td><code>segment</code></td><td>The 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.</td></tr>
|
||||
</table>
|
||||
</div>
|
||||
<p>A busy photo has 10 to 30 segments, about 1 to 2 KB.</p>
|
||||
|
||||
<h3><code>mask</code>, back chunk, keyed by model name</h3>
|
||||
<p>This is the segment-id map, for drawing only. It is binary:</p>
|
||||
<pre><code>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></pre>
|
||||
<ul>
|
||||
<li><strong>Why run-length encoding:</strong> segment maps are large flat regions, so runs are few. It decodes in one pass with no dependencies, and an expected 5 to 40 KB per photo is a few percent of a typical JPEG.</li>
|
||||
<li><strong>The alternative:</strong> a palette PNG through stb_image_write. It is smaller on noisy maps but encodes and decodes more slowly, and RLE is simpler to test. Revisit only if mask chunks turn out large.</li>
|
||||
<li><strong>Why at most 1024 on the long side:</strong> it is plenty for an overlay, and the model's own output is at a quarter of its input size, so storing more would invent detail. The overlay scales it up.</li>
|
||||
</ul>
|
||||
|
||||
<h3>Code changes</h3>
|
||||
<ul>
|
||||
<li><code>modules/Capture_File/segment_chunks.jai</code>, next to <code>label_chunks.jai</code>:
|
||||
<ul>
|
||||
<li>The tags <code>MCAP_SEGMENTS_TAG "segs"</code> and <code>MCAP_MASK_TAG "mask"</code>.</li>
|
||||
<li>The text read and write functions, <code>mcap_segments_text</code> and <code>mcap_segments_parse</code>.</li>
|
||||
<li>The mask functions, <code>mcap_mask_bytes</code> and <code>mcap_mask_parse</code>.</li>
|
||||
</ul></li>
|
||||
<li><code>mcap_store</code> gains a <code>back: [] Mcap_Chunk</code> parameter, so <code>segs</code> and <code>mask</code> are written in <strong>one</strong> rewrite. Today a back chunk needs its own call (<code>mcap_store_chunk(..., back = true)</code>), which would rewrite the file twice and could leave one without the other.</li>
|
||||
<li><code>mcap_read_without_original</code> already reads the back chunks, so the overlay needs no new reader.</li>
|
||||
<li>The format README gets a "Segmentation chunks" section.</li>
|
||||
</ul>
|
||||
|
||||
<h2 id="search">6. Catalog, search, Stats</h2>
|
||||
<ul>
|
||||
<li><strong>Catalog.</strong> <code>Catalog_Entry</code> gains:
|
||||
<ul>
|
||||
<li><code>segments: [..] Catalog_Segment</code>: class, area and score for each segment, from the front <code>segs</code> chunk keyed by the current model.</li>
|
||||
<li><code>segment_classes</code>: the hash, so stale results can be found.</li>
|
||||
<li><code>find_segments</code>: the class names, lowercased, one per line.</li>
|
||||
</ul>
|
||||
<code>catalog_entry_free</code> frees them. The mask is never loaded into the catalog.</li>
|
||||
<li><strong>Search.</strong> A new field, <code>has:</code> (<code>Search_Field.HAS</code>), matching a whole class name: <code>has:dog</code>, <code>has:"traffic light"</code>.
|
||||
<ul>
|
||||
<li>Plain words do not match segment classes, as with <code>folder:</code>, so <code>dog</code> keeps meaning text and titles, and <code>has:dog</code> means it is in the picture.</li>
|
||||
<li>Later: counts and sizes, such as <code>has:person:3+</code> or <code>has:dog>20%</code>. The front chunk already holds what they need.</li>
|
||||
</ul></li>
|
||||
<li><strong>Stats.</strong> Add "Segmented: N of M images shown" and a top-classes list: how many captures contain each class. Clicking a class searches <code>has:class</code>, as the top labels search <code>label:</code>.</li>
|
||||
</ul>
|
||||
<div class="note"><code>has:</code> and <code>label:</code> 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: <code>label:beach has:dog</code>.</div>
|
||||
|
||||
<h2 id="overlay">7. Capture window overlay</h2>
|
||||
<ul>
|
||||
<li>A <strong>Segments</strong> row under the labels line, shown when the capture has <code>segs</code> from the current model:
|
||||
<ul>
|
||||
<li>A toggle and an opacity slider.</li>
|
||||
<li>A legend: each class with its colour swatch and a checkbox to hide it. Clicking a class name searches <code>has:class</code>.</li>
|
||||
</ul></li>
|
||||
<li><strong>The mask texture.</strong> When the toggle is on, a viewer worker reads the capture's <code>mask</code> chunk (<code>mcap_read_without_original</code>, which does not read the original). It turns the run-length map into RGBA, colour per class and alpha from the slider, at the map's resolution. The main thread uploads it with <code>gpu_texture_create</code>, as it does the full image. Hidden classes are drawn transparent; changing which are hidden or the opacity rebuilds the texture on the worker, at most a few times a second.</li>
|
||||
<li><strong>Drawing.</strong> In <code>viewer_full_window</code>, after <code>Image(...)</code>, next to <code>ocr_ui_capture_overlay</code>: <code>draw.AddImage(mask_texture, image_min, image_min + source * zoom)</code>. The GPU scales it up, and the map's aspect ratio matches the image.</li>
|
||||
<li><strong>Hover.</strong> The pixel under the mouse is mapped into the map, and its segment's class and score are shown as a tooltip ("dog · 0.91 · 8% of the image").</li>
|
||||
<li><strong>Colours.</strong> A fixed palette by class index, the same in every capture, so a class keeps its colour. Pick it for contrast in both themes.</li>
|
||||
</ul>
|
||||
|
||||
<h2 id="phases">8. Phases</h2>
|
||||
<p>Each phase ends with something to run and check, and is a commit (0.0.11 onwards).</p>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 1: Export and check the model <span class="tag decided">Done</span></h3>
|
||||
<p><code>tools/export_mask2former.py</code> 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.</p>
|
||||
<p><strong>Done when:</strong> an .onnx that loads in onnxruntime 1.30 with CUDA and passes the IoU check, plus a note of the numbers.</p>
|
||||
</div>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 2: <code>modules/Segment</code> and the probe <span class="tag decided">Done</span></h3>
|
||||
<p>The module (create, segment, panoptic merge, cancel), <code>models.config</code>, the copy step in <code>first.jai</code>, and <code>tools/segment_probe.jai</code>. <code>fetch_models.jai</code> waits on where the model will be hosted (open question 2); until then <code>tools/export_mask2former.py</code> makes it.</p>
|
||||
<p><strong>Done when:</strong> <code>segment-probe</code> on the phase 1 photos writes masks that match phase 1's, and runs on both CPU and CUDA.</p>
|
||||
<div class="note"><strong>Result (2026-10-10).</strong>
|
||||
<ul>
|
||||
<li><strong>Agreement:</strong> 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.</li>
|
||||
<li><strong>Differences:</strong> 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.</li>
|
||||
<li><strong>Release build, RTX 5080:</strong> 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.</li>
|
||||
</ul>
|
||||
<p>Things learned on the way:</p>
|
||||
<ul>
|
||||
<li><strong>The score threshold.</strong> transformers 5.19's default is 0.5, not the 0.8 this plan assumed, and the check compares with that default. It's <code>score_threshold</code> in <code>models.config</code>, to raise if weak segments prove noisy in search.</li>
|
||||
<li><strong>Stuff isn't merged</strong> per class by default, so two patches of sky can be two segments.</li>
|
||||
<li><strong>Making segments</strong> now takes longer than the model. Spreading it across threads can roughly double throughput if the scan needs it.</li>
|
||||
<li><strong>Overlay colours.</strong> The overlay colours by class, so neighbouring people share a colour. Phase 6 should vary the shade per instance.</li>
|
||||
</ul></div>
|
||||
</div>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 3: Chunks and storing</h3>
|
||||
<p><code>segment_chunks.jai</code> (text, RLE), <code>mcap_store</code> with back chunks, the README section, and tests:</p>
|
||||
<ul>
|
||||
<li>Round trip.</li>
|
||||
<li>RLE edge cases: one segment covering everything, a map with no segments, runs wrapping across rows, the largest id.</li>
|
||||
<li>Storing <code>segs</code> and <code>mask</code> together leaves the original and every other chunk as they were.</li>
|
||||
</ul>
|
||||
<p><strong>Done when:</strong> the tests pass and a probe run can store its results into a scratch library's files.</p>
|
||||
</div>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 4: Per-folder opt-in and the scan</h3>
|
||||
<p>The pieces:</p>
|
||||
<ul>
|
||||
<li><code>.library-folders</code> (load, save, keep unknown keys).</li>
|
||||
<li>The Segmentation window's folder table.</li>
|
||||
<li>The Import window's checkbox.</li>
|
||||
<li>The <code>SEGMENT</code> task and its item selection.</li>
|
||||
<li>The import follow-up scan.</li>
|
||||
<li><code>-segment</code> and <code>-erase segments</code>.</li>
|
||||
<li>The shutdown status line.</li>
|
||||
</ul>
|
||||
<p><strong>Done when:</strong> 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.</p>
|
||||
</div>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 5: Catalog, <code>has:</code> and Stats</h3>
|
||||
<p><strong>Done when:</strong> <code>has:person</code> finds the right captures in the scratch library, catalog tests cover the new field, and Stats shows the top classes, each clickable.</p>
|
||||
</div>
|
||||
|
||||
<div class="phase">
|
||||
<h3>Phase 6: The overlay</h3>
|
||||
<p>The mask texture on a worker, the legend, hiding classes, opacity and hover.</p>
|
||||
<p><strong>Done when:</strong> 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.</p>
|
||||
</div>
|
||||
|
||||
<h2 id="questions">9. Open questions</h2>
|
||||
<ol>
|
||||
<li><strong>COCO panoptic or ADE20K?</strong> 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.</li>
|
||||
<li><strong>Where to host the exported .onnx</strong> for <code>fetch_models.jai</code>: 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.</li>
|
||||
<li><strong>The mask resolution cap.</strong> 1024 on the long side is a guess. Phase 3 should measure chunk sizes on real photos at 512, 1024 and 2048.</li>
|
||||
<li><strong>Snagit and screenshot folders.</strong> 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 <code>Format</code> and <code>AppName</code>.</li>
|
||||
<li><strong>Re-scanning on a model change.</strong> 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.</li>
|
||||
</ol>
|
||||
|
||||
<h2 id="later">10. Later: video frames <span class="tag later">Shelved</span></h2>
|
||||
<p>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:</p>
|
||||
<ul>
|
||||
<li><strong>Sampling:</strong> scene changes rather than a fixed rate. Compare small thumbnails of successive decoded frames, or use FFmpeg's scene score. Cap the frames per minute.</li>
|
||||
<li><strong>Storage:</strong> keep the frames themselves out of the .mcap except as small thumbnails. Store per-frame results in chunks keyed by model and frame time (for example a key like <code>mask2former…@83.40</code>), or in one chunk holding every frame's results.</li>
|
||||
<li>The per-folder option would grow a second key, <code>segment_videos</code>, in the same <code>.library-folders</code> file.</li>
|
||||
</ul>
|
||||
|
||||
</main>
|
||||
</body>
|
||||
</html>
|
||||
204
tools/export_mask2former.py
Normal file
204
tools/export_mask2former.py
Normal file
@ -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("<HH", w, h) + seg.astype(np.uint8).tobytes())
|
||||
with open(os.path.join(out, name + ".ref.txt"), "w", encoding="utf-8") as f:
|
||||
for s in result["segments_info"]:
|
||||
f.write(f"{s['id']} {s['label_id']} {s['score']:.6f}\n")
|
||||
print(f"wrote references for {len(files)} images to {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
236
tools/segment_probe.jai
Normal file
236
tools/segment_probe.jai
Normal file
@ -0,0 +1,236 @@
|
||||
// segment-probe: what modules/Segment makes of image files. For each FILE, its segments (class, score,
|
||||
// area); with -out FOLDER, the segments drawn over the image as a PNG, to look at; with -check FOLDER,
|
||||
// how they compare with PyTorch's own segmentation of the same image (NAME.ref and NAME.ref.txt, from
|
||||
// tools/export_mask2former.py --reference).
|
||||
//
|
||||
// Build: jai first.jai - segment_probe Run: bin/segment-probe FILE... [-out FOLDER] [-check FOLDER] [-cpu] [-dump-input FOLDER]
|
||||
|
||||
main :: () {
|
||||
args := get_command_line_arguments();
|
||||
files: [..] string;
|
||||
out_folder, check_folder, dump_folder: string;
|
||||
models_folder := tprint("%segment-models", path_strip_filename(get_path_of_running_executable()));
|
||||
model_name := "mask2former-swin-base-coco-panoptic";
|
||||
prefer_gpu := true;
|
||||
i := 1;
|
||||
while i < args.count {
|
||||
if args[i] == "-out" && i + 1 < args.count { out_folder = args[i + 1]; i += 2; continue; }
|
||||
if args[i] == "-check" && i + 1 < args.count { check_folder = args[i + 1]; i += 2; continue; }
|
||||
if args[i] == "-models" && i + 1 < args.count { models_folder = args[i + 1]; i += 2; continue; }
|
||||
if args[i] == "-model" && i + 1 < args.count { model_name = args[i + 1]; i += 2; continue; }
|
||||
if args[i] == "-cpu" { prefer_gpu = false; i += 1; continue; }
|
||||
if args[i] == "-dump-input" && i + 1 < args.count { dump_folder = args[i + 1]; i += 2; continue; }
|
||||
array_add(*files, args[i]);
|
||||
i += 1;
|
||||
}
|
||||
if !files.count { print("Usage: segment-probe FILE... [-out FOLDER] [-check FOLDER] [-cpu] [-model NAME] [-models FOLDER]\n"); exit(2); }
|
||||
if out_folder make_directory_if_it_does_not_exist(out_folder, recursive = true);
|
||||
|
||||
start := current_time_monotonic();
|
||||
engine, ok, error := segment_create(models_folder, model_name, prefer_gpu);
|
||||
if !ok { print("Cannot start %: %\n", model_name, error); exit(1); }
|
||||
defer segment_destroy(engine);
|
||||
print("% on %, % classes (made in % ms)%\n", model_name, engine.provider, engine.classes.count, ms(start),
|
||||
ifx engine.gpu_explanation then tprint(": %", engine.gpu_explanation) else "");
|
||||
|
||||
failures := 0;
|
||||
checked, matched, extra: s64;
|
||||
worst_iou: float64 = 1;
|
||||
agree, total_pixels, missed_confident: s64;
|
||||
model_ms, decode_ms: float64;
|
||||
done := 0;
|
||||
GROUP :: 8; // images decoded and segmented together
|
||||
position := 0;
|
||||
while position < files.count {
|
||||
group := array_view(files, position, min(GROUP, files.count - position));
|
||||
position += group.count;
|
||||
|
||||
images: [..] Image;
|
||||
names: [..] string;
|
||||
t := current_time_monotonic();
|
||||
for path: group {
|
||||
data, read_ok := read_entire_file(path, log_errors = false);
|
||||
if !read_ok { print("%: cannot be read\n", path); failures += 1; continue; }
|
||||
image, decoded, decode_error := image_decode(cast([] u8) data); // turned by its EXIF orientation, as mcap shows it
|
||||
free(data);
|
||||
if !decoded { print("%: cannot decode: %\n", path, decode_error); failures += 1; continue; }
|
||||
array_add(*images, image);
|
||||
array_add(*names, path);
|
||||
}
|
||||
decode_ms += ms(t);
|
||||
if !images.count continue;
|
||||
|
||||
inputs := NewArray(images.count, Segment_Image,, temp);
|
||||
for images inputs[it_index] = .{it.pixels, it.width, it.height, it.width * 4};
|
||||
if dump_folder { // the model's input, as raw floats, to compare with the checkpoint's processor
|
||||
make_directory_if_it_does_not_exist(dump_folder, recursive = true);
|
||||
size := engine.model.image_size;
|
||||
floats := NewArray(3 * size * size, float,, temp);
|
||||
for inputs {
|
||||
segment_prepare(engine, it, floats.data);
|
||||
write_entire_file(tprint("%/%.input", dump_folder, path_filename(names[it_index])), floats.data, floats.count * size_of(float));
|
||||
}
|
||||
}
|
||||
t = current_time_monotonic();
|
||||
results := segment_images(engine, inputs,, temp);
|
||||
model_ms += ms(t);
|
||||
|
||||
for result, n: results {
|
||||
path := names[n];
|
||||
image := images[n];
|
||||
done += 1;
|
||||
if result.failed { print("%: failed: %\n", path, result.error); failures += 1; continue; }
|
||||
print("% (% x %, map % x %): % segments\n", path_filename(path), image.width, image.height, result.width, result.height, result.segments.count);
|
||||
map_area := cast(float64) result.width * result.height;
|
||||
for result.segments print(" % % %\n", formatFloat(it.score, trailing_width = 3), pad(engine.classes[it.class].name, 22), percent(it.pixels / map_area));
|
||||
|
||||
if out_folder write_overlay(engine, image, result, tprint("%/%.segments.png", out_folder, path_strip_extension(path_filename(path))));
|
||||
if check_folder {
|
||||
c := compare(engine, result, tprint("%/%.ref", check_folder, path_filename(path)), tprint("%/%.ref.txt", check_folder, path_filename(path)));
|
||||
if !c.ok { print(" check: %\n", c.error); failures += 1; continue; }
|
||||
print(" check: pixels agreeing on their class %; segments: % of % matched, worst IoU %%\n", percent(c.agree / cast(float64) c.pixels),
|
||||
c.matched, c.checked, formatFloat(c.worst, trailing_width = 4), ifx c.extra then tprint(", % extra", c.extra) else "");
|
||||
checked += c.checked; matched += c.matched; extra += c.extra; worst_iou = min(worst_iou, c.worst);
|
||||
agree += c.agree; total_pixels += c.pixels; missed_confident += c.missed_confident;
|
||||
}
|
||||
}
|
||||
for * images image_free(it);
|
||||
reset_temporary_storage();
|
||||
}
|
||||
|
||||
print("\n% images: decoding % ms each, segmenting % ms each (model and post-processing, batch %)\n", done,
|
||||
formatFloat(decode_ms / max(done, 1), trailing_width = 0), formatFloat(model_ms / max(done, 1), trailing_width = 0),
|
||||
ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu);
|
||||
per := max(engine.stats.images, 1) * 1000.0;
|
||||
print(" of which: preparing % ms, the model % ms, making segments % ms, per image\n", formatFloat(engine.stats.prepare_us / per, trailing_width = 1),
|
||||
formatFloat(engine.stats.model_us / per, trailing_width = 1), formatFloat(engine.stats.post_us / per, trailing_width = 1));
|
||||
if check_folder && total_pixels {
|
||||
print("check: pixels agreeing on their class %; segments: % of % matched, % extra, worst IoU %\n",
|
||||
percent(agree / cast(float64) total_pixels), matched, checked, extra, formatFloat(worst_iou, trailing_width = 4));
|
||||
// Pass: nearly every pixel the same class, and every segment PyTorch is fairly sure of found. A
|
||||
// decoder or a rounding a step apart can tip a segment near a threshold either way.
|
||||
print("check: % confident segments (score at least %) missed\n", missed_confident, CONFIDENT);
|
||||
if cast(float64) agree < total_pixels * 0.995 || missed_confident failures += 1;
|
||||
}
|
||||
exit(xx ifx failures then 1 else 0);
|
||||
}
|
||||
|
||||
CONFIDENT :: 0.6; // a segment PyTorch scores this or more must be matched
|
||||
|
||||
Comparison :: struct {
|
||||
ok: bool;
|
||||
error: string;
|
||||
pixels, agree: s64; // pixels, and those given the same class (or none) by both
|
||||
checked, matched, extra: s64; // PyTorch's segments, those with one of ours of the same class at IoU >= 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";
|
||||
Loading…
Reference in New Issue
Block a user