Musa Mahmood
8c8b2c6e38
OCR filled the GPU: 9 GB for the mobile model on gpu-bench's 30 images, past 16 GB for the server model, which Windows then paged into shared memory (1.5 to 11 s a capture). onnxruntime's logged arena extensions showed why: the arena keeps every block it takes from CUDA and never joins or frees them, and OCR's inputs grow from run to run (recognizer batches went narrowest first; images in any order across calls), so each larger run left another block behind: 5.7 GB of recognizer blocks where the widest batch needs 0.9 GB. - Detection runs the largest images first, recognition the widest batches. - The detector's and recognizer's arenas shrink once each stage of an ocr_process call is done, by a run on a small input with memory.enable_memory_arena_shrinkage. Shrinking after every run cost 33 to 59% of the speed. - The worker threads are made with the engine and kept, not made for every stage of every call (onnxruntime keeps GPU state for each thread that runs a session, microsoft/onnxruntime#33037); 6% faster. - OCR's arenas grow by what is asked (kSameAsRequested). - ort.log_level sets what the engines' onnxruntime environments log; INFO shows every arena extension. Mobile OCR on gpu-bench (RTX 5080 laptop, 4 workers): peak 5.7 GB over the baseline, 2 GB between calls, 142 ms a capture, against 9.0 GB and 126 ms. The text read differs by a line in 2,986 (batches are grouped differently). The server model is not yet measured again.
606 lines
27 KiB
Plaintext
606 lines
27 KiB
Plaintext
// 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.log_level, "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 result as the chunks an .mcap stores, keyed by the model's name (modules/Capture_File/segment_chunks.jai):
|
|
// segs, with each segment's box in the pixels of the image read (image_width x image_height; which
|
|
// is "rendered" or "original"), largest first; and the mask. Temp.
|
|
segment_mcap_chunks :: (engine: *Segment_Engine, result: Segment_Result, which: string, image_width: s32, image_height: s32) -> segs: Mcap_Chunk, mask: Mcap_Chunk {
|
|
s: Mcap_Segments;
|
|
s.engine = tprint("onnxruntime %, %", ort.version, ifx engine.provider == .CUDA then "CUDA" else "CPU");
|
|
s.date = mcap_ocr_now();
|
|
s.image = which;
|
|
s.image_width, s.image_height = image_width, image_height;
|
|
s.classes = segment_classes_text(engine);
|
|
s.segments.allocator = temp;
|
|
sx := cast(float64) image_width / result.width;
|
|
sy := cast(float64) image_height / result.height;
|
|
map_area := cast(float64) result.width * result.height;
|
|
for result.segments {
|
|
x := cast(s32) (it.x0 * sx);
|
|
y := cast(s32) (it.y0 * sy);
|
|
array_add(*s.segments, .{id = it.id, score = it.score, area = cast(float) (it.pixels / map_area), x = x, y = y,
|
|
width = cast(s32) (it.x1 * sx + 0.5) - x, height = cast(s32) (it.y1 * sy + 0.5) - y, class = engine.classes[it.class].name});
|
|
}
|
|
quick_sort(s.segments, (a: Mcap_Segment, b: Mcap_Segment) -> s64 { return ifx a.area > b.area then -1 else ifx a.area < b.area then 1 else 0; });
|
|
name := engine.model.name;
|
|
return mcap_segments_chunk(name, s), mcap_mask_chunk(name, result.width, result.height, result.map);
|
|
}
|
|
|
|
// The class list's hash as segs stores it: 16 hex digits. Temp.
|
|
segment_classes_text :: (engine: *Segment_Engine) -> string {
|
|
return tprint("%", formatInt(engine.classes_hash, base = 16, minimum_digits = 16));
|
|
}
|
|
|
|
// 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 "Sort";
|
|
#import "Capture_File"; // segment_mcap_chunks
|
|
#import "onnxruntime";
|