Musa Mahmood
3e6e733a9d
- .library-folders in the library (src/app/library_folders.jai): each library folder's options, starting with segment: true. Every folder starts off; unknown keys are kept when it is rewritten. - The Segmentation window (Segments, above the grid): each folder's Segment checkbox, how many of its images are segmented, and Erase per folder; a button to segment the folders set to segment. - The segmentation scan (src/app/segment_ui.jai) runs on the OCR thread as its own task, with the OCR scan's progress, pause and cancel: images only (never videos), not yet segmented, in folders set to segment. The engine loads with the first scan (2.4 GB of GPU memory). Each image's segs and mask are stored in one rewrite and its catalog entry loaded again. - Importing into a folder set to segment segments the new images; the Import window's checkbox is the folder's setting. - -segment segments and exits; -erase segments, and -erase all now includes segments. - The catalog reads each capture's segs (class list hash, segment count). The exit overlay names the segmentation scan and model. segment_mcap_chunks moved into modules/Segment, shared with segment-probe. The quoting helpers of Capture_File's text formats are public. - Tests: .library-folders read and written back (quoted names, unknown keys, versions).
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_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 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";
|