Musa Mahmood
56588ff5e3
- The Capture window draws a segmented image's segments over it: each class in its colour, a shade of it for each instance, edges darker; under the OCR boxes, at any fit and zoom. A viewer worker reads the segs and mask chunks (never the images); the main thread makes the texture from the mask when the capture, the opacity or the classes hidden change. - Above the image: Segments on or off, the opacity, and a legend of the classes (largest first, with counts): untick one to hide it, right click it to search has: for it. Hovering the image shows the segment's class, score and share of the image. - Colours: a class's own is a hue from its name's hash. Its legend swatch is a colour picker (ImGui ColorEdit3) to choose another, kept in %APPDATA%\mcap\segment-colours.config for every image and library; right click it to go back to its own. - Tests: the overlay's pixels (shades, edges, hidden classes, opacity), reading a file's segments for it, and chosen colours. - This completes segmentation-plan.html's six phases.
735 lines
32 KiB
Plaintext
735 lines
32 KiB
Plaintext
// Segmentation in the app (modules/Segment): the segmentation scan, the Segmentation window, and
|
|
// -segment. A scan reads only images in the library folders set to segment (app/library_folders.jai),
|
|
// never a video, and stores each one's segments in its .mcap: a segs chunk in front, which the
|
|
// catalog reads, and a mask chunk at the back (modules/Capture_File/segment_chunks.jai).
|
|
//
|
|
// The scan runs on the OCR thread (app/ocr_ui.jai), as a task of its own, so it shares the OCR scan's
|
|
// progress, pause and cancel, and one thread uses the GPU at a time. The engine is made by the first
|
|
// scan, not at startup: it takes some 2.4 GB of GPU memory, and most sessions never segment anything.
|
|
|
|
SEGMENT_WINDOW :: "Segmentation";
|
|
SEGMENT_MODEL :: "mask2former-swin-base-coco-panoptic"; // the one segmentation model; its chunks are keyed by this
|
|
SEGMENT_SCAN_BATCH :: 8; // images read and segmented together (the engine runs them 4 at a time on the GPU)
|
|
|
|
segment_ui: struct {
|
|
// Owned by the OCR thread while ocr_ui.thread_busy, like ocr_ui.engine.
|
|
engine: *Segment_Engine;
|
|
engine_error: string;
|
|
|
|
// Main thread.
|
|
window_open: bool;
|
|
scan_requested: bool; // -segment has started its scan
|
|
logged_step: s64; // -segment's last progress line, in twentieths
|
|
erase_folder: string; // the folder whose segments the erase popup asks about; heap
|
|
erase_asked: bool;
|
|
|
|
// The window's counts per library folder, worked out again when the catalog changes.
|
|
counted_generation: s64 = -1;
|
|
counts: [..] Segment_Folder_Count;
|
|
|
|
overlay_state: Segment_Overlay_State; // the Capture window's overlay; main thread
|
|
}
|
|
|
|
Segment_Overlay_State :: struct {
|
|
show := true;
|
|
opacity: float = 0.45;
|
|
hidden: [..] string; // classes not drawn, by name, for every capture; heap
|
|
colours: [..] Segment_Colour; // classes given their own colour in the legend; kept in SEGMENT_COLOURS_FILE
|
|
overlay: *Segment_Overlay; // the capture's shown, or being replaced
|
|
entry: *Catalog_Entry; // the catalog entry it was asked for: a new one is read again
|
|
loading: bool;
|
|
texture: Gpu_Texture;
|
|
dirty: bool; // the texture is out of date
|
|
}
|
|
|
|
Segment_Folder_Count :: struct {
|
|
folder: string; // a view of viewer.folders' string, good until they change
|
|
images: s64; // captures that are not videos
|
|
segmented: s64; // of them, with this model's segments
|
|
catalogued: bool; // every capture in it has its catalog entry
|
|
}
|
|
|
|
// A class's colour, chosen in the Capture window's legend. Heap.
|
|
Segment_Colour :: struct {
|
|
class: string;
|
|
rgb: [3] u8;
|
|
}
|
|
|
|
// The colours chosen, one line each: colour: "person" #e05a47. In the app's data folder, beside the
|
|
// label list: a way of looking, so the same in every library.
|
|
SEGMENT_COLOURS_FILE :: "segment-colours.config";
|
|
|
|
segment_colours_path :: () -> string {
|
|
return tprint("%/%", app_program_data_dir(), SEGMENT_COLOURS_FILE);
|
|
}
|
|
|
|
segment_colours_load :: () {
|
|
o := *segment_ui.overlay_state;
|
|
o.colours.allocator = context.default_allocator;
|
|
text, ok := read_entire_file(segment_colours_path(), log_errors = false,, temp);
|
|
if !ok return;
|
|
for raw: split(text, "\n",, temp) {
|
|
line := trim(trim_right(raw, "\r"));
|
|
if !begins_with(line, "colour:") continue; // @version 1, comments, and keys a later build adds
|
|
class, rest, class_ok := mcap_read_quoted(trim(advance(line, 7)),, temp);
|
|
rest = trim(rest);
|
|
if !class_ok || rest.count != 7 || rest[0] != #char "#" continue;
|
|
value, value_ok := string_to_int(advance(rest, 1), 16, u32);
|
|
if !value_ok continue;
|
|
segment_colour_set(class, .[cast(u8) (value >> 16), cast(u8) ((value >> 8) & 0xFF), cast(u8) (value & 0xFF)], save = false);
|
|
}
|
|
}
|
|
|
|
segment_colours_save :: () {
|
|
builder: String_Builder;
|
|
builder.allocator = temp;
|
|
append(*builder, "@version 1\n# The colours of segment classes in the Capture window, chosen in its legend. A class not\n# listed has its own colour, from its name.\n");
|
|
for segment_ui.overlay_state.colours {
|
|
print_to_builder(*builder, "colour: % #%\n", mcap_quote(it.class), formatInt((cast(u32) it.rgb[0] << 16) | (cast(u32) it.rgb[1] << 8) | it.rgb[2], base = 16, minimum_digits = 6));
|
|
}
|
|
if !write_entire_file(segment_colours_path(), *builder) log_error("[Segment] Could not write %.", segment_colours_path());
|
|
}
|
|
|
|
// Gives class its own colour; or, with default, takes it back to the one from its name.
|
|
segment_colour_set :: (class: string, rgb: [3] u8, default := false, save := true) {
|
|
o := *segment_ui.overlay_state;
|
|
for * o.colours {
|
|
if it.class != class continue;
|
|
if default {
|
|
free(it.class,, context.default_allocator);
|
|
array_ordered_remove_by_index(*o.colours, it_index);
|
|
} else {
|
|
it.rgb = rgb;
|
|
}
|
|
if save segment_colours_save();
|
|
o.dirty = true;
|
|
return;
|
|
}
|
|
if default return;
|
|
array_add(*o.colours, .{copy_string(class,, context.default_allocator), rgb});
|
|
if save segment_colours_save();
|
|
o.dirty = true;
|
|
}
|
|
|
|
segment_models_folder :: () -> string {
|
|
return tprint("%segment-models", path_strip_filename(get_path_of_running_executable()));
|
|
}
|
|
|
|
// ---------------------------------------------------------------- startup and shutdown
|
|
|
|
segment_ui_startup :: () {
|
|
segment_ui.counts.allocator = context.default_allocator;
|
|
segment_ui.overlay_state.hidden.allocator = context.default_allocator;
|
|
segment_colours_load();
|
|
library_folders_load();
|
|
}
|
|
|
|
// After ocr_ui_shutdown has stopped the OCR thread.
|
|
segment_ui_shutdown :: () {
|
|
o := *segment_ui.overlay_state; // the viewer's workers have stopped; the texture goes with the device
|
|
segment_overlay_free(o.overlay);
|
|
segment_overlay_free(viewer.segments_done);
|
|
for o.hidden free(it,, context.default_allocator);
|
|
array_free(o.hidden);
|
|
for o.colours free(it.class,, context.default_allocator);
|
|
array_free(o.colours);
|
|
segment_destroy(segment_ui.engine);
|
|
segment_ui.engine = null;
|
|
free(segment_ui.engine_error,, context.default_allocator);
|
|
free(segment_ui.erase_folder,, context.default_allocator);
|
|
array_free(segment_ui.counts);
|
|
}
|
|
|
|
// ---------------------------------------------------------------- the scan
|
|
|
|
// Starts a scan of the images not yet segmented in the folders set to segment, or of those among only
|
|
// (an import's). False, with a toast, when there is nothing to do or the OCR thread is busy.
|
|
segment_scan_start :: (only: [] s32 = .[], quiet := false) -> bool {
|
|
if ocr_ui.thread_busy return false;
|
|
scan := *ocr_ui.scan;
|
|
for scan.items free(it.path);
|
|
array_reset(*scan.items);
|
|
free(scan.model);
|
|
scan.* = .{};
|
|
scan.items.allocator = context.default_allocator;
|
|
scan.model = copy_string(SEGMENT_MODEL,, context.default_allocator);
|
|
scan.segmenting = true;
|
|
current := ifx segment_ui.engine then segment_classes_text(segment_ui.engine) else ""; // the OCR thread is idle: the engine is ours
|
|
|
|
indices: [..] s32;
|
|
indices.allocator = temp;
|
|
if only.count array_add(*indices, ..only);
|
|
else for viewer.captures array_add(*indices, cast(s32) it_index);
|
|
for indices {
|
|
capture := *viewer.captures[it];
|
|
if capture.removed || !library_folder_segments(capture.folder) continue;
|
|
entry := capture.catalog;
|
|
if !entry || entry.media == .VIDEO continue; // images only; the catalog says which
|
|
// Done already, unless with another class list than the engine's (when it is up to say).
|
|
if entry.segment_classes && (!current || entry.segment_classes == current) continue;
|
|
array_add(*scan.items, .{index = it, path = copy_string(capture.path,, context.default_allocator)});
|
|
}
|
|
scan.total = scan.items.count;
|
|
if !scan.total {
|
|
if !quiet toast_show(.Success, "Every image in the folders set to segment has been segmented");
|
|
return false;
|
|
}
|
|
scan.running = true;
|
|
scan.start = current_time_monotonic();
|
|
return ocr_ui_start(.{kind = .SEGMENT});
|
|
}
|
|
|
|
// On the OCR thread: the SEGMENT task. Makes the engine the first time, then segments the scan's items
|
|
// a batch at a time, storing each.
|
|
segment_ui_scan_task :: () {
|
|
scan := *ocr_ui.scan;
|
|
defer atomic_swap(*scan.running, false);
|
|
if !segment_ui.engine {
|
|
start := current_time_monotonic();
|
|
engine, ok, error := segment_create(segment_models_folder(), SEGMENT_MODEL);
|
|
free(segment_ui.engine_error);
|
|
segment_ui.engine_error = ifx ok then "" else copy_string(error);
|
|
if !ok {
|
|
log_error("[Segment] Could not start %: %", SEGMENT_MODEL, error);
|
|
atomic_swap(*scan.failed, scan.total);
|
|
atomic_swap(*scan.done, scan.total);
|
|
return;
|
|
}
|
|
segment_ui.engine = engine;
|
|
log("[Segment] % on % in % ms.%", SEGMENT_MODEL, engine.provider, to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
|
|
}
|
|
engine := segment_ui.engine;
|
|
// Cancelling (or closing the app) stops it between batches of the engine's, rather than waiting the
|
|
// scan's batch out. What it had not segmented is left as it was.
|
|
engine.cancel = *scan.cancel;
|
|
defer engine.cancel = null;
|
|
|
|
position := 0;
|
|
while position < scan.items.count {
|
|
while scan.paused && !scan.cancel sleep_milliseconds(100);
|
|
if scan.cancel break;
|
|
count := min(SEGMENT_SCAN_BATCH, scan.items.count - position);
|
|
errors := segment_read_and_store(engine, array_view(scan.items, position, count));
|
|
if scan.cancel { // cut short: what it stored stays stored and is counted; the rest is left, not failed
|
|
for errors {
|
|
if !it atomic_add(*scan.done, 1);
|
|
free(it);
|
|
}
|
|
break;
|
|
}
|
|
for errors if it { atomic_add(*scan.failed, 1); log_error("[Segment scan] % %", scan.items[position + it_index].path, it); free(it); }
|
|
atomic_add(*scan.done, count);
|
|
position += count;
|
|
reset_temporary_storage();
|
|
}
|
|
scan.elapsed_seconds = to_float64_seconds(current_time_monotonic() - scan.start) - scan.paused_seconds;
|
|
log("[Segment scan] % of % images segmented with % in % (% failed)%.", scan.done, scan.total, SEGMENT_MODEL,
|
|
ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx scan.cancel then ", cancelled" else "");
|
|
}
|
|
|
|
// On the OCR thread: segments the items' images and stores each one's segments, handing its new
|
|
// catalog entry to the viewer. An error message (heap) for each that failed, "" for the others. A
|
|
// video is skipped, not failed: the catalog that chose it may have been out of date.
|
|
segment_read_and_store :: (engine: *Segment_Engine, items: [] Scan_Item) -> [] string {
|
|
errors := NewArray(items.count, string);
|
|
Prepared :: struct { image: Image; which: string; ok: bool; }
|
|
prepared := NewArray(items.count, Prepared,, temp);
|
|
defer for * prepared if it.ok image_free(*it.image);
|
|
inputs: [..] Segment_Image;
|
|
inputs.allocator = temp;
|
|
input_of: [..] s64; // which item each input is
|
|
input_of.allocator = temp;
|
|
for item, n: items {
|
|
m, bytes, which, ok, error := mcap_read_shown_image(item.path,, temp);
|
|
if !ok { errors[n] = copy_string(error); continue; }
|
|
if contains(m.text, "VideoCodec:") continue; // a video (app/import.jai): images only
|
|
image, decoded, decode_error := image_decode(bytes); // turned by its EXIF orientation, as the Capture window shows it
|
|
if !decoded { errors[n] = sprint("its image cannot be decoded: %", decode_error); continue; }
|
|
prepared[n] = .{image, which, true};
|
|
array_add(*inputs, .{image.pixels, image.width, image.height, image.width * 4});
|
|
array_add(*input_of, n);
|
|
}
|
|
if !inputs.count return errors;
|
|
|
|
results := segment_images(engine, inputs,, temp);
|
|
for result, r: results {
|
|
n := input_of[r];
|
|
item := items[n];
|
|
if result.failed { errors[n] = copy_string(result.error); continue; }
|
|
p := prepared[n];
|
|
segs, mask := segment_mcap_chunks(engine, result, p.which, p.image.width, p.image.height);
|
|
stored: bool;
|
|
store_error: string;
|
|
for 1..OCR_STORE_ATTEMPTS {
|
|
stored, store_error = mcap_store(item.path, front = .[segs], back = .[mask]);
|
|
if stored break;
|
|
sleep_milliseconds(250);
|
|
}
|
|
if !stored { errors[n] = sprint("cannot be stored: %", store_error); continue; }
|
|
entry, load_error := catalog_load(item.path);
|
|
if entry {
|
|
lock(*viewer.mutex);
|
|
array_add(*viewer.catalog_updates, .{item.index, entry});
|
|
unlock(*viewer.mutex);
|
|
} else {
|
|
errors[n] = sprint("was segmented, but could not be read back: %", load_error);
|
|
}
|
|
}
|
|
return errors;
|
|
}
|
|
|
|
// -segment: segments the folders set to segment once the catalog has loaded, and ends the app when
|
|
// done (0, or 1 if any image failed).
|
|
segment_ui_segment_and_exit :: (busy: bool) {
|
|
finish :: (code: s32, format: string, args: ..Any) {
|
|
if code log_error(format, ..args);
|
|
else log(format, ..args);
|
|
app.exit_code = code;
|
|
app.should_exit = true;
|
|
app.segment_and_exit = false;
|
|
}
|
|
if viewer.scan_error { finish(1, "[Segment scan] %", viewer.scan_error); return; }
|
|
if !segment_ui.scan_requested {
|
|
// After the catalog, and after the engines made at startup, which use the same thread.
|
|
if busy || viewer.catalog_loaded < viewer.captures.count || !ocr_ui.engine_tried || !labels_ui.engine_tried return;
|
|
segment_ui.scan_requested = true;
|
|
segment_ui.window_open = true; // its progress bar
|
|
if !segment_scan_start(quiet = true) finish(0, "[Segment scan] -segment: nothing to segment (no image in a folder set to segment needs it).");
|
|
return;
|
|
}
|
|
scan := *ocr_ui.scan;
|
|
if scan.running {
|
|
done, total, failed, rate, left, elapsed := ocr_scan_progress();
|
|
step := max(total / 20, 1);
|
|
if done / step > segment_ui.logged_step {
|
|
segment_ui.logged_step = done / step;
|
|
log("[Segment scan] % segmented, % failed, % elapsed, % left, % images/s.", ocr_scan_progress_text(done, total), failed, ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1));
|
|
}
|
|
return;
|
|
}
|
|
if busy return;
|
|
finish(ifx scan.failed || segment_ui.engine_error then cast(s32) 1 else 0, "[Segment scan] -segment is done: % of % segmented in %, % failed.%",
|
|
scan.done - scan.failed, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx segment_ui.engine_error then tprint(" %", segment_ui.engine_error) else "");
|
|
}
|
|
|
|
// ---------------------------------------------------------------- the Segmentation window
|
|
|
|
segment_ui_frame :: () { using ImGui;
|
|
busy := ocr_ui.thread_busy;
|
|
if app.segment_and_exit segment_ui_segment_and_exit(busy);
|
|
segment_erase_popup();
|
|
if !segment_ui.window_open return;
|
|
SetNextWindowSize(.{640, 520}, .FirstUseEver);
|
|
if Begin(SEGMENT_WINDOW, *segment_ui.window_open) {
|
|
segment_window_contents();
|
|
}
|
|
End();
|
|
}
|
|
|
|
segment_window_contents :: () { using ImGui;
|
|
busy := ocr_ui.thread_busy;
|
|
scan := *ocr_ui.scan;
|
|
|
|
// The model.
|
|
model_file := tprint("%/%.onnx", segment_models_folder(), SEGMENT_MODEL);
|
|
if !busy && segment_ui.engine {
|
|
Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider == .CUDA then "the GPU" else "the CPU");
|
|
if segment_ui.engine.gpu_explanation TextColored(app_theme.palette.dim, "%", segment_ui.engine.gpu_explanation);
|
|
} else if !busy && segment_ui.engine_error {
|
|
TextColored(app_theme.palette.warning, "% could not start: %", SEGMENT_MODEL, segment_ui.engine_error);
|
|
} else if !file_exists(model_file) {
|
|
TextColored(app_theme.palette.warning, "The model is not there (%): make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild.", model_file);
|
|
} else {
|
|
Text("%", SEGMENT_MODEL);
|
|
TextColored(app_theme.palette.dim, "Loads when a scan starts (about 2.4 GB of GPU memory).");
|
|
}
|
|
TextColored(app_theme.palette.dim, "Finds what is where in each image: people, animals, things, and stuff such as sky, wall and grass. Only in the folders set to segment, and never in videos.");
|
|
Spacing();
|
|
|
|
// The library folders.
|
|
segment_update_counts();
|
|
if BeginTable("folders", 4, .Borders | .RowBg | .SizingStretchProp) {
|
|
TableSetupColumn("Library folder", .WidthStretch);
|
|
TableSetupColumn("Segment", .WidthFixed);
|
|
TableSetupColumn("Segmented", .WidthFixed);
|
|
TableSetupColumn("", .WidthFixed);
|
|
TableHeadersRow();
|
|
for segment_ui.counts {
|
|
PushID(cast(s32) it_index);
|
|
TableNextRow();
|
|
TableNextColumn();
|
|
TextUnformatted(ifx it.folder then it.folder else "(the library's top)");
|
|
TableNextColumn();
|
|
on := library_folder_segments(it.folder);
|
|
if busy && scan.segmenting BeginDisabled();
|
|
if Checkbox("##segment", *on) library_folder_set_segment(it.folder, on);
|
|
if busy && scan.segmenting EndDisabled();
|
|
if IsItemHovered(.AllowWhenDisabled) SetTooltip(ifx on then "Scans read this folder's images" else "Scans leave this folder alone");
|
|
TableNextColumn();
|
|
if it.catalogued Text("% of %", it.segmented, it.images);
|
|
else TextColored(app_theme.palette.dim, "% of % ...", it.segmented, it.images);
|
|
TableNextColumn();
|
|
can_erase := it.segmented > 0 && !busy;
|
|
if !can_erase BeginDisabled();
|
|
if SmallButton("Erase") {
|
|
free(segment_ui.erase_folder,, context.default_allocator);
|
|
segment_ui.erase_folder = copy_string(it.folder,, context.default_allocator);
|
|
segment_ui.erase_asked = true;
|
|
}
|
|
if !can_erase EndDisabled();
|
|
if IsItemHovered(.AllowWhenDisabled) SetTooltip("Takes this folder's segments out of its files (asks first)");
|
|
PopID();
|
|
}
|
|
EndTable();
|
|
}
|
|
gone := library_folders_gone();
|
|
for gone {
|
|
TextColored(app_theme.palette.dim, "\"%\" is set to segment, but the library has no folder by that name.", it);
|
|
SameLine();
|
|
PushID(temp_c_string(it));
|
|
if SmallButton("Forget it") library_folder_forget(it);
|
|
PopID();
|
|
}
|
|
Spacing();
|
|
|
|
// The scan.
|
|
if scan.running && scan.segmenting {
|
|
Text("Segmenting:");
|
|
ocr_scan_progress_ui();
|
|
return;
|
|
}
|
|
any_on := false;
|
|
for segment_ui.counts if library_folder_segments(it.folder) any_on = true;
|
|
catalog_ready := viewer.catalog_loaded >= viewer.captures.count;
|
|
can_scan := !busy && catalog_ready && any_on;
|
|
if !can_scan BeginDisabled();
|
|
if Button("Segment the folders set to segment") segment_scan_start();
|
|
if !can_scan EndDisabled();
|
|
if IsItemHovered(.AllowWhenDisabled) {
|
|
if !any_on SetTooltip("Tick Segment for a folder first");
|
|
else if !catalog_ready SetTooltip("Once the catalog has loaded");
|
|
else if busy SetTooltip("After what the OCR thread is doing (a scan, or loading a model)");
|
|
else SetTooltip("Reads every image in those folders not yet segmented, storing each as it goes: a scan can stop and start again");
|
|
}
|
|
if scan.segmenting && scan.total && scan.done {
|
|
TextColored(app_theme.palette.dim, "Last scan: % of % images in %, % failed.", scan.done, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed);
|
|
}
|
|
}
|
|
|
|
// The window's per-folder counts, worked out again only when the catalog has changed.
|
|
segment_update_counts :: () {
|
|
if segment_ui.counted_generation == viewer.catalog_generation && segment_ui.counts.count == viewer.folders.count return;
|
|
segment_ui.counted_generation = viewer.catalog_generation;
|
|
array_reset_keeping_memory(*segment_ui.counts);
|
|
for viewer.folders array_add(*segment_ui.counts, .{folder = it, catalogued = true});
|
|
for viewer.captures {
|
|
if it.removed continue;
|
|
count: *Segment_Folder_Count;
|
|
for * c: segment_ui.counts if c.folder == it.folder { count = c; break; }
|
|
if !count continue;
|
|
entry := it.catalog;
|
|
if !entry { count.catalogued = false; continue; }
|
|
if entry.media == .VIDEO continue;
|
|
count.images += 1;
|
|
if entry.segment_classes count.segmented += 1;
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------- the Capture window's overlay
|
|
//
|
|
// The segments of the capture shown, over its image: each segment in its class's colour (a shade of
|
|
// it for each instance, so two people side by side differ), its edges darker, with the class under
|
|
// the mouse in a tooltip. A viewer worker reads the segs and mask chunks (never the images); the
|
|
// main thread makes the texture from the mask, again only when the capture, the opacity or the
|
|
// classes hidden change.
|
|
|
|
// A capture's segments as read for the overlay. Heap, with everything in it.
|
|
Segment_Overlay :: struct {
|
|
index: s32 = -1;
|
|
error: string; // why there is none
|
|
width, height: s32; // the mask's
|
|
map: [] u8; // a segment id per pixel; 0: none
|
|
segments: [] Mcap_Segment; // largest first
|
|
by_id: [256] s16; // into segments, for each id; -1: none
|
|
}
|
|
|
|
// On a viewer worker: the capture's segments by SEGMENT_MODEL, or why not.
|
|
segment_overlay_read :: (path: string) -> *Segment_Overlay {
|
|
overlay := New(Segment_Overlay,, context.default_allocator);
|
|
for * overlay.by_id it.* = -1;
|
|
m, original, ok, error := mcap_read_without_original(path,, temp);
|
|
if !ok { overlay.error = copy_string(error,, context.default_allocator); return overlay; }
|
|
segs := mcap_find_chunk(m.front_chunks, mcap_tag(MCAP_SEGMENTS_TAG), SEGMENT_MODEL);
|
|
mask := mcap_find_chunk(m.back_chunks, mcap_tag(MCAP_MASK_TAG), SEGMENT_MODEL);
|
|
if !segs || !mask { overlay.error = copy_string("it has not been segmented",, context.default_allocator); return overlay; }
|
|
segments, segments_ok, segments_error := mcap_segments_parse(cast(string) segs.data,, temp);
|
|
map, map_ok, map_error := mcap_mask_parse(mask.data,, temp);
|
|
if !segments_ok || !map_ok { overlay.error = sprint("its segments cannot be read: %", ifx segments_error then segments_error else map_error,, context.default_allocator); return overlay; }
|
|
|
|
push_allocator(context.default_allocator);
|
|
overlay.width, overlay.height = map.width, map.height;
|
|
overlay.map = array_copy(map.map);
|
|
overlay.segments = NewArray(segments.segments.count, Mcap_Segment);
|
|
for segments.segments {
|
|
overlay.segments[it_index] = it;
|
|
overlay.segments[it_index].class = copy_string(it.class);
|
|
overlay.by_id[it.id] = cast(s16) it_index;
|
|
}
|
|
return overlay;
|
|
}
|
|
|
|
segment_overlay_free :: (overlay: *Segment_Overlay) {
|
|
if !overlay return;
|
|
push_allocator(context.default_allocator);
|
|
free(overlay.error);
|
|
array_free(overlay.map);
|
|
for overlay.segments free(it.class);
|
|
array_free(overlay.segments);
|
|
free(overlay);
|
|
}
|
|
|
|
// Main thread, each frame: takes what the workers read, asks for the capture shown, and makes the
|
|
// texture when it is out of date.
|
|
segment_overlay_frame :: () {
|
|
o := *segment_ui.overlay_state;
|
|
lock(*viewer.mutex);
|
|
done := viewer.segments_done;
|
|
arrived := viewer.segments_done_request != 0 && viewer.segments_done_request == viewer.segments_request;
|
|
viewer.segments_done = null;
|
|
viewer.segments_done_request = 0;
|
|
unlock(*viewer.mutex);
|
|
if arrived {
|
|
segment_overlay_free(o.overlay);
|
|
o.overlay = done;
|
|
o.loading = false;
|
|
o.dirty = true;
|
|
} else {
|
|
segment_overlay_free(done);
|
|
}
|
|
|
|
// Nothing shown: let it all go.
|
|
index := viewer.full_index;
|
|
entry: *Catalog_Entry;
|
|
if viewer.full_open && index >= 0 && index < viewer.captures.count && !viewer.captures[index].removed entry = viewer.captures[index].catalog;
|
|
wanted := o.show && entry && entry.segment_classes.count > 0;
|
|
if !wanted {
|
|
gpu_texture_release(*o.texture);
|
|
if !viewer.full_open { segment_overlay_free(o.overlay); o.overlay = null; o.entry = null; o.loading = false; }
|
|
return;
|
|
}
|
|
|
|
// The capture shown, read again when its catalog entry is new (stepped to, or segmented again).
|
|
if o.entry != entry {
|
|
o.entry = entry;
|
|
o.loading = true;
|
|
lock(*viewer.mutex);
|
|
viewer.segments_wanted = index;
|
|
viewer.segments_request += 1;
|
|
unlock(*viewer.mutex);
|
|
signal(*viewer.wake);
|
|
}
|
|
if o.dirty && o.overlay && o.overlay.index == index {
|
|
o.dirty = false;
|
|
segment_overlay_build();
|
|
}
|
|
}
|
|
|
|
// The device went (graphics/DX11.jai): the texture with it; made again from the mask held.
|
|
segment_overlay_device_lost :: () {
|
|
o := *segment_ui.overlay_state;
|
|
gpu_texture_release(*o.texture);
|
|
o.dirty = true;
|
|
}
|
|
|
|
// The row above the image: the toggle, the opacity, and the classes in it, each to hide or show
|
|
// (right click to search for it).
|
|
segment_overlay_controls :: (index: s32) { using ImGui;
|
|
if index < 0 || index >= viewer.captures.count return;
|
|
entry := viewer.captures[index].catalog;
|
|
if !entry || !entry.segment_classes return;
|
|
o := *segment_ui.overlay_state;
|
|
|
|
if Checkbox("Segments", *o.show) o.dirty = true;
|
|
if IsItemHovered() SetTooltip("What segmentation found in this image, drawn over it (%)", SEGMENT_MODEL);
|
|
if !o.show { SameLine(); TextColored(app_theme.palette.dim, "% found", entry.segment_count); return; }
|
|
SameLine();
|
|
SetNextItemWidth(140);
|
|
percent := o.opacity * 100;
|
|
if SliderFloat("##opacity", *percent, 10, 90, "opacity %.0f") { o.opacity = percent / 100; o.dirty = true; }
|
|
if o.loading { SameLine(); TextColored(app_theme.palette.dim, "Loading..."); }
|
|
if o.overlay && o.overlay.index == index && o.overlay.error { SameLine(); TextColored(app_theme.palette.warning, "%", o.overlay.error); }
|
|
|
|
// The legend: each class once, largest first, with its colour and how many of it.
|
|
right_edge := GetCursorScreenPos().x + GetContentRegionAvail().x;
|
|
first := true;
|
|
for segment, n: entry.segments {
|
|
seen := false;
|
|
for 0..n - 1 if entry.segments[it].class == segment.class { seen = true; break; }
|
|
if seen continue;
|
|
count := 0;
|
|
for entry.segments if it.class == segment.class count += 1;
|
|
label := ifx count > 1 then tprint("% x%##legend%", segment.class, count, n) else tprint("%##legend%", segment.class, n);
|
|
width := GetFrameHeight() + GetStyle().ItemInnerSpacing.x + GetFrameHeight() + GetStyle().ItemInnerSpacing.x + CalcTextSize(temp_c_string(label), null, true).x;
|
|
if !first {
|
|
SameLine();
|
|
if GetCursorScreenPos().x + width > right_edge NewLine();
|
|
}
|
|
first = false;
|
|
|
|
// Its colour: click to choose another (kept for every capture), right click to go back to its own.
|
|
PushID(temp_c_string(segment.class));
|
|
base := segment_class_base_colour(segment.class);
|
|
rgb := float.[base[0] / 255.0, base[1] / 255.0, base[2] / 255.0];
|
|
if ColorEdit3("##colour", *rgb, .NoInputs | .NoLabel | .NoAlpha | .NoOptions | .PickerHueWheel) {
|
|
segment_colour_set(segment.class, .[cast(u8) (rgb[0] * 255 + 0.5), cast(u8) (rgb[1] * 255 + 0.5), cast(u8) (rgb[2] * 255 + 0.5)], save = false);
|
|
}
|
|
if IsItemDeactivatedAfterEdit() segment_colours_save(); // once the picker closes, not every frame of a drag
|
|
chosen := false;
|
|
for o.colours if it.class == segment.class chosen = true;
|
|
if BeginPopupContextItem("##reset") {
|
|
if MenuItem("Back to the default colour", null, false, chosen) segment_colour_set(segment.class, .[0, 0, 0], default = true);
|
|
EndPopup();
|
|
}
|
|
PopID();
|
|
SameLine(0, GetStyle().ItemInnerSpacing.x);
|
|
shown := !array_find(o.hidden, segment.class);
|
|
if Checkbox(temp_c_string(label), *shown) {
|
|
if shown { found, at_hidden := array_find(o.hidden, segment.class); if found { free(o.hidden[at_hidden],, context.default_allocator); array_unordered_remove_by_index(*o.hidden, at_hidden); } }
|
|
else array_add(*o.hidden, copy_string(segment.class,, context.default_allocator));
|
|
o.dirty = true;
|
|
}
|
|
if IsItemClicked(.Right) viewer_set_search(tprint("has:\"%\"", segment.class));
|
|
if IsItemHovered() SetTooltip("Shows or hides %; right click to search has:\"%\"", segment.class, segment.class);
|
|
}
|
|
}
|
|
|
|
// Over the image drawn from image_min to image_max: the overlay's texture, and the segment under the mouse.
|
|
segment_overlay_draw :: (index: s32, image_min: ImGui.ImVec2, image_max: ImGui.ImVec2) { using ImGui;
|
|
o := *segment_ui.overlay_state;
|
|
overlay := o.overlay;
|
|
if !o.show || !overlay || overlay.index != index || !o.texture.handle return;
|
|
draw := GetWindowDrawList();
|
|
draw.AddImage(draw, gpu_texture_ref(o.texture), image_min, image_max);
|
|
|
|
if !IsItemHovered() return; // the image
|
|
mouse := GetMousePos();
|
|
w, h := image_max.x - image_min.x, image_max.y - image_min.y;
|
|
if w <= 0 || h <= 0 return;
|
|
x := cast(s32) ((mouse.x - image_min.x) / w * overlay.width);
|
|
y := cast(s32) ((mouse.y - image_min.y) / h * overlay.height);
|
|
if x < 0 || y < 0 || x >= overlay.width || y >= overlay.height return;
|
|
id := overlay.map[cast(s64) y * overlay.width + x];
|
|
at := overlay.by_id[id];
|
|
if !id || at < 0 return;
|
|
segment := overlay.segments[at];
|
|
if array_find(o.hidden, segment.class) return;
|
|
SetTooltip("% score % % of the image", segment.class, FF(segment.score, 2), join(tprint("%", FF(segment.area * 100, 1)), "%",, temp));
|
|
}
|
|
|
|
// A class's colour (the same in every capture), a shade of it for its instance'th segment in a
|
|
// capture: the next ones darker, then round again.
|
|
segment_class_colour :: (class: string, instance: s64) -> [3] u8 {
|
|
base := segment_class_base_colour(class);
|
|
shade := 1 - 0.18 * cast(float) (instance % 4);
|
|
return .[cast(u8) (base[0] * shade), cast(u8) (base[1] * shade), cast(u8) (base[2] * shade)];
|
|
}
|
|
|
|
// The colour chosen for class in the legend; else its own, from a hash of its name.
|
|
segment_class_base_colour :: (class: string) -> [3] u8 {
|
|
for segment_ui.overlay_state.colours if it.class == class return it.rgb;
|
|
return segment_class_default_colour(class);
|
|
}
|
|
|
|
// A hue from a hash of the name, at a fixed saturation and brightness.
|
|
segment_class_default_colour :: (class: string) -> [3] u8 {
|
|
hash := fnv1a_hash(class.data, class.count);
|
|
hue := cast(float) (hash % 360) / 60.0;
|
|
saturation := 0.65;
|
|
value := 0.95;
|
|
sector := cast(s32) hue;
|
|
f := hue - sector;
|
|
p := value * (1 - saturation);
|
|
q := value * (1 - saturation * f);
|
|
t := value * (1 - saturation * (1 - f));
|
|
rgb: [3] float;
|
|
if sector == {
|
|
case 0; rgb = .[value, t, p];
|
|
case 1; rgb = .[q, value, p];
|
|
case 2; rgb = .[p, value, t];
|
|
case 3; rgb = .[p, q, value];
|
|
case 4; rgb = .[t, p, value];
|
|
case; rgb = .[value, p, q];
|
|
}
|
|
return .[cast(u8) (rgb[0] * 255), cast(u8) (rgb[1] * 255), cast(u8) (rgb[2] * 255)];
|
|
}
|
|
|
|
// The overlay as RGBA pixels at the mask's size (temp): each pixel in its segment's colour at the
|
|
// opacity, its edges darker and stronger; hidden classes and no segment left clear.
|
|
segment_overlay_pixels :: (overlay: *Segment_Overlay, hidden: [] string, opacity: float) -> [] u8 {
|
|
w, h := overlay.width, overlay.height;
|
|
fill, edge: [256] [4] u8; // by segment id
|
|
alpha := cast(u8) clamp(opacity * 255, 0, 255);
|
|
edge_alpha := cast(u8) clamp(opacity * 255 * 1.6 + 40, 0, 255);
|
|
for segment, n: overlay.segments {
|
|
if array_find(hidden, segment.class) continue;
|
|
instance := 0;
|
|
for 0..n - 1 if overlay.segments[it].class == segment.class instance += 1;
|
|
c := segment_class_colour(segment.class, instance);
|
|
fill[segment.id] = .[c[0], c[1], c[2], alpha];
|
|
edge[segment.id] = .[c[0] / 3, c[1] / 3, c[2] / 3, edge_alpha];
|
|
}
|
|
pixels := NewArray(cast(s64) w * h * 4, u8,, temp);
|
|
for y: 0..h - 1 {
|
|
row := cast(s64) y * w;
|
|
for x: 0..w - 1 {
|
|
i := row + x;
|
|
id := overlay.map[i];
|
|
if !id continue;
|
|
on_edge := (x > 0 && overlay.map[i - 1] != id) || (x + 1 < w && overlay.map[i + 1] != id) ||
|
|
(y > 0 && overlay.map[i - w] != id) || (y + 1 < h && overlay.map[i + w] != id);
|
|
colour := ifx on_edge then edge[id] else fill[id];
|
|
memcpy(pixels.data + i * 4, colour.data, 4);
|
|
}
|
|
}
|
|
return pixels;
|
|
}
|
|
|
|
#scope_file
|
|
|
|
// The texture, from the mask held.
|
|
segment_overlay_build :: () {
|
|
o := *segment_ui.overlay_state;
|
|
gpu_texture_release(*o.texture);
|
|
overlay := o.overlay;
|
|
if !overlay || !overlay.map.count return;
|
|
pixels := segment_overlay_pixels(overlay, o.hidden, o.opacity);
|
|
o.texture = gpu_texture_create(pixels.data, overlay.width, overlay.height);
|
|
}
|
|
|
|
#scope_export
|
|
|
|
// The confirmation for erasing one folder's segments, as a modal popup; at the top level, each frame.
|
|
segment_erase_popup :: () { using ImGui;
|
|
POPUP :: "Erase segments?";
|
|
if segment_ui.erase_asked && !IsPopupOpen(POPUP) OpenPopup(POPUP);
|
|
if !BeginPopupModal(POPUP, null, .AlwaysAutoResize) return;
|
|
folder := segment_ui.erase_folder;
|
|
Text("Erase the segments of every image in %?", ifx folder then folder else "the library's top");
|
|
TextColored(app_theme.palette.dim, "Each file holding them is rewritten without them; its images, text, OCR and labels stay as they are.");
|
|
TextColored(app_theme.palette.warning, "Only segmenting the folder again brings them back.");
|
|
Spacing();
|
|
if Button("Erase") {
|
|
segment_ui.erase_asked = false;
|
|
CloseCurrentPopup();
|
|
segment_ui.window_open = true; // the progress bar is in the OCR window, opened by the erase
|
|
ocr_ui.window_open = true;
|
|
ocr_erase_start(.SEGMENTS, folder = folder, one_folder = true);
|
|
}
|
|
SameLine();
|
|
if Button("Cancel") || IsKeyPressed(.Escape) {
|
|
segment_ui.erase_asked = false;
|
|
CloseCurrentPopup();
|
|
}
|
|
EndPopup();
|
|
}
|