diff --git a/README.md b/README.md index 59688a4..d6be635 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,29 @@ grid (in its order, so through a search's matches). reading the images again, so changing the list is cheap. - **Erase labels and vectors from the library** (for testing), after asking. +### The Segmentation window + +**Segments** (above the grid) finds what is where in images: people, animals and things, each one +apart, and stuff such as sky, wall and grass (Mask2Former, `modules/Segment`, 133 COCO classes). It +works only in the library folders you choose, and never on videos. + +- **Library folders:** a table of every folder in the library: + - **Segment:** whether scans read that folder's images. It starts off for every folder, so adding a + large archive never starts a long GPU job by itself. + - **Segmented:** how many of its images hold segments. + - **Erase:** takes that folder's segments out of its files, after asking. + + The choice is kept in the library (`.library-folders`), not in the app's settings, so it goes + where the library goes. +- **Segment the folders set to segment** reads every image in them not yet segmented. Each is stored + as it is done, so a scan can stop and start again. It runs on the OCR thread, with the same + progress bar, Pause and Cancel. The model loads when the first scan starts and takes about 2.4 GB + of GPU memory; on an RTX 5080 laptop a scan does about 17 images a second. +- **Importing** into a folder set to segment segments the new images once they are in. The Import + window's "Segment the images in this folder" checkbox is the same setting. +- **Storage:** each image's segments are stored in its file: a short list (class, score, area, box) + for search, and a mask of about 5 KB for drawing. See `modules/Capture_File/README.md`. + ### The Stats window What OCR found in the captures shown (the whole library, or a search's matches), by the chosen OCR @@ -226,7 +249,8 @@ then in that order instead of by date. | `-search QUERY` | Starts with QUERY in the search box. | | `-scan` | Scans the library (the chosen OCR model, and labels when the label model is up) once it has loaded, then exits: 0, or 1 if any capture failed. | | `-stats` | Logs the Stats window's numbers once the catalog has loaded, then exits; with `-search`, for the matches, and the first ten are logged. | -| `-erase ocr`, `-erase labels`, `-erase all` | Erases that from every file in the library, then exits (for testing). | +| `-segment` | Segments the images not yet segmented in the library folders set to segment (see the Segmentation window), once the library has loaded, then exits: 0, or 1 if any image failed. | +| `-erase ocr`, `-erase labels`, `-erase segments`, `-erase all` | Erases that from every file in the library, then exits (for testing). | | `-allow-multiple-instances` | Runs beside another copy instead of bringing that one forward. | | `run_tests` | Runs the test suites (debug builds); with `-headless`, exits with 0 if all passed. | | `-headless` | No message boxes: a failure is logged and the app exits with 1. | @@ -245,6 +269,7 @@ scripts and for trying things on a copy of the library: `bin\mcap.exe -library D | Logs | `logs\\` beside the program | | The hash of every capture's image, for finding duplicates | `.import-index` in the library folder | | The folder each library folder was imported from | `.import-sources` in the library folder | +| Which library folders are segmented | `.library-folders` in the library folder | Nothing is written to the folders imported from (`D:\Snagit`, a phone's backup): the library is a converted copy. @@ -276,6 +301,9 @@ converted copy. image_probe`) shows what it makes of files. - `modules/Snagx`: `.snagx` to `.mcap`, copied from SnagitConverter. - `modules/Capture_File`: reading and writing `.mcap` files, with a video's original left on disk. +- `modules/Segment`: panoptic segmentation with Mask2Former, run with onnxruntime (see its README and + `segmentation-plan.html`). `src/app/segment_ui.jai` is the scan and the Segmentation window; + `src/app/library_folders.jai` the per-folder options. - `modules/ffmpeg`: FFmpeg 6.0's bindings and DLLs (`modules/ffmpeg/windows`, in git; copied beside the executable by the build), from the gyan.dev shared build, for reading videos. diff --git a/check_targets.jai b/check_targets.jai index 8e6ffc4..1aeb110 100644 --- a/check_targets.jai +++ b/check_targets.jai @@ -29,7 +29,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; GIT_BRANCH :: "check"; GIT_HASH :: "-"; -VERSION :: "0.0.12"; +VERSION :: "0.0.13"; BUILD_DATE :: "-"; DEBUG :: true; COMPILE_TESTS :: true; diff --git a/first.jai b/first.jai index 7b4dc3a..6b31404 100644 --- a/first.jai +++ b/first.jai @@ -8,7 +8,7 @@ APP_NAME :: "mcap"; APP_FOLDER_NAME :: "mcap"; APP_BUNDLE_IDENTIFIER :: "com.example.mcap"; -VERSION :: "0.0.12"; +VERSION :: "0.0.13"; // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated // builds never write over the ones you run by hand. build_directory := "bin"; diff --git a/modules/Capture_File/ocr_text.jai b/modules/Capture_File/ocr_text.jai index 5390e39..ca66e76 100644 --- a/modules/Capture_File/ocr_text.jai +++ b/modules/Capture_File/ocr_text.jai @@ -129,7 +129,7 @@ mcap_ocr_now :: () -> string { return tprint("%-%-%T%:%:%Z", t.year, two(t.month_starting_at_0 + 1), two(t.day_of_month_starting_at_0 + 1), two(t.hour), two(t.minute), two(t.second)); } -#scope_module // the quoting is shared with label_chunks.jai +// The quoting the text formats share (OCR text, labels, segments), and the app's own text files (.library-folders). // Quoted, a quote written twice, a control character as a space. Temp. mcap_quote :: (text: string) -> string { @@ -163,6 +163,8 @@ mcap_read_quoted :: (s: string) -> text: string, rest: string, ok: bool { return "", s, false; } +#scope_module + // Version 2: x y width height x0 y0 x1 y1 x2 y2 x3 y3 confidence "text". Version 1: x y width height confidence "text". ocr_parse_line :: (value: string, version: s64) -> Mcap_Ocr_Line, bool { result: Mcap_Ocr_Line; diff --git a/modules/Segment/module.jai b/modules/Segment/module.jai index 78793e0..11aa4b3 100644 --- a/modules/Segment/module.jai +++ b/modules/Segment/module.jai @@ -264,6 +264,36 @@ segment_result_free :: (result: *Segment_Result, allocator := context.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; @@ -570,4 +600,6 @@ segment_strip_comment :: (line: string) -> string { #import "Math"; #import "System"; #import "Hash"; +#import "Sort"; +#import "Capture_File"; // segment_mcap_chunks #import "onnxruntime"; diff --git a/segmentation-plan.html b/segmentation-plan.html index 686354c..cc183e4 100644 --- a/segmentation-plan.html +++ b/segmentation-plan.html @@ -78,7 +78,7 @@
  • Decided Images only, and only in the library folders you opt in. The setting is per folder, kept in a file in the library root, so it travels with the library.
  • Decided No .mcap version bump. Results go in two new chunk kinds: a small segs summary at the front for search, and a mask at the back for drawing.
  • It is built in six phases. The first proves the ONNX export matches PyTorch before anything else depends on it.
    • -
  • Progress Phases 1 to 3 are done (2026-10-10). The model is exported and checked, modules/Segment runs it in mcap at about 57 ms per image on the RTX 5080, and results are stored in .mcap files at about 5 KB per image. Next is phase 4: the per-folder opt-in and the scan.
    • +
  • Progress Phases 1 to 4 are done (2026-10-10). The model is exported and checked, and modules/Segment runs it in mcap at about 57 ms per image on the RTX 5080. Results are stored in .mcap files at about 5 KB per image. Library folders opt in through .library-folders, and the Segmentation window and -segment scan them. Next is phase 5: has: search and Stats.
  • Later Video, by treating sampled frames as images, is shelved.
    • @@ -367,7 +367,7 @@ then runs: (id, length) pairs over the whole map, row by row, wrapping a
    -

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

    +

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

    The pieces:

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

    +
    Result (2026-10-10). +
      +
    • The scratch library: three folders. Photos (16) and Videos (4 videos and 2 GIFs) are set to segment; Other (9 images) is not.
      • +
    • First -segment: it segmented exactly the 16 photos and 2 GIFs. The model loaded in 4.5 s and the scan took 6 s. Other and the videos were left alone.
      • +
    • Second -segment: it found nothing to do and exited 0.
      • +
    • -erase segments: it removed all 18.
      • +
    • Closing mid-scan: no crash, and exit took 0.3 to 1.8 s. Only whole engine batches were stored.
      • +
    • Tests: 14 new checks of .library-folders pass: quoted names, unknown keys kept, a later version refused.
      • +
    +

    Changes from the plan:

    +
      +
    • Cancel granularity. Cancelling between the engine's batches of 4 keeps any whole batch already stored, rather than storing nothing for the scan's batch of 8. Each stored image is complete; the counter counts those.
      • +
    • Stale results. A capture segmented with another class list is scanned again only once the engine is up. Until then a segs chunk counts as done.
      • +
    • The import follow-up runs whether or not "Read text and labels" is ticked, since the folder's setting is the opt-in. It waits for the new captures' catalog entries, so it knows which are videos. Headless -import does not segment.
      • +
    • Erase per folder uses the OCR window's erase task, filtered to the folder, and its progress shows there.
      • +
    +

    Not yet checked by eye: the Segmentation window's layout. It is drawn during the GUI -segment runs (scan and cancel), but no screenshot was taken.

    diff --git a/src/App.jai b/src/App.jai index 92d4ead..53c2626 100644 --- a/src/App.jai +++ b/src/App.jai @@ -24,6 +24,8 @@ USING_WINDOWS :: OS == .WINDOWS; #load "app/background.jai"; #load "app/labels_ui.jai"; #load "app/import.jai"; +#load "app/segment_ui.jai"; +#load "app/library_folders.jai"; #load "app/exit_shutdown.jai"; #load "base/circle_buffer.jai"; @@ -66,6 +68,7 @@ Shared_Channel :: #import "Shared_Memory_Channel"; #import,dir "../modules/jaison/unicode_utils"; // is_valid_utf8 #import "OCR"; #import "Labels"; +#import "Segment"; #import "Capture_Text"; #import "Base64"; #import "onnxruntime"; @@ -189,7 +192,8 @@ App :: struct { scan_and_exit: bool; // -scan: scan the library with the chosen OCR model once it has loaded, then exit (0, or 1 if any capture failed) stats_and_exit: bool; // -stats: log the Statistics window's numbers once the catalog has loaded, then exit start_search: string; // -search QUERY: the search box's text at startup - erase_and_exit: string; // -erase ocr|labels|all: erase that from the whole library once it has loaded, then exit + erase_and_exit: string; // -erase ocr|labels|segments|all: erase that from the whole library once it has loaded, then exit + segment_and_exit: bool; // -segment: segment the library folders set to segment once the catalog has loaded, then exit (0, or 1 if any image failed) import_source, import_dest: string; // -import SOURCE DEST: import SOURCE into the library folder DEST (app/import.jai), then exit last_window_restore_time: Apollo_Time; diff --git a/src/app/app_ui.jai b/src/app/app_ui.jai index 9372821..ce76b83 100644 --- a/src/app/app_ui.jai +++ b/src/app/app_ui.jai @@ -10,6 +10,7 @@ app_ui_startup :: () { import_startup(); ocr_ui_startup(); labels_ui_startup(); + segment_ui_startup(); // the library folders' options: after the viewer, which knows the library } // Every frame, before input is handed to ImGui. Poll background work here. @@ -26,6 +27,8 @@ app_ui_shutdown :: () { viewer_shutdown(); app_shutdown_step(.OCR); ocr_ui_shutdown(); + app_shutdown_step(.SEGMENT); + segment_ui_shutdown(); // after the OCR thread, which runs its scans app_shutdown_step(.LABELS); labels_ui_shutdown(); // after the OCR thread and the background workers have stopped app_shutdown_step(.STATS); @@ -33,7 +36,7 @@ app_ui_shutdown :: () { } // Which step of app_ui_shutdown is running, and since when, for the overlay a slow shutdown shows. -App_Shutdown_Step :: enum s32 { NONE; IMPORT; VIEWER; OCR; LABELS; STATS; } +App_Shutdown_Step :: enum s32 { NONE; IMPORT; VIEWER; OCR; SEGMENT; LABELS; STATS; } app_shutdown: struct { step: App_Shutdown_Step; step_started_ns: s64; // monotonic @@ -68,8 +71,10 @@ app_ui_shutdown_status :: () -> string { case .CREATE; text = "Waiting for the OCR model to finish loading."; case .LABELS_CREATE; text = "Waiting for the label model to finish loading."; case .LABEL_SET; text = "Waiting for the label list to finish."; + case .SEGMENT; text = "Stopping the segmentation scan: finishing the images it is segmenting."; case; text = "Stopping OCR."; } + case .SEGMENT; text = "Closing the segmentation model."; case .LABELS; text = "Closing the label model."; case .STATS; text = "Closing Stats."; } @@ -87,6 +92,7 @@ app_main_window_ui :: () { viewer_ui(); ocr_ui_frame(); labels_ui_frame(); + segment_ui_frame(); stats_ui_frame(); import_frame(); } diff --git a/src/app/catalog.jai b/src/app/catalog.jai index a9de18a..12c356b 100644 --- a/src/app/catalog.jai +++ b/src/app/catalog.jai @@ -20,6 +20,10 @@ Catalog_Entry :: struct { video_codec: string; // "h264", "hevc" ...; "" when not a video frame_rate: float64; + // The segmentation model's (SEGMENT_MODEL, app/segment_ui.jai) segs chunk, when it has one. + segment_classes: string; // its class list's hash; "" when not segmented + segment_count: s32; + // Lowercased copies, for matching without case. find_title, find_url, find_app, find_notes, find_ocr, find_labels, find_folder: string; } @@ -45,6 +49,7 @@ catalog_entry_free :: (entry: *Catalog_Entry) { array_free(entry.labels); free(entry.label_list); free(entry.video_codec); + free(entry.segment_classes); array_free(entry.embedding); free(entry.find_title); free(entry.find_url); free(entry.find_app); free(entry.find_notes); free(entry.find_ocr); free(entry.find_labels); free(entry.find_folder); free(entry); @@ -121,6 +126,14 @@ catalog_load :: (path: string) -> *Catalog_Entry, error: string { half, half_ok := mcap_embedding_parse(embedding_chunk.data); if half_ok entry.embedding = half; } + segments_chunk := mcap_find_chunk(front.front_chunks, mcap_tag(MCAP_SEGMENTS_TAG), SEGMENT_MODEL); + if segments_chunk { + segments, segments_ok := mcap_segments_parse(cast(string) segments_chunk.data,, temp); + if segments_ok { + entry.segment_classes = copy_string(segments.classes); + entry.segment_count = cast(s32) segments.segments.count; + } + } library := viewer_folder(); if path.count > library.count + 1 && begins_with_nocase(path, library) { diff --git a/src/app/import.jai b/src/app/import.jai index 3222f5c..04860f2 100644 --- a/src/app/import.jai +++ b/src/app/import.jai @@ -95,6 +95,7 @@ importer: struct { // Main thread. new_indices: [..] s32; // the captures this import added, for the scan after it scan_pending: bool; + segment_pending: bool; // the folder imported into is set to segment: its new images, after the scan flag_started: bool; // -import has started its import } @@ -918,11 +919,21 @@ import_frame :: () { using ImGui; if importer.error toast_show(.Error, "Import: %", importer.error); else toast_show(.Success, "Imported % into % (% duplicates, % skipped, % failed)", importer.imported, importer.dest_name, importer.duplicates, importer.skipped, importer.failed); importer.scan_pending = importer.read_after && importer.new_indices.count > 0 && !app.import_source; + importer.segment_pending = library_folder_segments(importer.dest_name) && importer.new_indices.count > 0 && !app.import_source; } if importer.scan_pending && !ocr_ui.thread_busy && ocr_ui.engine && labels_ui_settled() { importer.scan_pending = false; ocr_scan_start(importer.new_indices); } + // Then segmentation (app/segment_ui.jai), once the new captures' catalog entries say which are videos. + if importer.segment_pending && !importer.scan_pending && !ocr_ui.thread_busy && !ocr_ui.scan.running { + catalogued := true; + for importer.new_indices if !viewer.captures[it].catalog && !viewer.captures[it].removed { catalogued = false; break; } + if catalogued { + importer.segment_pending = false; + segment_scan_start(importer.new_indices, quiet = true); + } + } if !importer.window_open return; SetNextWindowSize(.{720, 520}, .FirstUseEver); @@ -1001,6 +1012,12 @@ import_window :: () { using ImGui; else TextColored(app_theme.palette.dim, "Into %/%", viewer_folder(), dest); } Checkbox("Read text and labels in what is imported, once it is in", *importer.read_after); + if dest && import_name_is_valid(dest) { + // The library folder's own option (app/library_folders.jai): it stays with the folder. + segment := library_folder_segments(dest); + if Checkbox("Segment the images in this folder (what is where: people, animals, things, sky...)", *segment) library_folder_set_segment(dest, segment); + if IsItemHovered() SetTooltip("Kept for the library folder: the Segmentation window has it too. Videos are never segmented."); + } if running EndDisabled(); Spacing(); diff --git a/src/app/library_folders.jai b/src/app/library_folders.jai new file mode 100644 index 0000000..431bdba --- /dev/null +++ b/src/app/library_folders.jai @@ -0,0 +1,146 @@ +// Options for each library folder (a top-level folder of the library, one per place imported from), +// kept in .library-folders in the library itself, so they go where the library goes: +// +// @version 1 +// +// [folder "Photos 2014"] +// segment: true +// +// segment: segmentation scans (app/segment_ui.jai) read the images in it. A folder not listed, or an +// option not given, is off. Keys this build does not know are kept as they are when the file is +// written again, so a later build's options survive an earlier one. Main thread only. + +LIBRARY_FOLDERS_NAME :: ".library-folders"; + +Library_Folder :: struct { + name: string; // as viewer captures name it: "SnagitLibraryWork"; "" for the library's top + segment: bool; + other_lines: [..] string; // key: value lines this build does not know, as they were +} + +library_folders: struct { + folders: [..] Library_Folder; // heap; in the file's order, new ones last + loaded: bool; +} + +library_folders_path :: () -> string { + return tprint("%/%", viewer_folder(), LIBRARY_FOLDERS_NAME); +} + +// Reads .library-folders; a missing file is every option off. Mistakes are logged and skipped. +library_folders_load :: () { + path := library_folders_path(); + text, ok := read_entire_file(path, log_errors = false,, temp); + library_folders_parse(ifx ok then text else "", path); +} + +// Writes .library-folders beside itself and swaps it in. +library_folders_save :: () { + path := library_folders_path(); + temp_path := tprint("%.tmp", path); + if !write_entire_file(temp_path, library_folders_text()) || !atomic_file_replace_or_just_rename_if_destination_file_does_not_exist(temp_path, path) { + log_error("[Library] Could not write %.", path); + } +} + +// The options in text, a .library-folders file's (path is for messages), in place of those held. +library_folders_parse :: (text: string, path: string) { + push_allocator(context.default_allocator); + library_folders_clear(); + library_folders.folders.allocator = context.default_allocator; + library_folders.loaded = true; + + folder: *Library_Folder; + version_seen := false; + for raw, line_index: split(text, "\n",, temp) { + line := trim(trim_right(raw, "\r")); + if !line.count || line[0] == #char "#" continue; + if begins_with(line, "@version") { + version, version_ok := string_to_int(trim(advance(line, 8))); + if !version_ok || version != 1 { log_error("[Library] % is version %, which this build does not read; its options are off.", path, trim(advance(line, 8))); library_folders_clear(); return; } + version_seen = true; + continue; + } + if !version_seen { log_error("[Library] % does not start with @version 1; its options are off.", path); return; } + if begins_with(line, "[folder ") && ends_with(line, "]") { + name, rest, name_ok := mcap_read_quoted(trim(slice(line, 8, line.count - 9))); + if !name_ok || trim(rest).count { log_error("[Library] %:%: expected [folder \"NAME\"].", path, line_index + 1); folder = null; continue; } + folder = library_folder_find_or_add(name); + continue; + } + colon := find_index_from_left(line, #char ":"); + if !folder || colon < 0 { log_error("[Library] %:%: expected key: value under a [folder \"NAME\"]. Ignored.", path, line_index + 1); continue; } + key := trim(slice(line, 0, colon)); + value := trim(advance(line, colon + 1)); + if key == "segment" && (value == "true" || value == "false") folder.segment = value == "true"; + else array_add(*folder.other_lines, copy_string(line)); + } +} + +// The options held, as a .library-folders file. Folders with no option on and nothing kept are left +// out. Temp. +library_folders_text :: () -> string { + builder: String_Builder; + builder.allocator = temp; + append(*builder, "@version 1\n"); + append(*builder, "# Options for each library folder, set in mcap's Segmentation window. See src/app/library_folders.jai.\n"); + for library_folders.folders { + if !it.segment && !it.other_lines.count continue; + print_to_builder(*builder, "\n[folder %]\n", mcap_quote(it.name)); + print_to_builder(*builder, "segment: %\n", ifx it.segment then "true" else "false"); + for line: it.other_lines print_to_builder(*builder, "%\n", line); + } + return builder_to_string(*builder,, temp); +} + +library_folder_segments :: (name: string) -> bool { + for library_folders.folders if it.name == name return it.segment; + return false; +} + +library_folder_set_segment :: (name: string, on: bool) { + folder := library_folder_find_or_add(name); + if folder.segment == on return; + folder.segment = on; + library_folders_save(); + log("[Library] Segmentation % for %.", ifx on then "on" else "off", ifx name then name else "the library's top"); +} + +// Folders the file has options for that the library no longer holds (renamed or deleted). Temp. +library_folders_gone :: () -> [..] string { + gone: [..] string; + gone.allocator = temp; + for library_folders.folders if (it.segment || it.other_lines.count) && !array_find(viewer.folders, it.name) array_add(*gone, it.name); + return gone; +} + +library_folder_forget :: (name: string) { + for library_folders.folders { + if it.name != name continue; + free(it.name,, context.default_allocator); + for line: it.other_lines free(line,, context.default_allocator); + array_free(it.other_lines); + array_ordered_remove_by_index(*library_folders.folders, it_index); + break; + } + library_folders_save(); +} + +#scope_file + +library_folder_find_or_add :: (name: string) -> *Library_Folder { + for * library_folders.folders if it.name == name return it; + folder := array_add(*library_folders.folders); + folder.* = .{name = copy_string(name,, context.default_allocator)}; + folder.other_lines.allocator = context.default_allocator; + return folder; +} + +library_folders_clear :: () { + for library_folders.folders { + free(it.name,, context.default_allocator); + for line: it.other_lines free(line,, context.default_allocator); + array_free(it.other_lines); + } + array_reset(*library_folders.folders); +} diff --git a/src/app/main.jai b/src/app/main.jai index 78af206..9108d3d 100644 --- a/src/app/main.jai +++ b/src/app/main.jai @@ -24,6 +24,7 @@ main :: () { while app.library_folder.count > 3 && app.library_folder[app.library_folder.count - 1] == #char "/" app.library_folder.count -= 1; } app.scan_and_exit = array_find(app.args, "-scan"); + app.segment_and_exit = array_find(app.args, "-segment"); app.stats_and_exit = array_find(app.args, "-stats"); for app.args if it == "-search" && it_index + 1 < app.args.count app.start_search = app.args[it_index + 1]; for app.args if it == "-erase" && it_index + 1 < app.args.count app.erase_and_exit = app.args[it_index + 1]; diff --git a/src/app/ocr_ui.jai b/src/app/ocr_ui.jai index e568250..37a07a8 100644 --- a/src/app/ocr_ui.jai +++ b/src/app/ocr_ui.jai @@ -49,7 +49,7 @@ ocr_ui: struct { toast_shown := false; } -Ocr_Task_Kind :: enum { NONE; CREATE; READ; SCAN; LABELS_CREATE; LABEL_SET; ERASE; } +Ocr_Task_Kind :: enum { NONE; CREATE; READ; SCAN; LABELS_CREATE; LABEL_SET; ERASE; SEGMENT; } Ocr_Task :: struct { kind: Ocr_Task_Kind; @@ -76,11 +76,12 @@ Ocr_Scan :: struct { pause_start: Apollo_Time; elapsed_seconds: float64; // once it has ended: the time spent scanning, pauses left out label_set: *Label_Set; // labels scored against this; null when the scan does not label - erase_ocr, erase_labels: bool; // an erase, not a scan: what it takes out of each file + erase_ocr, erase_labels, erase_segments: bool; // an erase, not a scan: what it takes out of each file + segmenting: bool; // a segmentation scan (app/segment_ui.jai), not an OCR one } // A capture to read, and what it needs. -Erase_What :: enum u8 { NONE; OCR; LABELS; ALL; } +Erase_What :: enum u8 { NONE; OCR; LABELS; SEGMENTS; ALL; } Scan_Item :: struct { index: s32; @@ -225,12 +226,17 @@ ocr_ui_thread_proc :: (thread: *Thread) -> s64 { ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx scan.cancel then ", cancelled" else ""); atomic_swap(*scan.running, false); + case .SEGMENT; + segment_ui_scan_task(); // app/segment_ui.jai + case .ERASE; scan := *ocr_ui.scan; // Not in temporary storage, which is reset after each capture. - label_tags := string.[MCAP_LABELS_TAG, MCAP_EMBEDDING_TAG]; - tags: [] string; - if scan.erase_labels tags = label_tags; + tag_list: [..] string; // in the default allocator: temporary storage is reset after each capture + defer array_free(tag_list); + if scan.erase_labels array_add(*tag_list, MCAP_LABELS_TAG, MCAP_EMBEDDING_TAG); + if scan.erase_segments array_add(*tag_list, MCAP_SEGMENTS_TAG, MCAP_MASK_TAG); + tags: [] string = tag_list; erased := 0; for item: scan.items { while scan.paused && !scan.cancel sleep_milliseconds(100); @@ -411,10 +417,11 @@ ocr_ui_erase_and_exit :: (busy: bool) { what: Erase_What; if app.erase_and_exit == { case "ocr"; what = .OCR; - case "labels"; what = .LABELS; - case "all"; what = .ALL; + case "labels"; what = .LABELS; + case "segments"; what = .SEGMENTS; + case "all"; what = .ALL; case; - log_error("[Erase] -erase takes ocr, labels or all, not %.", app.erase_and_exit); + log_error("[Erase] -erase takes ocr, labels, segments or all, not %.", app.erase_and_exit); app.exit_code = 2; app.should_exit = true; app.erase_and_exit = ""; return; } ocr_ui.scan_requested = true; @@ -431,14 +438,16 @@ ocr_ui_erase_and_exit :: (busy: bool) { // What an erase takes out, for people. Temp. erase_what_text :: (scan: *Ocr_Scan) -> string { - if scan.erase_ocr && scan.erase_labels return "OCR, labels and vectors"; + if scan.erase_ocr && scan.erase_labels && scan.erase_segments return "OCR, labels, vectors and segments"; if scan.erase_ocr return "OCR readings"; + if scan.erase_segments return "segments"; return "labels and vectors"; } -// Takes OCR readings, labels and vectors, or both, out of every capture in the library (not only -// those shown) that has any, on the OCR thread, as a scan runs: with its progress, pause and cancel. -ocr_erase_start :: (what: Erase_What) { +// Takes OCR readings, labels and vectors, segments, or all of them, out of every capture in the +// library (not only those shown), or with one_folder in that library folder only, that has any, on the +// OCR thread, as a scan runs: with its progress, pause and cancel. +ocr_erase_start :: (what: Erase_What, folder := "", one_folder := false) { if ocr_ui.thread_busy || what == .NONE return; scan := *ocr_ui.scan; for scan.items free(it.path); @@ -446,14 +455,16 @@ ocr_erase_start :: (what: Erase_What) { free(scan.model); scan.* = .{}; scan.items.allocator = context.default_allocator; - scan.erase_ocr = what == .OCR || what == .ALL; - scan.erase_labels = what == .LABELS || what == .ALL; + scan.erase_ocr = what == .OCR || what == .ALL; + scan.erase_labels = what == .LABELS || what == .ALL; + scan.erase_segments = what == .SEGMENTS || what == .ALL; for * viewer.captures { - if it.removed continue; + if it.removed || (one_folder && it.folder != folder) continue; entry := it.catalog; has := !entry // not known: looked at anyway || (scan.erase_ocr && entry.ocr.count > 0) - || (scan.erase_labels && (entry.labels.count > 0 || entry.embedding.count > 0 || entry.label_list.count > 0)); + || (scan.erase_labels && (entry.labels.count > 0 || entry.embedding.count > 0 || entry.label_list.count > 0)) + || (scan.erase_segments && entry.segment_classes.count > 0); if !has continue; array_add(*scan.items, .{index = cast(s32) it_index, path = copy_string(it.path,, context.default_allocator)}); } @@ -469,7 +480,7 @@ ocr_erase_popup :: () { using ImGui; POPUP :: "Erase from the library?"; if ocr_ui.erase_asked != .NONE && !IsPopupOpen(POPUP) OpenPopup(POPUP); if !BeginPopupModal(POPUP, null, .AlwaysAutoResize) return; - what := ifx ocr_ui.erase_asked == .OCR then "every OCR reading (all models)" else ifx ocr_ui.erase_asked == .LABELS then "every capture's labels and image vectors (all label models)" else "every OCR reading, label and image vector"; + what := ifx ocr_ui.erase_asked == .OCR then "every OCR reading (all models)" else ifx ocr_ui.erase_asked == .LABELS then "every capture's labels and image vectors (all label models)" else "every OCR reading, label, image vector and segment"; Text("Erase % from all % captures in %?", what, viewer_capture_count(), viewer_folder()); TextColored(app_theme.palette.dim, "Each file holding any is rewritten without them; its images, text and preview stay as they are."); TextColored(app_theme.palette.warning, "This cannot be undone: only scanning the library again brings them back."); @@ -573,7 +584,8 @@ ocr_scan_ui :: () { using ImGui; if !catalog_ready TextColored(app_theme.palette.dim, "(once the catalog has loaded)"); else TextColored(app_theme.palette.dim, "Reads every capture this model has not read, and labels those not yet labelled with this list. Each is stored as it is read, so a scan can stop and start again."); if scan.total && scan.done { - if scan.erase_ocr || scan.erase_labels TextColored(app_theme.palette.dim, "Last erase: % from % of % captures in %, % failed.", erase_what_text(scan), scan.done, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed); + if scan.erase_ocr || scan.erase_labels || scan.erase_segments TextColored(app_theme.palette.dim, "Last erase: % from % of % captures in %, % failed.", erase_what_text(scan), scan.done, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed); + else if scan.segmenting TextColored(app_theme.palette.dim, "Last scan: segmentation, % of % images in %, % failed (see the Segmentation window).", scan.done, scan.total, ocr_duration_text(scan.elapsed_seconds), scan.failed); else TextColored(app_theme.palette.dim, "Last scan: % of % read with % in %, % failed.", scan.done, scan.total, scan.model, ocr_duration_text(scan.elapsed_seconds), scan.failed); } if busy BeginDisabled(); @@ -588,11 +600,12 @@ ocr_scan_ui :: () { using ImGui; // The running scan's or erase's progress bar and buttons (the OCR and Labels windows show it). ocr_scan_progress_ui :: () { using ImGui; scan := *ocr_ui.scan; - if scan.erase_ocr || scan.erase_labels Text("Erasing %:", erase_what_text(scan)); + if scan.erase_ocr || scan.erase_labels || scan.erase_segments Text("Erasing %:", erase_what_text(scan)); done, total, failed, rate, left, elapsed := ocr_scan_progress(); ProgressBar(cast(float) done / max(total, 1), .{-FLOAT32_MIN, 0}, temp_c_string(ocr_scan_progress_text(done, total))); - if scan.erase_ocr || scan.erase_labels Text("% elapsed, % left, % captures/s, % failed", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed); + if scan.erase_ocr || scan.erase_labels || scan.erase_segments Text("% elapsed, % left, % captures/s, % failed", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed); + else if scan.segmenting Text("% elapsed, % left, % images/s, % failed", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed); // the engine is the OCR thread's while it runs else Text("% elapsed, % left, % captures/s, % failed, on %", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed, ifx ocr_ui.engine && ocr_ui.engine.provider == .CUDA then "the GPU" else "the CPU"); if scan.paused { if Button("Resume") { scan.paused_seconds += to_float64_seconds(current_time_monotonic() - scan.pause_start); atomic_swap(*scan.paused, false); } diff --git a/src/app/segment_ui.jai b/src/app/segment_ui.jai new file mode 100644 index 0000000..1cb8daa --- /dev/null +++ b/src/app/segment_ui.jai @@ -0,0 +1,376 @@ +// 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; +} + +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 +} + +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; + library_folders_load(); +} + +// After ocr_ui_shutdown has stopped the OCR thread. +segment_ui_shutdown :: () { + 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 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(); +} diff --git a/src/app/viewer.jai b/src/app/viewer.jai index 36581cf..96212a9 100644 --- a/src/app/viewer.jai +++ b/src/app/viewer.jai @@ -863,6 +863,9 @@ viewer_ui :: () { using ImGui; SameLine(); if Button("Labels") labels_ui.window_open = !labels_ui.window_open; // app/labels_ui.jai SameLine(); + if Button("Segments") segment_ui.window_open = !segment_ui.window_open; // app/segment_ui.jai + if IsItemHovered() SetTooltip("Segmentation: what is where in the images of the library folders you choose"); + SameLine(); if Button("Import") importer.window_open = !importer.window_open; // app/import.jai if IsItemHovered() SetTooltip("Bring in a folder of screenshots and photos"); SameLine(); @@ -883,7 +886,8 @@ viewer_ui :: () { using ImGui; done, total, failed, rate, left, elapsed := ocr_scan_progress(); ProgressBar(cast(float) done / max(total, 1), .{-FLOAT32_MIN, max(2.0, GetFontSize() * 0.15)}, ""); scan := *ocr_ui.scan; - doing := ifx scan.erase_ocr || scan.erase_labels then tprint("Erasing %", erase_what_text(scan)) else tprint("OCR scan with %", scan.model); + doing := ifx scan.erase_ocr || scan.erase_labels || scan.erase_segments then tprint("Erasing %", erase_what_text(scan)) + else ifx scan.segmenting then tprint("Segmentation scan with %", scan.model) else tprint("OCR scan with %", scan.model); if IsItemHovered() SetTooltip("%: %, % elapsed, % left, % failed%", doing, ocr_scan_progress_text(done, total), ocr_duration_text(elapsed), ocr_duration_text(left), failed, ifx scan.paused then " (paused)" else ""); } diff --git a/src/test/import_tests.jai b/src/test/import_tests.jai index 7bf1410..b259cd3 100644 --- a/src/test/import_tests.jai +++ b/src/test/import_tests.jai @@ -7,7 +7,7 @@ checks_run := 0; checks_failed := 0; -MAX_GROUPS :: 8; +MAX_GROUPS :: 12; group_names : [MAX_GROUPS] string; group_checks : [MAX_GROUPS] int; group_fails : [MAX_GROUPS] int; @@ -474,6 +474,53 @@ segment_chunk_tests :: () { file_delete(path); } +// .library-folders (app/library_folders.jai): each folder's options, read and written back with what +// this build does not know kept. +library_folders_tests :: () { + group("library folders"); + TEXT :: #string END +@version 1 +# a comment + +[folder "Photos 2014"] +segment: true +ocr: false + +[folder "My ""quoted"" folder"] +segment: true + +[folder "Screens"] +segment: false + +[folder ""] +segment: true +END; + library_folders_parse(TEXT, "test"); + expect(library_folder_segments("Photos 2014"), "a folder set to segment"); + expect(library_folder_segments("My \"quoted\" folder"), "a quoted name"); + expect(!library_folder_segments("Screens"), "a folder set not to"); + expect(library_folder_segments(""), "the library's top"); + expect(!library_folder_segments("Not listed"), "a folder not listed is off"); + + // Written back: the unknown key kept, under its folder; a folder with nothing on left out. + text := library_folders_text(); + expect(contains(text, "[folder \"My \"\"quoted\"\" folder\"]"), "the quoted name written back quoted"); + expect(contains(text, "[folder \"Photos 2014\"]\nsegment: true\nocr: false\n"), "a key this build does not know kept: %", text); + expect(!contains(text, "Screens"), "a folder with no option on is left out"); + library_folders_parse(text, "test"); + expect(library_folder_segments("Photos 2014") && library_folder_segments("My \"quoted\" folder") && library_folder_segments(""), "it reads back the same"); + expect(library_folders_text() == text, "and writes the same again"); + + library_folders_parse("@version 2\n[folder \"Photos\"]\nsegment: true\n", "test"); + expect(!library_folder_segments("Photos"), "a later version is not read: every option off"); + library_folders_parse("[folder \"Photos\"]\nsegment: true\n", "test"); + expect(!library_folder_segments("Photos"), "no @version: every option off"); + library_folders_parse("@version 1\nsegment: true\n[folder \"Photos\"]\nsegment: maybe\n", "test"); + expect(!library_folder_segments("Photos"), "a key outside a folder, or a value not true or false, sets nothing"); + library_folders_parse("", "test"); + expect(!library_folders.folders.count, "no file: no options"); +} + // What mcap_chunks_size counts, here where it cannot be seen. mcap_chunks_size_for_test :: (chunks: [] Mcap_Chunk) -> s64 { size: s64 = 4; @@ -495,5 +542,6 @@ run_import_tests :: (builder: *String_Builder = null) -> bool { video_tests(); mcap_tests(); segment_chunk_tests(); + library_folders_tests(); return report(); } diff --git a/tools/segment_probe.jai b/tools/segment_probe.jai index ee9cf7f..ae599b0 100644 --- a/tools/segment_probe.jai +++ b/tools/segment_probe.jai @@ -109,7 +109,7 @@ main :: () { if out_folder write_overlay(engine, image, result, tprint("%/%.segments.png", out_folder, path_strip_extension(path_filename(path)))); if store && ends_with_nocase(path, ".mcap") { - segs, mask := segment_chunks(engine, result, model_name, whiches[n], image.width, image.height); + segs, mask := segment_mcap_chunks(engine, result, whiches[n], image.width, image.height); stored, store_error := mcap_store(path, front = .[segs], back = .[mask]); if stored print(" stored: % bytes of segs, % bytes of mask\n", segs.data.count, mask.data.count); else { print(" cannot store: %\n", store_error); failures += 1; } @@ -144,29 +144,6 @@ main :: () { exit(xx ifx failures then 1 else 0); } -// A result as the chunks an .mcap stores (modules/Capture_File/segment_chunks.jai), for an image of -// image_width x image_height: boxes in its pixels, segments largest first. Temp. -segment_chunks :: (engine: *Segment_Engine, result: Segment_Result, model: string, 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 = tprint("%", formatInt(engine.classes_hash, base = 16, minimum_digits = 16)); - 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; }); - return mcap_segments_chunk(model, s), mcap_mask_chunk(model, result.width, result.height, result.map); -} - CONFIDENT :: 0.6; // a segment PyTorch scores this or more must be matched Comparison :: struct { @@ -286,6 +263,4 @@ pad :: (text: string, width: s64) -> string { #import "Image_Decode"; #import "Segment"; #import "Capture_File"; -#import "onnxruntime"; -#import "Sort"; #import "stb_image_write";