diff --git a/.gitignore b/.gitignore index a48abf7..c6b80c5 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,5 @@ Thumbs.db /modules/onnxruntime/macos/ /modules/onnxruntime/download/ /modules/OCR/models/ +# Labels: downloaded, not committed (modules/Labels/fetch_models.jai) +/modules/Labels/models/ diff --git a/README.md b/README.md index c7033bc..b520be7 100644 --- a/README.md +++ b/README.md @@ -31,12 +31,14 @@ and everything around them, ready for an app to be written into `src/app/app_ui. ## The capture viewer -`src/app/viewer.jai`, `catalog.jai` and `ocr_ui.jai`: every `.mcap` in `D:\SnagitIndex\mcap` as a grid -of previews, and one capture at full size in the Capture window. +`src/app/viewer.jai`, `catalog.jai`, `ocr_ui.jai` and `labels_ui.jai`: every `.mcap` in +`D:\SnagitIndex\mcap` as a grid of previews, and one capture at full size in the Capture window. - **Search** above the grid: every word must match, anywhere in the OCR text, window title, URL, app - name or annotation text; `"a phrase"` matches as one; `ocr:`, `title:`, `url:`, `app:` and `note:` - limit a word to one field. In the Capture window the matching OCR lines are highlighted, with Next + name, annotation text or labels; `"a phrase"` matches as one; `ocr:`, `title:`, `url:`, `app:`, + `note:` and `label:` limit a word to one field. `like:"a circuit board"` searches by meaning: the + captures whose image is nearest the phrase, best first (the label model's text encoder makes the + phrase's vector; each labelled capture's image vector is in its file). In the Capture window the matching OCR lines are highlighted, with Next and Previous. Matching runs on a pool of background threads (`app/background.jai`), as do the Stats window's numbers, so typing never waits on them; the grid shows the last results until the new ones are back, usually within a frame. @@ -45,7 +47,9 @@ of previews, and one capture at full size in the Capture window. - **OCR** (`modules/OCR`): the OCR window shows whether it runs on the GPU and, if not, what to install, and switches model (kept as `model` under `[ocr]` in the settings file). **Scan the library** reads every capture shown that the model has not read yet, storing each reading in its file as it goes, - so it can pause, cancel and pick up again. `-library FOLDER` shows another folder of `.mcap` files; + so it can pause, cancel and pick up again; with **Labels too** it labels them in the same pass. The + Capture window's **OCR** button reads one capture; a capture read by several models keeps each + reading, and the window can show any of them. `-library FOLDER` shows another folder of `.mcap` files; `-scan` scans it with the chosen model once it has loaded and exits (0, or 1 if any capture failed), for a scan without the UI, such as on a copy. - **Delete:** hovering a thumbnail shows a trash button in its corner. The confirmation offers moving @@ -59,8 +63,17 @@ of previews, and one capture at full size in the Capture window. - **Stats** above the grid: what OCR found in the captures shown, by the chosen model: captures with text, lines, words, characters, confidence, words per capture, a year table, the captures with the most words (click to open) and the commonest words (click to search). `-stats` logs the same - numbers once the catalog has loaded, and exits. The Capture window's **OCR** button reads one capture; - a capture read by several models keeps each reading, and the window can show any of them. + numbers once the catalog has loaded, and exits. `-search QUERY` starts with QUERY in the search box + (with `-stats`, the numbers are for its matches, and the first ten are logged). Stats also counts + each capture's best label (click one to search for it). +- **Labels** (`modules/Labels`, SigLIP 2 base at 384 px): each capture is scored against a label list, + `labels.config` in the app's data folder (`%APPDATA%\mcap`): a name and a prompt per label, a fuller + description scoring better than a word. A starter list is written the first time; edits are + picked up while the app runs. A scan stores each capture's five best labels and its image vector in + its file (`labl` and `embd` chunks, see `modules/Capture_File/README.md`), so a changed list is + scored again from the stored vectors without reading the images. The Labels window shows the list + and how much of the library is labelled; the Capture window shows the capture's labels. `label:` + matches a capture's best label, and any scoring at least half as much. ## Building diff --git a/check_targets.jai b/check_targets.jai index 213fc80..3891cb4 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.4"; +VERSION :: "0.0.5"; BUILD_DATE :: "-"; DEBUG :: true; COMPILE_TESTS :: true; diff --git a/first.jai b/first.jai index b216715..5421613 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.4"; +VERSION :: "0.0.5"; // 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"; @@ -283,6 +283,15 @@ copy_assets :: () { } else { log("OCR: the models have not been fetched; run jai fetch_models.jai in modules/OCR for OCR."); } + + // Labels: its models in bin/label-models, the same way. + make_directory_if_it_does_not_exist(tprint("%/label-models", build_directory)); + copy_if_changed("modules/Labels/models.config", tprint("%/label-models/models.config", build_directory)); + if file_exists("modules/Labels/models") { + for file_list("modules/Labels/models") copy_if_changed(it, tprint("%/label-models/%", build_directory, path_filename(it))); + } else { + log("Labels: the models have not been fetched; run jai fetch_models.jai in modules/Labels for labels."); + } } copy_if_changed :: (from: string, to: string) { diff --git a/modules/Capture_File/README.md b/modules/Capture_File/README.md index 634e0dc..45e87e5 100644 --- a/modules/Capture_File/README.md +++ b/modules/Capture_File/README.md @@ -1,36 +1,44 @@ # The .mcap file One capture in one file: a fixed magic, a file version stamp, the capture's data as -[Capture Text](capture_text.md), what OCR read in it, a small preview, then the images. There is no ZIP -and no compression around it; the images keep their own PNG compression. +[Capture Text](capture_text.md), what OCR read in it, what other models added (chunks), a small preview, +then the images, then any large chunks. There is no ZIP and no compression around it; the images keep +their own PNG compression. -The parts a library browser needs (header, text, OCR and preview) come first, so it reads the start of -each file and never touches the images; a search, which needs no preview, reads less still. +The parts a library browser needs (header, text, OCR, front chunks and preview) come first, so it reads +the start of each file and never touches the images; a search, which needs no preview, reads less still. -## Layout, version 2 +## Layout, version 3 All integers are little-endian. Sizes count bytes. | Field | Type | | | --- | --- | --- | | magic | 8 bytes | `89 4D 43 41 50 0D 0A 1A`, which is `\x89MCAP\r\n\x1A`. Says this is an .mcap, in every version. | -| file version | u32 | `2`. Version 1 files are still read; see below. | +| file version | u32 | `3`. Versions 1 and 2 are still read; see below. | | text size, text | u32, bytes | Capture Text, UTF-8: every JSON file of the capture, as `[[file]]` sections. | | OCR size, OCR | u32, bytes | OCR text, UTF-8 (below): what each OCR model read. Size 0 when none has. | +| front chunks | chunk list | Small things models added, such as labels and an image vector (below). | | preview width, height | u16, u16 | In pixels. | | preview size, preview | u32, bytes | PNG. Size 0 when there is none. | | original name size, name | u16, bytes | The original image's name in the .snagx, such as `{GUID}.png`; the page's `CaptureBackgroundImage` names it too. | | original size, original | u32, bytes | The capture as taken: PNG, exactly as it was in the .snagx. | | rendered kind | u8 | 0: none. 1: the same as the original (stored once). 2: its own, below. | | rendered size, rendered | u32, bytes | Only for kind 2. The capture with its annotations drawn in (Snagit's `thumbnail.png`): PNG, exactly as it was. | +| back chunks | chunk list | Large things models added, such as segmentation masks: after the images, so reading the front never reaches them. | The file ends there: anything after the last part is a mistake. The magic borrows PNG's: the first byte has its high bit set, so a channel that strips it shows; `\r\n` shows line endings being converted; and `\x1A` stops `type` on Windows. -**Version 1** is the same without the OCR size and OCR. A reader takes both; anything written is -version 2, so a version 1 file becomes version 2 the first time OCR is stored in it. +A **chunk list** is a u32 count, then each chunk: a 4-byte tag of ASCII letters saying what it is +(`labl`, `embd`), a u32 size, and that many bytes: a u16 key size, the key (whose it is: a model's +name), and the chunk's own data. A file holds at most one chunk per tag and key. + +**Version 2** is version 3 without the two chunk lists; **version 1** is version 2 without the OCR size +and OCR. A reader takes all three; anything written is version 3, so an older file becomes version 3 +the first time something is stored in it. ## The parts @@ -38,6 +46,10 @@ version 2, so a version 1 file becomes version 2 the first time OCR is stored in and `{GUID}.backup.json` when there is one. - **OCR.** One section per OCR model, so a capture read by several models keeps every reading; reading it again with a model replaces only that model's section. See *OCR text* below. +- **Chunks.** What other models add, each keyed by the model's name, so models do not overwrite each + other; storing a chunk again replaces only the one with its tag and key. A chunk a reader does not + know is kept as it is when the file is rewritten, so a later build can add kinds without a new file + version. See *Label chunks* below. - **Preview.** The rendered image (or the original, when there is no rendered one) scaled so its longest side is at most 320 pixels, as a PNG. An image already that small is used as it is, never scaled up. - **Original and rendered.** A capture with no annotations has a rendered image identical to the original, @@ -91,3 +103,32 @@ A line that cannot be read is skipped, as are keys this version does not know. * had no corners (`line: x y width height confidence "text"`); it is still read, with the box's corners. The OCR text carries its own version, so this did not change the .mcap file version. `mcap_store_ocr` (`modules/Capture_File/ocr_text.jai`) stores a section in a file. + +## Label chunks + +A label model (an image-text encoder such as SigLIP 2) adds two front chunks, keyed by its name. + +**`labl`**: the capture's best labels from a label list, as text in the style above: + +``` +@version 1 +engine: "onnxruntime 1.30.0, CUDA" +date: 2026-10-08T21:04:11Z +list: 9c4f0b2d81e3a675 +label: 0.6720 "Settings" +label: 0.0951 "Installer" +``` + +| Key | Value | +| --- | --- | +| `engine`, `date` | As in OCR text. | +| `list` | Which label list they were scored against: a 64-bit FNV-1a hash of its names and prompts. Labels from another list are out of date. | +| `label` | A score, 0 to 1 (the model's own sigmoid; small, but their order is what counts), and the label's name, quoted with `""` for a quote. Best first. | + +**`embd`**: the image's vector, binary: a u32 length, a u8 format (1: 16-bit floats), then the values, +little-endian. It is a unit vector, so any phrase's vector (from the same model's text encoder) scores +against it with a dot product: labels from a new list, or a search by meaning, need no new pass over +the images. 768 values, 1.5 KB, for SigLIP 2 base. + +`mcap_store` (`modules/Capture_File/module.jai`) stores an OCR section and chunks in one rewrite; +`label_chunks.jai` reads and writes these two. diff --git a/modules/Capture_File/label_chunks.jai b/modules/Capture_File/label_chunks.jai new file mode 100644 index 0000000..19aa479 --- /dev/null +++ b/modules/Capture_File/label_chunks.jai @@ -0,0 +1,111 @@ +// The chunks a label model adds to a capture (.mcap version 3), each keyed by the model's name: +// +// "labl", in front: its labels, as text in the style of the Ex4 text formats +// +// @version 1 +// engine: "onnxruntime 1.30.0, CUDA" +// date: 2026-10-08T21:04:11Z +// list: 9c4f0b2d81e3a675 +// label: 0.0647 "Terminal" +// label: 0.0234 "Schematic" +// +// list: says which label list they were scored against (a hash of its names and prompts), so labels from an +// older list can be told apart. label: a score (0 to 1, the model's sigmoid) and a name, quoted +// with "" for a quote; best first. A line that cannot be read is skipped, as are unknown keys. +// +// "embd", in front: the image's vector, binary: u32 length, u8 format (1: 16-bit floats), then the +// values, little-endian. A unit vector, so a phrase's vector scores against it with a dot product. + +MCAP_LABELS_TAG :: "labl"; +MCAP_EMBEDDING_TAG :: "embd"; +MCAP_LABELS_TEXT_VERSION :: 1; +MCAP_EMBEDDING_HALF :: 1; + +Mcap_Label :: struct { + score: float; + name: string; +} + +Mcap_Labels :: struct { + engine: string; // what ran it, such as "onnxruntime 1.30.0, CUDA" + date: string; // when, UTC: 2026-10-08T21:04:11Z + list: string; // the label list's hash + labels: [..] Mcap_Label; // best first +} + +// A labl chunk's text. Temp. +mcap_labels_text :: (labels: Mcap_Labels) -> string { + builder: String_Builder; + builder.allocator = temp; + print_to_builder(*builder, "@version %\n", MCAP_LABELS_TEXT_VERSION); + print_to_builder(*builder, "engine: %\n", mcap_quote(labels.engine)); + print_to_builder(*builder, "date: %\n", labels.date); + print_to_builder(*builder, "list: %\n", labels.list); + for labels.labels print_to_builder(*builder, "label: % %\n", formatFloat(it.score, trailing_width = 4, zero_removal = .NO), mcap_quote(it.name)); + return builder_to_string(*builder,, temp); +} + +// A labl chunk's labels; strings and arrays from the context allocator. +mcap_labels_parse :: (text: string) -> Mcap_Labels, ok: bool, error: string { + result: Mcap_Labels; + version_seen := false; + for raw: split(text, "\n",, temp) { + line := trim(raw); + if !line.count || line[0] == #char "#" continue; + if begins_with(line, "@version") { + rest := trim(advance(line, 8)); + version, ok := string_to_int(rest); + if !ok || version < 1 || version > MCAP_LABELS_TEXT_VERSION return result, false, tprint("its labels are version %, but this build reads up to %", rest, MCAP_LABELS_TEXT_VERSION); + version_seen = true; + continue; + } + if !version_seen return result, false, "its labels do not start with @version"; + colon := find_index_from_left(line, #char ":"); + if colon < 0 continue; + key := trim(slice(line, 0, colon)); + value := trim(advance(line, colon + 1)); + if key == { + case "engine"; + quoted, _, ok := mcap_read_quoted(value); + if ok result.engine = copy_string(quoted); + case "date"; result.date = copy_string(value); + case "list"; result.list = copy_string(value); + case "label"; + space := find_index_from_left(value, #char " "); + if space < 0 continue; + score, score_ok := string_to_float(slice(value, 0, space)); + name, after, name_ok := mcap_read_quoted(trim_left(advance(value, space))); + if score_ok && name_ok && !trim(after).count array_add(*result.labels, .{score, copy_string(name)}); + } + } + return result, true, ""; +} + +// An embd chunk's data for a vector stored as 16-bit floats. Temp. +mcap_embedding_bytes :: (half: [] u16) -> [] u8 { + bytes := NewArray(5 + half.count * 2, u8,, temp); + (cast(*u32) bytes.data).* = cast(u32) half.count; + bytes[4] = MCAP_EMBEDDING_HALF; + memcpy(bytes.data + 5, half.data, half.count * 2); + return bytes; +} + +// An embd chunk's vector, as 16-bit floats; a copy from the context allocator. +mcap_embedding_parse :: (data: [] u8) -> [] u16, ok: bool, error: string { + if data.count < 5 return .[], false, "its embedding is cut short"; + count := cast(s64) (cast(*u32) data.data).*; + if data[4] != MCAP_EMBEDDING_HALF return .[], false, tprint("its embedding is in format %, which this build does not read", data[4]); + if data.count != 5 + count * 2 return .[], false, "its embedding's length does not match its size"; + half := NewArray(count, u16); + memcpy(half.data, data.data + 5, count * 2); // the data need not be aligned + return half, true, ""; +} + +// The chunk for labels by model, ready to store with mcap_store_chunk. Its data is temp. +mcap_labels_chunk :: (model: string, labels: Mcap_Labels) -> Mcap_Chunk { + return .{tag = mcap_tag(MCAP_LABELS_TAG), key = model, data = cast([] u8) mcap_labels_text(labels)}; +} + +mcap_embedding_chunk :: (model: string, half: [] u16) -> Mcap_Chunk { + return .{tag = mcap_tag(MCAP_EMBEDDING_TAG), key = model, data = mcap_embedding_bytes(half)}; +} diff --git a/modules/Capture_File/module.jai b/modules/Capture_File/module.jai index 9de3899..5292924 100644 --- a/modules/Capture_File/module.jai +++ b/modules/Capture_File/module.jai @@ -1,5 +1,10 @@ // .mcap: one capture in one file. A fixed magic, a file version stamp, the capture's Capture Text, -// its OCR text (version 2), a small preview, then the images. See docs/mcap.md for the layout. +// its OCR text (version 2), its front chunks (version 3), a small preview, the images, then its back +// chunks (version 3). See docs/mcap.md for the layout. +// +// Chunks hold what models add (labels, embeddings, later masks), each named by a tag and a key (the +// model's name): small ones in front, where a library browser reads them without the images, large +// ones at the back. Chunks this build does not know are kept as they are when a file is rewritten. // // Reading gives views into the bytes it was given; nothing is copied. Images stay PNG. // @@ -7,7 +12,7 @@ // the same files. MCAP_MAGIC :: u8.[0x89, #char "M", #char "C", #char "A", #char "P", 0x0D, 0x0A, 0x1A]; -MCAP_VERSION :: 2; // what is written; every version from 1 up is read +MCAP_VERSION :: 3; // what is written; every version from 1 up is read MCAP_EXTENSION :: "mcap"; MCAP_PREVIEW_MAX_SIDE :: 320; // the preview's longest side, in pixels @@ -23,15 +28,47 @@ Mcap_Image :: struct { png: [] u8; // empty when there is none } +// What a model added to a capture, in its own format. A file holds one chunk per tag and key. +Mcap_Chunk :: struct { + tag: [4] u8; // what it is, in four ASCII letters, such as "labl" or "embd" + key: string; // whose: the model's name + data: [] u8; +} + Mcap :: struct { version: u32 = MCAP_VERSION; // as read; always MCAP_VERSION when written text: string; // Capture Text ocr: string; // OCR text (ocr_text.jai), one section per model; "" when none, and always in version 1 + front_chunks: [] Mcap_Chunk; // small: read with the front. None before version 3 preview: Mcap_Image; original_name: string; // the original image's name in the .snagx, such as {GUID}.png original: [] u8; // PNG rendered_kind: Mcap_Rendered_Kind; rendered: [] u8; // PNG; OWN only, see mcap_rendered + back_chunks: [] Mcap_Chunk; // large: after the images. None before version 3 +} + +// "labl" as a tag. +mcap_tag :: (text: string) -> [4] u8 { + assert(text.count == 4, "a chunk tag is four letters"); + tag: [4] u8; + for 0..3 tag[it] = text[it]; + return tag; +} + +mcap_tag_text :: (tag: [4] u8) -> string { + return tprint("%", to_string(tag.data, 4)); +} + +// The chunk with tag and key in chunks, or null. +mcap_find_chunk :: (chunks: [] Mcap_Chunk, tag: [4] u8, key: string) -> *Mcap_Chunk { + for * chunks if mcap_same_tag(it.tag, tag) && it.key == key return it; + return null; +} + +mcap_same_tag :: (a: [4] u8, b: [4] u8) -> bool { + for a if it != b[it_index] return false; + return true; } // The rendered image (Snagit's thumbnail.png) whichever way it is stored; empty when there is none. @@ -48,8 +85,9 @@ mcap_rendered :: (m: Mcap) -> [] u8 { // The file's bytes. A part too large for its size field fails the whole file rather than writing one // that cannot be read back. mcap_serialize :: (m: Mcap) -> Serializer, ok: bool, error: string { - size := MCAP_HEADER_SIZE + 4 + m.text.count + 4 + m.ocr.count + 8 + m.preview.png.count + 2 + m.original_name.count + 4 + m.original.count + 1; + size := MCAP_HEADER_SIZE + 4 + m.text.count + 4 + m.ocr.count + mcap_chunks_size(m.front_chunks) + 8 + m.preview.png.count + 2 + m.original_name.count + 4 + m.original.count + 1; if m.rendered_kind == .OWN size += 4 + m.rendered.count; + size += mcap_chunks_size(m.back_chunks); s := new_serializer(size); magic := MCAP_MAGIC; @@ -58,6 +96,8 @@ mcap_serialize :: (m: Mcap) -> Serializer, ok: bool, error: string { if !AddString(*s, m.text, u32) return s, false, "the text is too large for its size field"; if !AddString(*s, m.ocr, u32) return s, false, "the OCR text is too large for its size field"; + ok, error := mcap_add_chunks(*s, m.front_chunks); + if !ok return s, false, error; Add(*s, m.preview.width); Add(*s, m.preview.height); if !AddSizedArray(*s, m.preview.png, u32) return s, false, "the preview is too large for its size field"; @@ -65,6 +105,8 @@ mcap_serialize :: (m: Mcap) -> Serializer, ok: bool, error: string { if !AddSizedArray(*s, m.original, u32) return s, false, "the original image is too large for its size field"; Add(*s, m.rendered_kind); if m.rendered_kind == .OWN && !AddSizedArray(*s, m.rendered, u32) return s, false, "the rendered image is too large for its size field"; + ok, error = mcap_add_chunks(*s, m.back_chunks); + if !ok return s, false, error; assert(s.count == size); return s, true, ""; @@ -94,14 +136,18 @@ mcap_parse :: (data: [] u8) -> Mcap, ok: bool, error: string { m.original = ReadSizedArray(*d, u8, u32); m.rendered_kind = Read(*d, Mcap_Rendered_Kind); if d.ok && m.rendered_kind == .OWN m.rendered = ReadSizedArray(*d, u8, u32); + if d.ok && m.version >= 3 { + m.back_chunks, ok, error = mcap_read_chunks(*d); + if !ok return m, false, error; + } if !d.ok return m, false, "the file ends early (it is truncated or damaged)"; if cast(u8) m.rendered_kind > cast(u8) Mcap_Rendered_Kind.OWN return m, false, tprint("the rendered image is of unknown kind %", cast(u8) m.rendered_kind); if d.cursor != d.count return m, false, tprint("% bytes follow the last part", d.count - d.cursor); return m, true, ""; } -// Only the front of the file, as a library browser needs it: the header, text, OCR text and, with -// preview, the preview. One read of the start of the file, and a few more only when a part runs past +// Only the front of the file, as a library browser needs it: the header, text, OCR text, front chunks +// and, with preview, the preview. One read of the start of the file, and a few more only when a part runs past // it. The images are left unread and empty. The parts are views into a buffer from the context allocator. mcap_read_front :: (path: string, preview := true) -> Mcap, ok: bool, error: string { m: Mcap; @@ -112,7 +158,7 @@ mcap_read_front :: (path: string, preview := true) -> Mcap, ok: bool, error: str if !length_ok return m, false, "cannot be read"; wanted := min(length, MCAP_FRONT_READ_SIZE); - for attempt: 1..5 { // each further read reaches at least one more size field + for attempt: 1..MCAP_FRONT_READS { buffer := NewArray(wanted, u8, initialized = false); if !file_read(file, buffer.data, wanted) { array_free(buffer); return m, false, "cannot be read"; } @@ -121,22 +167,25 @@ mcap_read_front :: (path: string, preview := true) -> Mcap, ok: bool, error: str if ok return m, true, ""; if d.ok || wanted == length { array_free(buffer); return m, false, error; } // not a matter of too few bytes - // Short: the sizes read so far say how much is needed. + // Short: the sizes read so far say how far the front reaches at least. Read that and as much + // again as the first read, so the size fields after it (the OCR, the chunk count, each chunk's + // size) come in the same read rather than one read each. needed := mcap_front_size(buffer, preview); array_free(buffer); if needed <= wanted || needed > length return m, false, "the file ends early (it is truncated or damaged)"; - wanted = needed; + wanted = min(needed + MCAP_FRONT_READ_SIZE, length); if !file_set_position(file, 0) return m, false, "cannot be read"; } - return m, false, "cannot be read"; + return m, false, tprint("its front was not all read in % reads", MCAP_FRONT_READS); } #scope_module MCAP_HEADER_SIZE :: 12; // magic and version MCAP_FRONT_READ_SIZE :: 64 * 1024; // holds the header, text and preview of nearly every capture +MCAP_FRONT_READS :: 16; // reads of the front at most: each reaches past at least one more size field -// Header, text, OCR text and, with preview, the preview. +// Header, text, OCR text, front chunks and, with preview, the preview. mcap_read_front_part :: (d: *Deserializer, m: *Mcap, preview := true) -> ok: bool, error: string { magic := ReadArrayView(d, MCAP_MAGIC.count, u8); if !d.ok return false, "is not an .mcap file (it is shorter than the magic)"; @@ -148,6 +197,12 @@ mcap_read_front_part :: (d: *Deserializer, m: *Mcap, preview := true) -> ok: boo m.text = ReadSizedString(d, u32); if m.version >= 2 m.ocr = ReadSizedString(d, u32); // version 1 has no OCR block + if d.ok && m.version >= 3 { // nor versions before 3 chunks + ok: bool; + error: string; + m.front_chunks, ok, error = mcap_read_chunks(d); + if !ok return false, error; + } if !preview { if !d.ok return false, "the file ends early (it is truncated or damaged)"; return true, ""; @@ -176,13 +231,131 @@ mcap_front_size :: (prefix: [] u8, preview := true) -> s64 { } if !skip_sized(prefix, *position) return position; // the text if version >= 2 && !skip_sized(prefix, *position) return position; // the OCR text + if version >= 3 { // the front chunks + if position + 4 > prefix.count return position + 4; + d.cursor = position; + count := Read(*d, u32); + position += 4; + for 1..count { + position += 4; // the tag + if !skip_sized(prefix, *position) return position; // its size and payload + } + } if !preview return position; position += 4; // the preview's width and height skip_sized(prefix, *position); // the preview return position; } +// Chunks: a u32 count, then each as its tag, a u32 size, and that many bytes: a u16 key size, the +// key, and the data. +mcap_chunks_size :: (chunks: [] Mcap_Chunk) -> s64 { + size: s64 = 4; + for chunks size += 4 + 4 + 2 + it.key.count + it.data.count; + return size; +} + +mcap_add_chunks :: (s: *Serializer, chunks: [] Mcap_Chunk) -> ok: bool, error: string { + Add(s, cast(u32) chunks.count); + for chunks { + payload := 2 + it.key.count + it.data.count; + if it.key.count > 0xFFFF || payload > 0xFFFF_FFFF return false, tprint("the % chunk of % is too large for its size field", mcap_tag_text(it.tag), it.key); + tag := it.tag; + AddArray_NoSize(s, tag); + Add(s, cast(u32) payload); + AddString(s, it.key, u16); + AddArray_NoSize(s, it.data); + } + return true, ""; +} + +// Views into d's bytes; the list itself from the context allocator. +mcap_read_chunks :: (d: *Deserializer) -> [] Mcap_Chunk, ok: bool, error: string { + count := Read(d, u32); + if !d.ok return .[], false, "the file ends early (it is truncated or damaged)"; + // Running out of bytes clears d.ok, as Read does: a front read then reads further, a whole parse fails. + if cast(s64) count * 10 > d.count - d.cursor { d.ok = false; return .[], false, tprint("it says it has % chunks, more than its bytes can hold", count); } + chunks := NewArray(count, Mcap_Chunk); + for * chunk: chunks { + tag := ReadArrayView(d, 4, u8); + payload := Read(d, u32); + if !d.ok || payload < 2 return chunks, false, "a chunk runs past the end of the file (it is truncated or damaged)"; + if cast(s64) payload > d.count - d.cursor { d.ok = false; return chunks, false, "a chunk runs past the end of the file (it is truncated or damaged)"; } + for 0..3 chunk.tag[it] = tag[it]; + chunk.key = ReadSizedString(d, u16); + if !d.ok || 2 + chunk.key.count > cast(s64) payload return chunks, false, "a chunk's key runs past the chunk"; + chunk.data = ReadArrayView(d, payload - 2 - chunk.key.count, u8); + } + if !d.ok return chunks, false, "the file ends early (it is truncated or damaged)"; + return chunks, true, ""; +} + +#scope_export + +// Stores what a scan adds to the .mcap at path in one rewrite: with ocr, that OCR section in place of +// any by the same model (the others kept); and each of front, in place of any front chunk with the +// same tag and key (every other chunk kept, known to this build or not). Upgrades the file to the +// current version. Rewritten beside itself and swapped in (mcap_write), so a failure leaves it as it was. +mcap_store :: (path: string, ocr: *Mcap_Ocr_Section = null, front: [] Mcap_Chunk = .[]) -> ok: bool, error: string { + data, read_ok := read_entire_file(path, log_errors = false); + if !read_ok return false, "cannot be read"; + defer free(data); + m, parsed, parse_error := mcap_parse(cast([] u8) data); + if !parsed return false, parse_error; + + push_allocator(temp); + if ocr { + sections, ok, error := mcap_ocr_parse(m.ocr); + if !ok return false, error; + replaced := false; + for * sections if it.model == ocr.model { it.* = ocr.*; replaced = true; break; } + if !replaced array_add(*sections, ocr.*); + m.ocr = mcap_ocr_text(sections); + } + m.front_chunks = mcap_chunks_with(m.front_chunks, front); + m.version = MCAP_VERSION; + + serializer, serialized, serialize_error := mcap_serialize(m); + if !serialized return false, serialize_error; + written, write_error := mcap_write(path, serializer); + return written, write_error; +} + +// Stores chunk alone: in front, or with back after the images. See mcap_store. +mcap_store_chunk :: (path: string, chunk: Mcap_Chunk, back := false) -> ok: bool, error: string { + if !back { + ok, error := mcap_store(path, front = .[chunk]); + return ok, error; + } + data, read_ok := read_entire_file(path, log_errors = false); + if !read_ok return false, "cannot be read"; + defer free(data); + m, parsed, parse_error := mcap_parse(cast([] u8) data); + if !parsed return false, parse_error; + push_allocator(temp); + m.back_chunks = mcap_chunks_with(m.back_chunks, .[chunk]); + m.version = MCAP_VERSION; + serializer, serialized, serialize_error := mcap_serialize(m); + if !serialized return false, serialize_error; + written, write_error := mcap_write(path, serializer); + return written, write_error; +} + +// chunks with each of added in place of the one with its tag and key, or after them. Temp. +mcap_chunks_with :: (chunks: [] Mcap_Chunk, added: [] Mcap_Chunk) -> [] Mcap_Chunk { + list: [..] Mcap_Chunk; + list.allocator = temp; + array_add(*list, ..chunks); + for added { + existing := mcap_find_chunk(list, it.tag, it.key); + if existing existing.* = it; + else array_add(*list, it); + } + return list; +} + #load "ocr_text.jai"; +#load "label_chunks.jai"; #import "Basic"; #import "File"; diff --git a/modules/Capture_File/ocr_text.jai b/modules/Capture_File/ocr_text.jai index 16addb4..5390e39 100644 --- a/modules/Capture_File/ocr_text.jai +++ b/modules/Capture_File/ocr_text.jai @@ -75,7 +75,7 @@ mcap_ocr_parse :: (text: string) -> [..] Mcap_Ocr_Section, ok: bool, error: stri value := trim(advance(line, colon + 1)); if key == { case "engine"; - quoted, _, ok := ocr_read_quoted(value); + quoted, _, ok := mcap_read_quoted(value); if ok section.engine = copy_string(quoted); case "image"; parts := split(value, " ",, temp); @@ -102,13 +102,13 @@ mcap_ocr_text :: (sections: [] Mcap_Ocr_Section) -> string { print_to_builder(*builder, "@version %\n", MCAP_OCR_TEXT_VERSION); for sections { print_to_builder(*builder, "\n[ocr %]\n", it.model); - print_to_builder(*builder, "engine: %\n", ocr_quote(it.engine)); + print_to_builder(*builder, "engine: %\n", mcap_quote(it.engine)); print_to_builder(*builder, "image: % % %\n", it.image, it.image_width, it.image_height); print_to_builder(*builder, "date: %\n", it.date); for line: it.lines { c := mcap_ocr_corners(line); print_to_builder(*builder, "line: % % % % % % % % % % % % % %\n", line.x, line.y, line.width, line.height, - c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], formatFloat(line.confidence, trailing_width = 3, zero_removal = .NO), ocr_quote(line.text)); + c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], formatFloat(line.confidence, trailing_width = 3, zero_removal = .NO), mcap_quote(line.text)); } } return builder_to_string(*builder,, temp); @@ -118,25 +118,8 @@ mcap_ocr_text :: (sections: [] Mcap_Ocr_Section) -> string { // others, and upgrades the file to the current version. The file is rewritten beside itself and // swapped in (mcap_write), so a failure leaves it as it was. mcap_store_ocr :: (path: string, section: Mcap_Ocr_Section) -> ok: bool, error: string { - data, read_ok := read_entire_file(path, log_errors = false); - if !read_ok return false, "cannot be read"; - defer free(data); - m, parsed, parse_error := mcap_parse(cast([] u8) data); - if !parsed return false, parse_error; - - push_allocator(temp); - sections, ok, error := mcap_ocr_parse(m.ocr); - if !ok return false, error; - replaced := false; - for * sections if it.model == section.model { it.* = section; replaced = true; break; } - if !replaced array_add(*sections, section); - m.ocr = mcap_ocr_text(sections); - m.version = MCAP_VERSION; - - serializer, serialized, serialize_error := mcap_serialize(m); - if !serialized return false, serialize_error; - written, write_error := mcap_write(path, serializer); - return written, write_error; + ok, error := mcap_store(path, *section); + return ok, error; } // The time now as the date: key writes it, UTC. Temp. @@ -146,10 +129,10 @@ 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_file +#scope_module // the quoting is shared with label_chunks.jai // Quoted, a quote written twice, a control character as a space. Temp. -ocr_quote :: (text: string) -> string { +mcap_quote :: (text: string) -> string { builder: String_Builder; builder.allocator = temp; append(*builder, #char "\""); @@ -164,7 +147,7 @@ ocr_quote :: (text: string) -> string { } // One quoted string ("" for a quote) at the start of s: its text (context allocator) and what follows. -ocr_read_quoted :: (s: string) -> text: string, rest: string, ok: bool { +mcap_read_quoted :: (s: string) -> text: string, rest: string, ok: bool { if !s.count || s[0] != #char "\"" return "", s, false; builder: String_Builder; i := 1; @@ -195,7 +178,7 @@ ocr_parse_line :: (value: string, version: s64) -> Mcap_Ocr_Line, bool { it.* = number; rest = advance(rest, space); } - text, after, ok := ocr_read_quoted(trim_left(rest)); + text, after, ok := mcap_read_quoted(trim_left(rest)); if !ok || trim(after).count return result, false; result.x, result.y = cast(s32) numbers[0], cast(s32) numbers[1]; result.width, result.height = cast(s32) numbers[2], cast(s32) numbers[3]; diff --git a/modules/Labels/README.md b/modules/Labels/README.md new file mode 100644 index 0000000..b432c5d --- /dev/null +++ b/modules/Labels/README.md @@ -0,0 +1,52 @@ +# Labels + +Images and text in one space, from an image-text encoder: SigLIP 2 base at 384 px (Google, Apache +2.0), as converted to ONNX by onnx-community. An image can then be scored against any phrase: a label +list, or words typed into a search. It runs ONNX models through onnxruntime +([`modules/onnxruntime`](../onnxruntime)), sharing its session and CUDA code with +[`modules/OCR`](../OCR). + +## Setting up + +``` +cd modules/onnxruntime && jai fetch.jai # if OCR has not fetched it already +cd modules/Labels && jai fetch_models.jai # the vision model (fp32, 373 MB), the text model (int8, 283 MB), tokenizer.json (34 MB) +``` + +Each file's SHA-256 is checked. Building the app copies the models into `bin/label-models`. + +## Using it + +```jai +#import "Labels"; + +engine, ok, error := labels_create(models_folder, "siglip2-base-patch16-384"); // prefer_gpu = true +images := labels_embed_images(engine, pictures); // [] Labels_Image (RGBA): one unit vector each, engine.dimensions long +phrases := labels_embed_texts(engine, .["a photo of a cat", "a screenshot of a terminal"]); +score := labels_score(engine, images.data, phrases.data); // 0 to 1: the model's sigmoid +labels_destroy(engine); +``` + +- **Images** are resized to 384 x 384 (not keeping their shape) with PIL's bilinear filter, which + averages over each output pixel's footprint when shrinking, then scaled to -1..1: as SigLIP's own + processor does. The vision model runs on CUDA when it can, else on the CPU. +- **Text** is lowercased (ASCII letters) and tokenized with Gemma's tokenizer (`tokenizer.jai`: BPE + with byte fallback, read from `tokenizer.json`), cut or padded to 64 tokens. The text model runs + on the CPU: short inputs, a few at a time. +- **Scores** are `sigmoid(e^logit_scale * cosine + logit_bias)`, with the trained model's + `logit_scale` and `logit_bias` (in `models.config`; the ONNX encoders do not have them). They are + small for most phrases; their order is what counts. Descriptive prompts ("a screenshot of a code + editor showing source code") score far better than single words. +- `labels_to_half` and `labels_from_half` turn a vector into 16-bit floats, as `.mcap` files store + them, and back. + +## Checked against the reference + +- **The tokenizer** gives the same ids as Hugging Face's `tokenizers` (0.23.2) on the same + `tokenizer.json`, for 336 strings: label phrases, punctuation, digits, accents, Chinese, Japanese, + emoji, tabs and newlines, and text past 64 tokens. +- **The vectors** match SigLIP's own preprocessing through the same ONNX files (Python, onnxruntime + and PIL): image vectors at a cosine of 0.9998 or more, text vectors exactly, scores within 0.001. + +On the RTX 5080 laptop, an image takes about 18 ms on the GPU (preparing it included) and 190 ms on +the CPU; a phrase about 5 ms. diff --git a/modules/Labels/fetch_models.jai b/modules/Labels/fetch_models.jai new file mode 100644 index 0000000..7e7ec30 --- /dev/null +++ b/modules/Labels/fetch_models.jai @@ -0,0 +1,67 @@ +// Downloads the files models.config names into models/, checking each one's SHA-256. The models +// are not in git; run this after cloning. Files already there with the right hash are skipped. +// +// jai fetch_models.jai (about 690 MB) +// +// SigLIP 2 (Google, Apache 2.0), as converted to ONNX by onnx-community: +// https://huggingface.co/onnx-community/siglip2-base-patch16-384-ONNX + +SOURCE :: "https://huggingface.co/onnx-community/siglip2-base-patch16-384-ONNX/resolve/main"; + +Model_File :: struct { name: string; source_path: string; sha256: string; } + +FILES :: Model_File.[ + .{"siglip2-base-patch16-384-vision.onnx", "onnx/vision_model.onnx", "49ae4958b1098ca995e929d646f7be05a69c65e6344beae07d58c6598ffc5210"}, + .{"siglip2-base-patch16-384-text-int8.onnx", "onnx/text_model_int8.onnx", "30a94b43b9ddc39c04ab67f80cba4b808262682b2fa1a2a8b924bf1e7bd84c61"}, + .{"siglip2-tokenizer.json", "tokenizer.json", "cb9140fae3ac5122c972d37adf83e1248471a38147ad76f8215c8872c6fd8322"}, +]; + +#run,stallable { if !fetch() exit(1); } + +fetch :: () -> bool { + set_build_options_dc(.{do_output = false}); + here := trim_right(path_strip_filename(#file), "/\\"); + folder := tprint("%/models", here); + make_directory_if_it_does_not_exist(folder); + + for FILES { + path := tprint("%/%", folder, it.name); + if file_exists(path) && sha256_of(path) == it.sha256 { log("% is already here.", it.name); continue; } + + url := tprint("%/%", SOURCE, it.source_path); + log("Downloading %...", url); + result := run_command("curl", "-L", "--fail", "--progress-bar", "-o", path, url); + if result.exit_code != 0 { log_error("curl failed for % (exit code %).", it.name, result.exit_code); return false; } + hash := sha256_of(path); + if hash != it.sha256 { + log_error("% has SHA-256 %, expected %; it was deleted.", it.name, hash, it.sha256); + file_delete(path); + return false; + } + } + log("The models are in %.", folder); + return true; +} + +// Lowercase hex, or "" when it could not be worked out. +sha256_of :: (path: string) -> string { + #if OS == .WINDOWS { + result, output := run_command("certutil", "-hashfile", path, "SHA256", capture_and_return_output = true); + if result.exit_code != 0 return ""; + lines := split(output, "\n"); + if lines.count < 2 return ""; + return to_lower_copy(replace(trim(lines[1]), " ", "")); // line 2 is the hash + } else { + command := ifx OS == .MACOS then string.["shasum", "-a", "256", path] else string.["sha256sum", path]; + result, output := run_command(..command, capture_and_return_output = true); + if result.exit_code != 0 return ""; + return to_lower_copy(trim(split(output, " ")[0])); + } +} + +#import "Basic"; +#import "Compiler"; +#import "File"; +#import "File_Utilities"; +#import "Process"; +#import "String"; diff --git a/modules/Labels/models.config b/modules/Labels/models.config new file mode 100644 index 0000000..2809b75 --- /dev/null +++ b/modules/Labels/models.config @@ -0,0 +1,23 @@ +@version 1 +# Label models: image-text encoders that place images and text in one space, so an image can be +# scored against any phrase. Each [model NAME] names the files it needs from this folder +# (fetch_models.jai downloads them) and how it is run; labels_create takes the NAME. +# +# title shown to people +# vision, text, tokenizer file names, quoted: the image encoder, the text encoder, tokenizer.json +# image_size images are resized to this square for the vision model (px) +# max_tokens text is cut or padded to this many tokens +# logit_scale, logit_bias a phrase's score for an image is sigmoid(e^logit_scale * cosine + logit_bias), +# from the trained model (they are not in the ONNX encoders) +# image_batch images run through the vision model together + +[model siglip2-base-patch16-384] +title: "SigLIP 2 base, 384 px" +vision: "siglip2-base-patch16-384-vision.onnx" +text: "siglip2-base-patch16-384-text-int8.onnx" +tokenizer: "siglip2-tokenizer.json" +image_size: 384 +max_tokens: 64 +logit_scale: 4.726024 +logit_bias: -16.772243 +image_batch: 8 diff --git a/modules/Labels/module.jai b/modules/Labels/module.jai new file mode 100644 index 0000000..c36cd51 --- /dev/null +++ b/modules/Labels/module.jai @@ -0,0 +1,381 @@ +// Labels: images and text in one space, from an image-text encoder (SigLIP 2), so an image can be +// scored against any phrase: a list of labels, or words typed into a search. +// +// engine, ok, error := labels_create(models_folder, "siglip2-base-patch16-384"); +// images := labels_embed_images(engine, pictures); // one unit vector per picture +// phrases := labels_embed_texts(engine, .["code editor", "schematic"]); +// score := labels_score(engine, image_vector, phrase_vector); // 0 to 1 +// +// The models are listed in models_folder/models.config (fetch_models.jai downloads them) and run +// with onnxruntime (modules/onnxruntime): the vision model on an NVIDIA GPU through CUDA when it +// can, else on the CPU; the text model, small inputs and only a few at a time, on the CPU. + +Labels_Provider :: enum { CPU; CUDA; } + +// Tightly packed or not: stride is the bytes from one row to the next. +Labels_Image :: struct { + pixels: *u8; // RGBA + width, height, stride: s32; +} + +Labels_Model_Info :: struct { + name: string; + title: string; + vision, text, tokenizer: string; // file names in the models folder + image_size: s32 = 384; + max_tokens: s32 = 64; + logit_scale: float = 0; + logit_bias: float = 0; + image_batch: s32 = 8; +} + +Labels_Engine :: struct { + model: Labels_Model_Info; + provider: Labels_Provider; // the vision model's + gpu_explanation: string; // why the GPU is not in use; "" when it is + env: *OrtEnv; + memory_info: *OrtMemoryInfo; + vision, text: Ort_Session; + tokenizer: *Gemma_Tokenizer; + dimensions: s32; // of every vector +} + +LABELS_MODELS_CONFIG :: "models.config"; + +// The models models.config lists, in its order; from the context allocator. +labels_list_models :: (models_folder: string) -> [] Labels_Model_Info, error: string { + models: [..] Labels_Model_Info; + path := tprint("%/%", models_folder, LABELS_MODELS_CONFIG); + text, ok := read_entire_file(path, log_errors = false,, temp); + if !ok return models, tprint("Cannot read %.", path); + + model: *Labels_Model_Info; + version_seen := false; + line_number := 0; + for raw: split(text, "\n",, temp) { + line_number += 1; + line := trim(labels_strip_comment(raw)); + if !line.count continue; + if begins_with(line, "@version") { version_seen = true; continue; } + if !version_seen return models, tprint("% does not start with @version 1.", path); + if begins_with(line, "[model ") && ends_with(line, "]") { + model = array_add(*models); + model.* = .{}; + model.name = copy_string(trim(slice(line, 7, line.count - 8))); + continue; + } + colon := find_index_from_left(line, #char ":"); + if colon < 0 || !model { log_error("[Labels] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; } + key := trim(slice(line, 0, colon)); + value := trim(advance(line, colon + 1)); + if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\"" value = slice(value, 1, value.count - 2); + number, number_ok := string_to_float(value); + if key == { + case "title"; model.title = copy_string(value); + case "vision"; model.vision = copy_string(value); + case "text"; model.text = copy_string(value); + case "tokenizer"; model.tokenizer = copy_string(value); + case "image_size"; if number_ok model.image_size = cast(s32) number; + case "max_tokens"; if number_ok model.max_tokens = cast(s32) number; + case "logit_scale"; if number_ok model.logit_scale = number; + case "logit_bias"; if number_ok model.logit_bias = number; + case "image_batch"; if number_ok model.image_batch = cast(s32) number; + case; log_error("[Labels] %:%: unknown key %. Ignored.", path, line_number, key); + } + } + return models, ""; +} + +// An engine for the model named in models_folder/models.config; its vision model on CUDA with +// prefer_gpu when every library that needs is present, else on the CPU. +labels_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Labels_Engine, ok: bool, error: string { + push_allocator(context.default_allocator); + loaded, load_error := ort_load(); + if !loaded return null, false, load_error; + + models, models_error := labels_list_models(models_folder); + if models_error return null, false, models_error; + model: *Labels_Model_Info; + for * models if it.name == model_name { model = it; break; } + if !model return null, false, tprint("% has no [model %].", LABELS_MODELS_CONFIG, model_name); + + engine := New(Labels_Engine); + engine.model = model.*; + api := ort.api; + ok, error := ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "labels", *engine.env)); + if !ok { labels_destroy(engine); return null, false, error; } + ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info)); + + engine.tokenizer, ok, error = tokenizer_load(tprint("%/%", models_folder, model.tokenizer)); + if !ok { labels_destroy(engine); return null, false, error; } + + HINT :: "run `jai fetch_models.jai` in modules/Labels and rebuild"; + cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1); + vision_path := tprint("%/%", models_folder, model.vision); + use_gpu := false; + if prefer_gpu { + available, explanation := ort_cuda_check("Labelling"); + use_gpu = available; + if !available engine.gpu_explanation = copy_string(explanation); + } else { + engine.gpu_explanation = copy_string("Labelling runs on the CPU, as it was asked to."); + } + if use_gpu { + ort_gpu_lock(); // one CUDA engine made or released at a time (modules/onnxruntime/session.jai) + engine.vision, ok, error = ort_create_session(engine.env, vision_path, true, 2, output = "pooler_output", missing_hint = HINT); + if ok ok, error = labels_warm_up(engine); + if !ok ort_release_session(*engine.vision); + ort_gpu_unlock(); + if ok { + engine.provider = .CUDA; + } else { + engine.gpu_explanation = sprint("Labelling runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", error); + use_gpu = false; + } + } + if !use_gpu { + engine.vision, ok, error = ort_create_session(engine.env, vision_path, false, cores, output = "pooler_output", missing_hint = HINT); + if !ok { labels_destroy(engine); return null, false, error; } + engine.provider = .CPU; + } + engine.text, ok, error = ort_create_session(engine.env, tprint("%/%", models_folder, model.text), false, max(cores / 2, 1), output = "pooler_output", missing_hint = HINT); + if !ok { labels_destroy(engine); return null, false, error; } + + // The vectors' size, from a run of the text model on an empty text. + probe := labels_embed_texts(engine, .[""],, temp); + if !engine.dimensions { labels_destroy(engine); return null, false, "The text model gave no vector."; } + return engine, true, ""; +} + +labels_destroy :: (engine: *Labels_Engine) { + if !engine return; + if engine.provider == .CUDA ort_gpu_lock(); + ort_release_session(*engine.vision); + if engine.provider == .CUDA ort_gpu_unlock(); + ort_release_session(*engine.text); + tokenizer_free(engine.tokenizer,, context.default_allocator); + if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info); + if engine.env ort.api.ReleaseEnv(engine.env); + free(engine.gpu_explanation,, context.default_allocator); + free(engine,, context.default_allocator); +} + +labels_gpu_status :: (engine: *Labels_Engine) -> in_use: bool, explanation: string { + return engine.provider == .CUDA, engine.gpu_explanation; +} + +// One unit vector per image, engine.dimensions long, one after another; from allocator. An image +// that fails gives a zero vector, and error says why (the last error). +labels_embed_images :: (engine: *Labels_Engine, images: [] Labels_Image, allocator := context.allocator) -> [] float, error: string { + size := engine.model.image_size; + dims := engine.dimensions; + result := NewArray(images.count * dims, float, initialized = true,, allocator); + error: string; + plane := cast(s64) size * size; + batch := cast(s64) max(engine.model.image_batch, 1); + input := NewArray(plane * 3 * batch, float,, context.default_allocator); + defer array_free(input,, context.default_allocator); + + first := 0; + while first < images.count { + count := min(batch, images.count - first); + for 0..count - 1 labels_prepare_image(images[first + it], size, input.data + plane * 3 * it); + shape := s64.[count, 3, size, size]; + out, ok, run_error := ort_run_float(*engine.vision, engine.memory_info, array_view(input, 0, plane * 3 * count), shape); + if !ok { + error = tprint("%", run_error); + } else { + if out.dim_count == 2 && out.dims[0] == count && out.dims[1] == dims { + for n: 0..count - 1 labels_normalize_into(out.data + n * dims, dims, result.data + (first + n) * dims); + } else { + error = tprint("The vision model gave an output of shape %, not [%, %].", array_view(out.dims, 0, out.dim_count), count, dims); + } + ort_release_output(*out); + } + first += count; + } + return result, error; +} + +// One unit vector per text, as labels_embed_images gives them; from allocator. +labels_embed_texts :: (engine: *Labels_Engine, texts: [] string, allocator := context.allocator) -> [] float, error: string { + if !texts.count return .[], ""; + tokens := engine.model.max_tokens; + ids := NewArray(texts.count * tokens, s64,, temp); + for text, n: texts { + encoded := tokenizer_encode(engine.tokenizer, text, tokens); + memcpy(ids.data + n * tokens, encoded.data, tokens * size_of(s64)); + } + shape := s64.[texts.count, tokens]; + out, ok, run_error := ort_run(*engine.text, engine.memory_info, ids.data, ids.count * size_of(s64), .ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64, shape); + if !ok return .[], tprint("%", run_error); + defer ort_release_output(*out); + if out.dim_count != 2 || out.dims[0] != texts.count return .[], tprint("The text model gave an output of shape %.", array_view(out.dims, 0, out.dim_count)); + dims := cast(s32) out.dims[1]; + if !engine.dimensions engine.dimensions = dims; + result := NewArray(texts.count * dims, float,, allocator); + for n: 0..texts.count - 1 labels_normalize_into(out.data + n * dims, dims, result.data + n * dims); + return result, ""; +} + +// How well a phrase fits an image, 0 to 1, from their unit vectors: the model's own sigmoid. +labels_score :: (engine: *Labels_Engine, image: *float, text: *float) -> float { + return labels_score(engine.model, image, text, engine.dimensions); +} + +labels_score :: (model: Labels_Model_Info, image: *float, text: *float, dimensions: s64) -> float { + cosine: float; + for 0..dimensions - 1 cosine += image[it] * text[it]; + logit := exp(model.logit_scale) * cosine + model.logit_bias; + return 1.0 / (1.0 + exp(-logit)); +} + +// A vector as half floats, for storing; and back. Temp / allocator. +labels_to_half :: (vector: [] float) -> [] u16 { + result := NewArray(vector.count, u16,, temp); + for vector result[it_index] = labels_float_to_half(it); + return result; +} + +labels_from_half :: (half: [] u16, allocator := context.allocator) -> [] float { + result := NewArray(half.count, float,, allocator); + for half result[it_index] = float16_to_float32(cast(Float16) it); + return result; +} + +// A float as 16-bit float bits: the Float16 module's float_to_half_fast3 (Fabian Giesen's, public +// domain), copied with #no_aoc. It relies on unsigned wraparound, which a build with arithmetic +// overflow checks (the app's release build) stops at; the module's own is not marked. +labels_float_to_half :: (value: float32) -> u16 #no_aoc { + f32_infinity :: cast(u32) 255 << 23; + f16_infinity :: cast(u32) 31 << 23; + magic : float32 : 0h0780_0000; + sign_mask : u32 : 0x8000_0000; + round_mask : u32 : 0xffff_f000; + + using bits: union { f: float32; u: u32; } + f = value; + sign := u & sign_mask; + u ^= sign; + half: u16; + if u >= f32_infinity { + half = ifx u > f32_infinity then cast(u16) 0x7e00 else 0x7c00; // NaN to a quiet NaN, infinity to infinity + } else { // normal, denormal or zero + u &= round_mask; + f *= magic; + u -= round_mask; + if u > f16_infinity u = f16_infinity; // too large: infinity + half = cast(u16) (u >> 13); + } + return half | cast(u16) (sign >> 16); +} + +#scope_module + +// Runs the vision model once, on a blank image, when it is made. As OCR's engines are (see +// ocr_warm_up): with cuBLASLt 13.8, the process's first CUDA session destroyed without ever having +// run left the next one to crash on its first run. +labels_warm_up :: (engine: *Labels_Engine) -> ok: bool, error: string { + size := engine.model.image_size; + input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator); + defer array_free(input,, context.default_allocator); + out, ok, error := ort_run_float(*engine.vision, engine.memory_info, input, s64.[1, 3, size, size]); + if !ok return false, tprint("The first run on the GPU failed: %", error); + ort_release_output(*out); + return true, ""; +} + +// The image resized to size x size (not keeping its shape, as SigLIP's processor does) with PIL's +// bilinear filter, which averages over each output pixel's footprint when shrinking; then as planar +// RGB, scaled to -1..1, at dest. +labels_prepare_image :: (image: Labels_Image, size: s32, dest: *float) { + // Horizontal pass into rows of size floats per channel, then vertical. + across := labels_resample_weights(image.width, size); + down := labels_resample_weights(image.height, size); + rows := NewArray(cast(s64) image.height * size * 3, float,, temp); + for y: 0..image.height - 1 { + row := image.pixels + cast(s64) y * image.stride; + for x: 0..size - 1 { + w := *across[x]; + r, g, b: float; + for k: 0..w.weights.count - 1 { + p := row + (w.first + k) * 4; + r += p[0] * w.weights[k]; g += p[1] * w.weights[k]; b += p[2] * w.weights[k]; + } + at := (cast(s64) y * size + x) * 3; + rows[at] = r; rows[at + 1] = g; rows[at + 2] = b; + } + } + plane := cast(s64) size * size; + for y: 0..size - 1 { + w := *down[y]; + for x: 0..size - 1 { + r, g, b: float; + for k: 0..w.weights.count - 1 { + at := (cast(s64) (w.first + k) * size + x) * 3; + r += rows[at] * w.weights[k]; g += rows[at + 1] * w.weights[k]; b += rows[at + 2] * w.weights[k]; + } + i := cast(s64) y * size + x; + dest[i] = clamp(r, 0, 255) / 127.5 - 1; // (v / 255 - 0.5) / 0.5 + dest[i + plane] = clamp(g, 0, 255) / 127.5 - 1; + dest[i + plane * 2] = clamp(b, 0, 255) / 127.5 - 1; + } + } +} + +Resample_Weights :: struct { first: s32; weights: [] float; } + +// For each of out positions, the input positions it takes and their weights: a triangle filter +// widened by the scale when shrinking, as PIL's bilinear resize. Temp. +labels_resample_weights :: (in: s32, out: s32) -> [] Resample_Weights { + result := NewArray(out, Resample_Weights,, temp); + scale := cast(float) in / out; + filter_scale := max(scale, 1.0); + support := filter_scale; // the triangle reaches 1 unit, scaled + for i: 0..out - 1 { + centre := (i + 0.5) * scale; + first := max(cast(s32) (centre - support + 0.5), 0); + last := min(cast(s32) (centre + support + 0.5), in); + weights := NewArray(max(last - first, 1), float,, temp); + total: float; + for k: 0..last - first - 1 { + x := (first + k - centre + 0.5) / filter_scale; + w := max(0.0, 1 - abs(x)); + weights[k] = w; + total += w; + } + if total > 0 for * weights it.* /= total; + else { first = clamp(cast(s32) centre, 0, in - 1); weights[0] = 1; } + result[i] = .{first, weights}; + } + return result; +} + +labels_normalize_into :: (source: *float, dims: s64, dest: *float) { + length: float; + for 0..dims - 1 length += source[it] * source[it]; + length = sqrt(length); + scale := ifx length > 0 then 1 / length else 0; + for 0..dims - 1 dest[it] = source[it] * scale; +} + +labels_strip_comment :: (line: string) -> string { + in_quotes := false; + for i: 0..line.count - 1 { + c := line[i]; + if c == #char "\"" in_quotes = !in_quotes; + else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t") return slice(line, 0, i); + } + return line; +} + +#load "tokenizer.jai"; + +#import "Basic"; +#import "String"; +#import "File"; +#import "Math"; +#import "System"; +#import "Hash_Table"; +#import "Float16"; +#import "onnxruntime"; diff --git a/modules/Labels/tokenizer.jai b/modules/Labels/tokenizer.jai new file mode 100644 index 0000000..b78efb6 --- /dev/null +++ b/modules/Labels/tokenizer.jai @@ -0,0 +1,248 @@ +// The tokenizer SigLIP 2's text model was trained with: Gemma's, a SentencePiece-style BPE with +// byte fallback, read from its tokenizer.json. What it does, as that file says: +// +// - spaces become "▁" (the normalizer), and the whole text is then one word; +// - each character is a piece, or when the vocabulary has no such piece, its UTF-8 bytes as +// <0xXX> pieces (byte fallback); +// - neighbouring pieces are merged, the pair with the lowest rank in the merges list first and the +// leftmost of equal ranks, until no listed pair is left; +// - goes at the end, and fills it out to max_length (64, as SigLIP 2 was trained). +// +// Before all that, text is lowercased, as SigLIP 2 was trained with lowercased text. Only ASCII +// letters are lowered here; other letters keep their case. + +Gemma_Tokenizer :: struct { + text: string; // tokenizer.json, which unescaped pieces point into + vocab: Table(string, s32); // piece to id + merges: Table(u64, Merge); // left id << 32 | right id, to the merge + byte_ids: [256] s32; // <0x00>..<0xFF> + unk, eos, pad: s32; +} + +Merge :: struct { rank: s32; result: s32; } + +GEMMA_SPACE :: "\xE2\x96\x81"; // ▁, U+2581 + +// Reads tokenizer.json. Everything from the context allocator; tokenizer_free frees it. +tokenizer_load :: (path: string) -> *Gemma_Tokenizer, ok: bool, error: string { + text, read := read_entire_file(path, log_errors = false); + if !read return null, false, tprint("Cannot read %.", path); + t := New(Gemma_Tokenizer); + t.text = text; + + // "vocab": { "": 0, ... } + p := find_index_from_left(text, "\"vocab\""); + if p < 0 { tokenizer_free(t); return null, false, tprint("% has no vocab.", path); } + rest := advance(text, p + 7); + if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "{") { tokenizer_free(t); return null, false, tprint("% has a vocab that is not an object.", path); } + table_ensure_space(*t.vocab, 300_000); + while true { + json_skip_space(*rest); + if rest.count && rest[0] == #char "}" { advance(*rest); break; } + piece, ok := json_string(*rest); + if !ok || !json_expect(*rest, #char ":") { tokenizer_free(t); return null, false, tprint("% has a vocab entry that cannot be read.", path); } + json_skip_space(*rest); + id, number_ok, after := string_to_int(rest, T = s64); + if !number_ok { tokenizer_free(t); return null, false, tprint("% has a vocab id that cannot be read.", path); } + rest = after; + table_set(*t.vocab, piece, cast(s32) id); + json_skip_space(*rest); + if rest.count && rest[0] == #char "," advance(*rest); + } + + // "merges": [ ["a", "b"], ... ] + p = find_index_from_left(rest, "\"merges\""); + if p < 0 { tokenizer_free(t); return null, false, tprint("% has no merges.", path); } + rest = advance(rest, p + 8); + if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has merges that are not a list.", path); } + table_ensure_space(*t.merges, 800_000); + rank: s32 = 0; + joined: [..] u8; // each merge's two pieces end to end, to look up their result + defer array_free(joined); + while true { + json_skip_space(*rest); + if rest.count && rest[0] == #char "]" { advance(*rest); break; } + if !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has a merge that is not a pair (older \"a b\" merges are not read).", path); } + left, left_ok := json_string(*rest); + json_expect(*rest, #char ","); + right, right_ok := json_string(*rest); + if !left_ok || !right_ok || !json_expect(*rest, #char "]") { tokenizer_free(t); return null, false, tprint("% has a merge that cannot be read.", path); } + json_skip_space(*rest); + if rest.count && rest[0] == #char "," advance(*rest); + + left_found, left_id := table_find(*t.vocab, left); + right_found, right_id := table_find(*t.vocab, right); + joined.count = 0; + array_add(*joined, ..cast([] u8) left); + array_add(*joined, ..cast([] u8) right); + result_found, result_id := table_find(*t.vocab, cast(string) joined); + if left_found && right_found && result_found { + table_set(*t.merges, ((cast(u64) left_id) << 32) | cast(u64) cast(u32) right_id, .{rank, result_id}); + } + rank += 1; + } + + found: bool; + found, t.unk = table_find(*t.vocab, ""); + if !found { tokenizer_free(t); return null, false, tprint("% has no .", path); } + found, t.eos = table_find(*t.vocab, ""); + if !found { tokenizer_free(t); return null, false, tprint("% has no .", path); } + found, t.pad = table_find(*t.vocab, ""); + if !found { tokenizer_free(t); return null, false, tprint("% has no .", path); } + for 0..255 { + byte_found, id := table_find(*t.vocab, tprint("<0x%>", formatInt(it, base = 16, minimum_digits = 2))); + t.byte_ids[it] = ifx byte_found then id else t.unk; + } + return t, true, ""; +} + +tokenizer_free :: (t: *Gemma_Tokenizer) { + if !t return; + // Pieces with escapes were copied; the others point into text. + for _, piece: t.vocab if !tokenizer_points_into(t, piece) free(piece); + table_deinit(*t.vocab); + table_deinit(*t.merges); + free(t.text); + free(t); +} + +// text as max_length token ids: its pieces (at most max_length - 1), , then . Temp. +tokenizer_encode :: (t: *Gemma_Tokenizer, text: string, max_length := 64) -> [] s64 { + // Lowercased, spaces as ▁. + normal: String_Builder; + normal.allocator = temp; + for 0..text.count - 1 { + c := text[it]; + if c == #char " " append(*normal, GEMMA_SPACE); + else if c >= #char "A" && c <= #char "Z" append(*normal, c + (#char "a" - #char "A")); + else append(*normal, c); + } + word := builder_to_string(*normal,, temp); + + // One piece per character; byte pieces for those the vocabulary lacks. + symbols: [..] s32; + symbols.allocator = temp; + i := 0; + while i < word.count { + length := utf8_length(word[i]); + if i + length > word.count length = word.count - i; // a cut-off sequence: its bytes + found, id := table_find(*t.vocab, slice(word, i, length)); + if found array_add(*symbols, id); + else for b: 0..length - 1 array_add(*symbols, t.byte_ids[word[i + b]]); + i += length; + } + + // Merges, lowest rank first, the leftmost of equal ranks. + while symbols.count > 1 { + best := -1; + best_merge: Merge; + for 0..symbols.count - 2 { + found, merge := table_find(*t.merges, ((cast(u64) symbols[it]) << 32) | cast(u64) cast(u32) symbols[it + 1]); + if found && (best < 0 || merge.rank < best_merge.rank) { best = it; best_merge = merge; } + } + if best < 0 break; + symbols[best] = best_merge.result; + array_ordered_remove_by_index(*symbols, best + 1); + } + + ids := NewArray(max_length, s64,, temp); + kept := min(symbols.count, max_length - 1); + for 0..kept - 1 ids[it] = symbols[it]; + ids[kept] = t.eos; + for kept + 1..max_length - 1 ids[it] = t.pad; + return ids; +} + +#scope_file + +// The length of the UTF-8 sequence a byte starts: 1 for ASCII and stray continuation bytes. +utf8_length :: (lead: u8) -> s64 { + if lead < 0xC0 return 1; + if lead < 0xE0 return 2; + if lead < 0xF0 return 3; + return 4; +} + +tokenizer_points_into :: (t: *Gemma_Tokenizer, s: string) -> bool { + return s.data >= t.text.data && s.data < t.text.data + t.text.count; +} + +json_skip_space :: (s: *string) { + while s.count && (s.*[0] == #char " " || s.*[0] == #char "\n" || s.*[0] == #char "\r" || s.*[0] == #char "\t") advance(s); +} + +json_expect :: (s: *string, c: u8) -> bool { + json_skip_space(s); + if !s.count || s.*[0] != c return false; + advance(s); + return true; +} + +// A JSON string at the start of s: a view into s when it has no escapes, else unescaped into a +// copy from the context allocator. +json_string :: (s: *string) -> string, ok: bool { + json_skip_space(s); + if !s.count || s.*[0] != #char "\"" return "", false; + advance(s); + end := 0; + escaped := false; + while end < s.count && s.*[end] != #char "\"" { + if s.*[end] == #char "\\" { escaped = true; end += 1; } + end += 1; + } + if end >= s.count return "", false; + raw := slice(s.*, 0, end); + advance(s, end + 1); + if !escaped return raw, true; + + builder: String_Builder; + i := 0; + while i < raw.count { + c := raw[i]; + if c != #char "\\" { append(*builder, c); i += 1; continue; } + if i + 1 >= raw.count return "", false; + e := raw[i + 1]; + i += 2; + if e == { + case #char "n"; append(*builder, #char "\n"); + case #char "t"; append(*builder, #char "\t"); + case #char "r"; append(*builder, #char "\r"); + case #char "b"; append(*builder, 8); + case #char "f"; append(*builder, 12); + case #char "u"; + code, ok := hex4(raw, i); + if !ok return "", false; + i += 4; + if code >= 0xD800 && code < 0xDC00 && i + 6 <= raw.count && raw[i] == #char "\\" && raw[i + 1] == #char "u" { + low, low_ok := hex4(raw, i + 2); + if low_ok && low >= 0xDC00 && low < 0xE000 { code = 0x10000 + ((code - 0xD800) << 10) + (low - 0xDC00); i += 6; } + } + append_utf8(*builder, code); + case; append(*builder, e); // \" \\ \/ + } + } + return builder_to_string(*builder), true; +} + +hex4 :: (s: string, at: s64) -> u32, ok: bool { + if at + 4 > s.count return 0, false; + value: u32; + for 0..3 { + c := s[at + it]; + digit: u32; + if c >= #char "0" && c <= #char "9" digit = c - #char "0"; + else if c >= #char "a" && c <= #char "f" digit = c - #char "a" + 10; + else if c >= #char "A" && c <= #char "F" digit = c - #char "A" + 10; + else return 0, false; + value = value * 16 + digit; + } + return value, true; +} + +append_utf8 :: (builder: *String_Builder, code: u32) { + if code < 0x80 { append(builder, cast(u8) code); return; } + if code < 0x800 { append(builder, cast(u8) (0xC0 | (code >> 6))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; } + if code < 0x10000 { append(builder, cast(u8) (0xE0 | (code >> 12))); append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; } + append(builder, cast(u8) (0xF0 | (code >> 18))); append(builder, cast(u8) (0x80 | ((code >> 12) & 0x3F))); + append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F))); +} diff --git a/modules/OCR/README.md b/modules/OCR/README.md index ab50c72..9fed55a 100644 --- a/modules/OCR/README.md +++ b/modules/OCR/README.md @@ -79,7 +79,7 @@ new size, and the server model ran 5 times slower. Several copies of the session CUDA streams, made both models slower. GPU memory still peaks near full with the server model; that is for a later change. -`ocr_create` looks for each of these first (`cuda_check.jai`) and says which are missing, rather +`ocr_create` looks for each of these first (`modules/onnxruntime/cuda.jai`, shared with Labels) and says which are missing, rather than letting onnxruntime fail later and obscurely. If they are all there but CUDA still cannot start, it runs on the CPU and keeps onnxruntime's message. diff --git a/modules/OCR/models.jai b/modules/OCR/models.jai index 9ab6db7..f1c83a6 100644 --- a/modules/OCR/models.jai +++ b/modules/OCR/models.jai @@ -83,13 +83,6 @@ ocr_strip_comment :: (line: string) -> string { return line; } -Ocr_Session :: struct { - session: *OrtSession; - input_name, output_name: *u8; // heap C strings - input_shape: [4] s64; // as the model declares it: -1 for a size it leaves open - input_rank: s64; -} - // The model's sessions on the GPU or the CPU: detector, recognizer and, when it has one, orientation // classifier. On failure, the reason; the sessions made so far are left for ocr_release_sessions. ocr_create_sessions :: (engine: *Ocr_Engine, models_folder: string, gpu: bool, intra_op_threads: s32) -> ok: bool, error: string { @@ -141,102 +134,16 @@ ocr_warm_up :: (engine: *Ocr_Engine) -> ok: bool, error: string { return true, ""; } -// A session for the model file; with CUDA when gpu is set. On failure, the reason. +// A session for the model file, set up as Ocr_Options says (modules/onnxruntime/session.jai). ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, ocr_options: Ocr_Options) -> Ocr_Session, ok: bool, error: string { - result: Ocr_Session; - api := ort.api; - options: *OrtSessionOptions; - ok, error := ort_check(api.CreateSessionOptions(*options)); - if !ok return result, false, error; - defer api.ReleaseSessionOptions(options); - - ort_check(api.SetSessionGraphOptimizationLevel(options, .ORT_ENABLE_ALL)); - ort_check(api.SetIntraOpNumThreads(options, intra_op_threads)); - ort_check(api.SetInterOpNumThreads(options, 1)); - - if gpu { - cuda: *OrtCUDAProviderOptionsV2; - ok, error = ort_check(api.CreateCUDAProviderOptions(*cuda)); - if !ok return result, false, error; - defer api.ReleaseCUDAProviderOptions(cuda); - // Every image is a different size: an exhaustive cuDNN search would be redone for each new shape. - keys := (*u8).["device_id", "cudnn_conv_algo_search", "arena_extend_strategy", "cudnn_conv_use_max_workspace"]; - values := (*u8).["0", "HEURISTIC", temp_c_string(ocr_options.cuda_arena_extend_strategy), temp_c_string(ifx ocr_options.cudnn_max_workspace then "1" else "0")]; - ok, error = ort_check(api.UpdateCUDAProviderOptions(cuda, keys.data, values.data, xx keys.count)); - if !ok return result, false, error; - ok, error = ort_check(api.SessionOptionsAppendExecutionProvider_CUDA_V2(options, cuda)); - if !ok return result, false, error; - } - - if !file_exists(path) return result, false, tprint("% is not there; run `jai fetch_models.jai` in modules/OCR and rebuild.", path); - ok, error = ort_check(api.CreateSession(env, ort_path(path), options, *result.session)); - if !ok return result, false, error; - - allocator: *OrtAllocator; - ort_check(api.GetAllocatorWithDefaultOptions(*allocator)); - name: *u8; - ort_check(api.SessionGetInputName(result.session, 0, allocator, *name)); - result.input_name = to_c_string(to_string(name),, context.default_allocator); - ort_check(api.AllocatorFree(allocator, name)); - ort_check(api.SessionGetOutputName(result.session, 0, allocator, *name)); - result.output_name = to_c_string(to_string(name),, context.default_allocator); - ort_check(api.AllocatorFree(allocator, name)); - - input_type: *OrtTypeInfo; - if ort_check(api.SessionGetInputTypeInfo(result.session, 0, *input_type)) { - tensor_info: *OrtTensorTypeAndShapeInfo; // belongs to input_type - if ort_check(api.CastTypeInfoToTensorInfo(input_type, *tensor_info)) && tensor_info { - rank: u64; - ort_check(api.GetDimensionsCount(tensor_info, *rank)); - result.input_rank = min(cast(s64) rank, result.input_shape.count); - ort_check(api.GetDimensions(tensor_info, result.input_shape.data, xx result.input_rank)); - } - api.ReleaseTypeInfo(input_type); - } - return result, true, ""; + settings := Ort_Gpu_Settings.{arena_extend_strategy = ocr_options.cuda_arena_extend_strategy, cudnn_max_workspace = ocr_options.cudnn_max_workspace}; + session, ok, error := ort_create_session(env, path, gpu, intra_op_threads, settings, missing_hint = "run `jai fetch_models.jai` in modules/OCR and rebuild"); + return session, ok, error; } -ocr_release_session :: (s: *Ocr_Session) { - if s.session ort.api.ReleaseSession(s.session); - free(s.input_name,, context.default_allocator); - free(s.output_name,, context.default_allocator); - s.* = .{}; -} - -// A run's output: its values stay valid until ocr_release_output. -Ocr_Output :: struct { - value: *OrtValue; - data: *float; - dims: [4] s64; - dim_count: s64; -} - -// Runs the model on a float tensor of shape (the session may be shared between threads). -ocr_run :: (s: *Ocr_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64) -> Ocr_Output, ok: bool, error: string { - out: Ocr_Output; - api := ort.api; - input: *OrtValue; - ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data.data, xx (data.count * size_of(float)), shape.data, xx shape.count, .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, *input)); - if !ok return out, false, error; - defer api.ReleaseValue(input); - - ok, error = ort_check(api.Run(s.session, null, *s.input_name, *input, 1, *s.output_name, 1, *out.value)); - if !ok return out, false, error; - - raw: *void; - ort_check(api.GetTensorMutableData(out.value, *raw)); - out.data = cast(*float) raw; - info: *OrtTensorTypeAndShapeInfo; - ort_check(api.GetTensorTypeAndShape(out.value, *info)); - count: u64; - ort_check(api.GetDimensionsCount(info, *count)); - out.dim_count = min(cast(s64) count, out.dims.count); - ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count)); - api.ReleaseTensorTypeAndShapeInfo(info); - return out, true, ""; -} - -ocr_release_output :: (out: *Ocr_Output) { - if out.value ort.api.ReleaseValue(out.value); - out.* = .{}; -} +// onnxruntime's session and run, under the names the OCR code has always used. +Ocr_Session :: Ort_Session; +Ocr_Output :: Ort_Output; +ocr_run :: ort_run_float; +ocr_release_output :: ort_release_output; +ocr_release_session :: ort_release_session; diff --git a/modules/OCR/module.jai b/modules/OCR/module.jai index 549c0cd..5ce4fbb 100644 --- a/modules/OCR/module.jai +++ b/modules/OCR/module.jai @@ -107,7 +107,7 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op use_gpu := false; if prefer_gpu { - available, explanation := ocr_cuda_check(); + available, explanation := ort_cuda_check("OCR"); // modules/onnxruntime/cuda.jai use_gpu = available; if !available engine.gpu_explanation = copy_string(explanation); } else { @@ -117,13 +117,15 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op if use_gpu { // A few CPU threads prepare images and decode results while the GPU runs the networks. engine.workers = max(options.gpu_workers, 1); + ort_gpu_lock(); // one CUDA engine made or released at a time (modules/onnxruntime/session.jai) sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, true, 2); if sessions_ok sessions_ok, sessions_error = ocr_warm_up(engine); + if !sessions_ok ocr_release_sessions(engine); + ort_gpu_unlock(); if sessions_ok { engine.provider = .CUDA; engine.rec_batch = max(engine.model.rec_batch_gpu, 1); } else { - ocr_release_sessions(engine); engine.gpu_explanation = sprint("OCR runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", sessions_error); use_gpu = false; } @@ -151,7 +153,9 @@ ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string { ocr_destroy :: (engine: *Ocr_Engine) { if !engine return; + if engine.provider == .CUDA ort_gpu_lock(); ocr_release_sessions(engine); + if engine.provider == .CUDA ort_gpu_unlock(); if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info); if engine.env ort.api.ReleaseEnv(engine.env); free(engine.gpu_explanation,, context.default_allocator); @@ -356,7 +360,6 @@ run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void #load "models.jai"; #load "detect.jai"; #load "recognize.jai"; -#load "cuda_check.jai"; #import "Basic"; #import "String"; diff --git a/modules/OCR/cuda_check.jai b/modules/onnxruntime/cuda.jai similarity index 73% rename from modules/OCR/cuda_check.jai rename to modules/onnxruntime/cuda.jai index b0823bc..b9ca657 100644 --- a/modules/OCR/cuda_check.jai +++ b/modules/onnxruntime/cuda.jai @@ -1,4 +1,5 @@ // Whether onnxruntime's CUDA provider can run here, and if not, what is missing and how to fix it. +// Shared by the modules that run models (OCR, Labels); what names the feature in the explanations. // The provider needs the NVIDIA driver, parts of the CUDA 13 toolkit and cuDNN 9 for CUDA 13; it // fails late and obscurely without them, so each library is looked for first, by the name the // provider loads it by (from onnxruntime_providers_cuda.dll's imports and the names it loads). @@ -46,23 +47,23 @@ Cuda_Part :: enum { DRIVER; TOOLKIT; CUDNN; } } // Whether the CUDA provider should be tried, and when not, an explanation for people: what is -// missing, and the steps that fix it. Temp. -ocr_cuda_check :: () -> available: bool, explanation: string { +// missing, and the steps that fix it, starting " runs on the CPU". Temp. +ort_cuda_check :: (what: string) -> available: bool, explanation: string { #if OS == .MACOS { - return false, "OCR runs on the CPU: onnxruntime has no CUDA on macOS."; + return false, tprint("% runs on the CPU: onnxruntime has no CUDA on macOS.", what); } provider_path := tprint("%/%", path_strip_filename(ort.library_path), CUDA_PROVIDER_LIBRARY); if !file_exists(provider_path) { - return false, tprint("OCR runs on the CPU: the onnxruntime beside the program is the CPU-only build (% is not there).\nTo use an NVIDIA GPU, run `jai fetch.jai` in modules/onnxruntime (without `cpu`) and rebuild.", CUDA_PROVIDER_LIBRARY); + return false, tprint("% runs on the CPU: the onnxruntime beside the program is the CPU-only build (% is not there).\nTo use an NVIDIA GPU, run `jai fetch.jai` in modules/onnxruntime (without `cpu`) and rebuild.", what, CUDA_PROVIDER_LIBRARY); } - #if OS == .WINDOWS ocr_add_cudnn_to_path(); + #if OS == .WINDOWS ort_add_cudnn_to_path(); missing_driver, missing_toolkit, missing_cudnn: [..] string; missing_driver.allocator, missing_toolkit.allocator, missing_cudnn.allocator = temp, temp, temp; for CUDA_LIBRARIES { - if ocr_library_loads(it.name) continue; + if ort_library_loads(it.name) continue; if it.part == { case .DRIVER; array_add(*missing_driver, it.name); case .TOOLKIT; array_add(*missing_toolkit, it.name); @@ -74,13 +75,13 @@ ocr_cuda_check :: () -> available: bool, explanation: string { builder: String_Builder; builder.allocator = temp; if missing_driver.count { - append(*builder, "OCR runs on the CPU: no NVIDIA driver was found ("); + print_to_builder(*builder, "% runs on the CPU: no NVIDIA driver was found (", what); append(*builder, join(..missing_driver, ", ",, temp)); append(*builder, "), so there is no NVIDIA GPU to use."); return false, builder_to_string(*builder,, temp); } - append(*builder, "OCR runs on the CPU. The NVIDIA GPU can run it once these are installed:\n"); + print_to_builder(*builder, "% runs on the CPU. The NVIDIA GPU can run it once these are installed:\n", what); step := 1; if missing_toolkit.count { print_to_builder(*builder, "\n%. The CUDA Toolkit 13.x, which has % that are missing.\n", step, join(..missing_toolkit, ", ",, temp)); @@ -107,39 +108,37 @@ ocr_cuda_check :: () -> available: bool, explanation: string { #scope_module -ocr_library_loads :: (name: string) -> bool { +// Whether the library loads; one that does stays loaded, for the life of the process. Freeing it +// again let a check on one thread unload cuBLASLt or cuDNN while onnxruntime on another thread was +// loading or using it, which crashed there (a DEP access violation in the unloaded library's memory). +// They are about to be used anyway. +ort_library_loads :: (name: string) -> bool { #if OS == .WINDOWS { // The search order the CUDA provider's own loads use: beside the program, the system folders, PATH. - module := LoadLibraryW(utf8_to_wide(name,, temp)); - if !module return false; - FreeLibrary(module); - return true; + return LoadLibraryW(utf8_to_wide(name,, temp)) != null; } else { - handle := dlopen(temp_c_string(name), RTLD_LAZY); - if !handle return false; - dlclose(handle); - return true; + return dlopen(temp_c_string(name), RTLD_LAZY) != null; } } #if OS == .WINDOWS { // When cuDNN does not load, puts the newest Program Files\NVIDIA\CUDNN\v9.*\bin\13.*\x64 (or // bin\13.*, as older releases lay it out) holding cudnn64_9.dll at the front of this process's PATH. - ocr_add_cudnn_to_path :: () { - if ocr_library_loads("cudnn64_9.dll") return; - program_files := ocr_environment_variable("ProgramW6432"); - if !program_files program_files = ocr_environment_variable("ProgramFiles"); + ort_add_cudnn_to_path :: () { + if ort_library_loads("cudnn64_9.dll") return; + program_files := ort_environment_variable("ProgramW6432"); + if !program_files program_files = ort_environment_variable("ProgramFiles"); if !program_files return; - versions := ocr_subfolders_newest_first(tprint("%\\NVIDIA\\CUDNN", program_files), "v9."); + versions := ort_subfolders_newest_first(tprint("%\\NVIDIA\\CUDNN", program_files), "v9."); for version: versions { - for cuda: ocr_subfolders_newest_first(tprint("%\\bin", version), "13.") { + for cuda: ort_subfolders_newest_first(tprint("%\\bin", version), "13.") { for folder: string.[tprint("%\\x64", cuda), cuda] { if !file_exists(tprint("%\\cudnn64_9.dll", folder)) continue; found := replace(folder, "/", "\\",, temp); - path := ocr_environment_variable("PATH"); - OcrSetEnvironmentVariableW(utf8_to_wide("PATH",, temp), utf8_to_wide(tprint("%;%", found, path),, temp)); - log("[OCR] cuDNN is not on PATH; using the one in %.", found); + path := ort_environment_variable("PATH"); + OrtSetEnvironmentVariableW(utf8_to_wide("PATH",, temp), utf8_to_wide(tprint("%;%", found, path),, temp)); + log("[onnxruntime] cuDNN is not on PATH; using the one in %.", found); return; } } @@ -148,7 +147,7 @@ ocr_library_loads :: (name: string) -> bool { // The folders in parent whose names start with prefix, the highest version first ("v9.26" above // "v9.3"). Temp. - ocr_subfolders_newest_first :: (parent: string, prefix: string) -> [..] string { + ort_subfolders_newest_first :: (parent: string, prefix: string) -> [..] string { Found :: struct { names: [..] string; prefix: string; } found: Found; found.names.allocator = temp; @@ -156,12 +155,12 @@ ocr_library_loads :: (name: string) -> bool { visit_files(parent, recursive = false, *found, (info: *File_Visit_Info, found: *Found) { if begins_with(info.short_name, found.prefix) array_add(*found.names, copy_string(info.full_name,, temp)); }, visit_files = false, visit_directories = true); - quick_sort(found.names, (a: string, b: string) -> s64 { return ocr_compare_versions(path_filename(b), path_filename(a)); }); + quick_sort(found.names, (a: string, b: string) -> s64 { return ort_compare_versions(path_filename(b), path_filename(a)); }); return found.names; } // Compares the numbers in two names one by one: "v9.26" > "v9.3". - ocr_compare_versions :: (a: string, b: string) -> s64 { + ort_compare_versions :: (a: string, b: string) -> s64 { next_number :: (s: *string) -> s64, found: bool { while s.count && (s.*[0] < #char "0" || s.*[0] > #char "9") advance(s); if !s.count return 0, false; @@ -180,16 +179,16 @@ ocr_library_loads :: (name: string) -> bool { } // Temp; "" when it is not set. - ocr_environment_variable :: (name: string) -> string { + ort_environment_variable :: (name: string) -> string { wide_name := utf8_to_wide(name,, temp); - size := OcrGetEnvironmentVariableW(wide_name, null, 0); + size := OrtGetEnvironmentVariableW(wide_name, null, 0); if !size return ""; buffer := NewArray(size, u16,, temp); - written := OcrGetEnvironmentVariableW(wide_name, buffer.data, size); + written := OrtGetEnvironmentVariableW(wide_name, buffer.data, size); if !written || written >= size return ""; return wide_to_utf8(buffer.data, cast(s32) written,, temp); } - OcrGetEnvironmentVariableW :: (name: *u16, buffer: *u16, size: u32) -> u32 #foreign kernel32 "GetEnvironmentVariableW"; - OcrSetEnvironmentVariableW :: (name: *u16, value: *u16) -> s32 #foreign kernel32 "SetEnvironmentVariableW"; + OrtGetEnvironmentVariableW :: (name: *u16, buffer: *u16, size: u32) -> u32 #foreign kernel32 "GetEnvironmentVariableW"; + OrtSetEnvironmentVariableW :: (name: *u16, value: *u16) -> s32 #foreign kernel32 "SetEnvironmentVariableW"; } diff --git a/modules/onnxruntime/module.jai b/modules/onnxruntime/module.jai index 049b594..b52981b 100644 --- a/modules/onnxruntime/module.jai +++ b/modules/onnxruntime/module.jai @@ -65,6 +65,8 @@ ort_path :: (path: string) -> *ORTCHAR_T { } #load "bindings.jai"; +#load "session.jai"; +#load "cuda.jai"; #scope_module @@ -106,3 +108,5 @@ Get_Api_Base :: #type () -> *OrtApiBase #c_call; #import "File"; #import "File_Utilities"; #import "System"; +#import "Sort"; +#import "Atomics"; diff --git a/modules/onnxruntime/session.jai b/modules/onnxruntime/session.jai new file mode 100644 index 0000000..045c213 --- /dev/null +++ b/modules/onnxruntime/session.jai @@ -0,0 +1,157 @@ +// Sessions: one model file each, run on an NVIDIA GPU through CUDA or on the CPU, and running them on +// a tensor. Shared by the modules that run models (OCR, Labels). + +// CUDA sessions are made, first run and released one at a time, process-wide: hold ort_gpu_lock +// around it. Done on two threads at once (onnxruntime 1.30.0, CUDA 13, cuBLASLt 13.8, RTX 5080 +// laptop), it crashed with a DEP access violation in cuBLASLt or near it, every time in a debug build; +// one at a time, never. Running sessions on many threads at once is fine. A spin lock, needing no +// setup, as engines can be made before anything else runs; it is held for a second or so at a time. +ort_gpu_lock :: () { + while atomic_swap(*ort_gpu_held, 1) != 0 sleep_milliseconds(2); +} + +ort_gpu_unlock :: () { + atomic_swap(*ort_gpu_held, 0); +} + +#scope_file +ort_gpu_held: s32; +#scope_export + +Ort_Session :: struct { + session: *OrtSession; + input_name, output_name: *u8; // heap C strings + input_shape: [4] s64; // as the model declares it: -1 for a size it leaves open + input_rank: s64; +} + +// How a CUDA session uses GPU memory. The defaults came from measuring OCR on the sample captures +// (RTX 5080 laptop): with nearly every input a new size, onnxruntime's exact-size arena filled GPU +// memory, and cuDNN's largest workspaces made it worse. +Ort_Gpu_Settings :: struct { + arena_extend_strategy := "kNextPowerOfTwo"; // reuses freed blocks for other sizes; "kSameAsRequested" grows by exactly what is asked + cudnn_max_workspace := false; // let cuDNN take as much working memory as it finds best (onnxruntime's own default) +} + +// A session for the model file at path; with CUDA when gpu is set. output names the output that +// ort_run returns; "" for the first. missing_hint is added when the file is not there, to say how +// to get it. On failure, the reason. +ort_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, settings := Ort_Gpu_Settings.{}, output := "", missing_hint := "") -> Ort_Session, ok: bool, error: string { + result: Ort_Session; + api := ort.api; + options: *OrtSessionOptions; + ok, error := ort_check(api.CreateSessionOptions(*options)); + if !ok return result, false, error; + defer api.ReleaseSessionOptions(options); + + ort_check(api.SetSessionGraphOptimizationLevel(options, .ORT_ENABLE_ALL)); + ort_check(api.SetIntraOpNumThreads(options, intra_op_threads)); + ort_check(api.SetInterOpNumThreads(options, 1)); + + if gpu { + cuda: *OrtCUDAProviderOptionsV2; + ok, error = ort_check(api.CreateCUDAProviderOptions(*cuda)); + if !ok return result, false, error; + defer api.ReleaseCUDAProviderOptions(cuda); + // Inputs come in many sizes: an exhaustive cuDNN search would be redone for each new shape. + keys := (*u8).["device_id", "cudnn_conv_algo_search", "arena_extend_strategy", "cudnn_conv_use_max_workspace"]; + values := (*u8).["0", "HEURISTIC", temp_c_string(settings.arena_extend_strategy), temp_c_string(ifx settings.cudnn_max_workspace then "1" else "0")]; + ok, error = ort_check(api.UpdateCUDAProviderOptions(cuda, keys.data, values.data, xx keys.count)); + if !ok return result, false, error; + ok, error = ort_check(api.SessionOptionsAppendExecutionProvider_CUDA_V2(options, cuda)); + if !ok return result, false, error; + } + + if !file_exists(path) return result, false, tprint("% is not there%.", path, ifx missing_hint then tprint("; %", missing_hint) else ""); + ok, error = ort_check(api.CreateSession(env, ort_path(path), options, *result.session)); + if !ok return result, false, error; + + allocator: *OrtAllocator; + ort_check(api.GetAllocatorWithDefaultOptions(*allocator)); + name: *u8; + ort_check(api.SessionGetInputName(result.session, 0, allocator, *name)); + result.input_name = to_c_string(to_string(name),, context.default_allocator); + ort_check(api.AllocatorFree(allocator, name)); + + // The output: the one named, else the first. + output_count: u64; + ort_check(api.SessionGetOutputCount(result.session, *output_count)); + names: [..] string; + names.allocator = temp; + for i: 0..cast(s64) output_count - 1 { + ort_check(api.SessionGetOutputName(result.session, xx i, allocator, *name)); + array_add(*names, copy_string(to_string(name),, temp)); + ort_check(api.AllocatorFree(allocator, name)); + } + chosen := ifx names.count then names[0] else ""; + if output { + found := array_find(names, output); + if !found { ort_release_session(*result); return result, false, tprint("% has no output %; it has %.", path, output, names); } + chosen = output; + } + result.output_name = to_c_string(chosen,, context.default_allocator); + + input_type: *OrtTypeInfo; + if ort_check(api.SessionGetInputTypeInfo(result.session, 0, *input_type)) { + tensor_info: *OrtTensorTypeAndShapeInfo; // belongs to input_type + if ort_check(api.CastTypeInfoToTensorInfo(input_type, *tensor_info)) && tensor_info { + rank: u64; + ort_check(api.GetDimensionsCount(tensor_info, *rank)); + result.input_rank = min(cast(s64) rank, result.input_shape.count); + ort_check(api.GetDimensions(tensor_info, result.input_shape.data, xx result.input_rank)); + } + api.ReleaseTypeInfo(input_type); + } + return result, true, ""; +} + +ort_release_session :: (s: *Ort_Session) { + if s.session ort.api.ReleaseSession(s.session); + free(s.input_name,, context.default_allocator); + free(s.output_name,, context.default_allocator); + s.* = .{}; +} + +// A run's output: its values stay valid until ort_release_output. +Ort_Output :: struct { + value: *OrtValue; + data: *float; // the values, for a float output (every model here has one) + dims: [4] s64; + dim_count: s64; +} + +// Runs the session on one input tensor of shape, bytes long at data, with element type element. +// A session may be run from several threads at once. +ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s64, element: ONNXTensorElementDataType, shape: [] s64) -> Ort_Output, ok: bool, error: string { + out: Ort_Output; + api := ort.api; + input: *OrtValue; + ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data, xx bytes, shape.data, xx shape.count, element, *input)); + if !ok return out, false, error; + defer api.ReleaseValue(input); + + ok, error = ort_check(api.Run(s.session, null, *s.input_name, *input, 1, *s.output_name, 1, *out.value)); + if !ok return out, false, error; + + raw: *void; + ort_check(api.GetTensorMutableData(out.value, *raw)); + out.data = cast(*float) raw; + info: *OrtTensorTypeAndShapeInfo; + ort_check(api.GetTensorTypeAndShape(out.value, *info)); + count: u64; + ort_check(api.GetDimensionsCount(info, *count)); + out.dim_count = min(cast(s64) count, out.dims.count); + ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count)); + api.ReleaseTensorTypeAndShapeInfo(info); + return out, true, ""; +} + +ort_run_float :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64) -> Ort_Output, ok: bool, error: string { + out, ok, error := ort_run(s, memory_info, data.data, data.count * size_of(float), .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, shape); + return out, ok, error; +} + +ort_release_output :: (out: *Ort_Output) { + if out.value ort.api.ReleaseValue(out.value); + out.* = .{}; +} diff --git a/src/App.jai b/src/App.jai index 4628f1a..29fd8b4 100644 --- a/src/App.jai +++ b/src/App.jai @@ -11,6 +11,7 @@ USING_WINDOWS :: OS == .WINDOWS; #load "test/logger_tests.jai"; #load "test/catalog_tests.jai"; #load "test/ocr_tests.jai"; + #load "test/labels_tests.jai"; } #load "app/main.jai"; // entry point @@ -20,6 +21,7 @@ USING_WINDOWS :: OS == .WINDOWS; #load "app/catalog.jai"; #load "app/stats_ui.jai"; #load "app/background.jai"; +#load "app/labels_ui.jai"; #load "app/exit_shutdown.jai"; #load "base/circle_buffer.jai"; @@ -46,16 +48,19 @@ USING_WINDOWS :: OS == .WINDOWS; #import "Unicode"; #import "File"; #import "Random"; // the noise viewer_overwrite_and_delete writes +#import "Float16"; // label vectors are stored as 16-bit floats #import "File_Utilities"; #import "Math"; #import "System"; #import "Thread"; #import "Hash_Table"; +#import "Hash"; // fnv1a_hash, for the label list's hash Shared_Channel :: #import "Shared_Memory_Channel"; #import "stb_image"; #import "stb_image_resize"; #import "Capture_File"; // the .mcap format, copied from D:\SnagitIndex\SnagitConverter #import "OCR"; +#import "Labels"; #import "Capture_Text"; #import "Base64"; #import "onnxruntime"; @@ -178,6 +183,7 @@ App :: struct { library_folder: string; // -library FOLDER: the .mcap folder the viewer shows, in place of VIEWER_FOLDER 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 last_window_restore_time: Apollo_Time; imgui_context: *ImGui.ImGuiContext; diff --git a/src/app/app_ui.jai b/src/app/app_ui.jai index ecb9dc9..610c445 100644 --- a/src/app/app_ui.jai +++ b/src/app/app_ui.jai @@ -8,6 +8,7 @@ APP_MAIN_WINDOW_LABEL :: "Main"; app_ui_startup :: () { viewer_startup(); ocr_ui_startup(); + labels_ui_startup(); } // Every frame, before input is handed to ImGui. Poll background work here. @@ -19,6 +20,7 @@ app_ui_frame_start :: () { app_ui_shutdown :: () { viewer_shutdown(); ocr_ui_shutdown(); + labels_ui_shutdown(); // after the OCR thread and the background workers have stopped stats_ui_shutdown(); } @@ -26,5 +28,6 @@ app_main_window_ui :: () { push_imgui_window(APP_MAIN_WINDOW_LABEL); viewer_ui(); ocr_ui_frame(); + labels_ui_frame(); stats_ui_frame(); } diff --git a/src/app/background.jai b/src/app/background.jai index 9c5a459..ec6b95b 100644 --- a/src/app/background.jai +++ b/src/app/background.jai @@ -7,6 +7,10 @@ // UI thread, and reads nothing else: never viewer.captures, which a refresh can grow. Catalog // entries replaced while jobs are out are freed once none are (background_retire). A newer search // makes the older ones stop early, and their results are dropped. +// +// A search with like: terms first turns their phrases into vectors with the label model's text +// encoder (a job of its own), then ranks the captures matching the rest by how near their image +// vectors are, best first. BACKGROUND_SEARCH_CHUNK :: 256; // captures matched per job @@ -14,6 +18,7 @@ Search_Item :: struct { index: s32; entry: *Catalog_Entry; // null while not loaded name: string; // the capture's; never freed + score: float; // for like: terms, the mean cosine of its image vector with theirs } background: struct { @@ -55,7 +60,7 @@ background_snapshot_visible :: () -> [] Search_Item { items := NewArray(viewer.visible.count, Search_Item,, context.default_allocator); for viewer.visible { capture := *viewer.captures[it]; - items[it_index] = .{it, capture.catalog, capture.name}; + items[it_index] = .{index = it, entry = capture.catalog, name = capture.name}; } return items; } @@ -68,6 +73,10 @@ background_frame :: () { job := cast(*Background_Job) work; background.jobs_out -= 1; if job.kind == { + case .SEARCH_EMBED; + request := job.request; + if request.generation == background.search_generation search_dispatch_chunks(request); + else search_request_free(request); case .SEARCH_CHUNK; request := job.request; request.chunk_matches[job.chunk] = job.matches; @@ -111,8 +120,10 @@ background_frame :: () { // Matches the search against every capture not removed, in chunks on the workers; the results // replace viewer.visible when they are all back. background_search_start :: () { - // While the catalog loads under a search, matches arrive every frame: a few searches a second do. - if viewer.search_terms.count && !viewer.sorted_by_date && to_float64_seconds(current_time_monotonic() - viewer.visible_rebuilt_at) < 0.25 return; + // Under a search, catalog entries arriving (as the catalog loads, or a scan stores readings) would + // start one every frame: a few searches a second do. Typing, a refresh and a delete reset + // visible_rebuilt_at, so they are searched at once. + if viewer.search_terms.count && to_float64_seconds(current_time_monotonic() - viewer.visible_rebuilt_at) < 0.25 return; viewer.visible_dirty = false; viewer.visible_rebuilt_at = current_time_monotonic(); @@ -128,16 +139,35 @@ background_search_start :: () { for viewer.search_terms request.terms[it_index] = .{it.field, copy_string(it.text)}; items: [..] Search_Item; array_reserve(*items, viewer_capture_count()); - for * viewer.captures if !it.removed array_add(*items, .{cast(s32) it_index, it.catalog, it.name}); + for * viewer.captures if !it.removed array_add(*items, .{index = cast(s32) it_index, entry = it.catalog, name = it.name}); request.items = items; - chunks := max((items.count + BACKGROUND_SEARCH_CHUNK - 1) / BACKGROUND_SEARCH_CHUNK, 1); - request.chunk_matches = NewArray(chunks, [..] Search_Item); + // like: phrases need their vectors first, from the label model's text encoder. + likes: [..] string; + likes.allocator = temp; + for request.terms if it.field == .LIKE array_add(*likes, it.text); + if likes.count && labels_ui.engine { + request.like_texts = NewArray(likes.count, string); + for likes request.like_texts[it_index] = copy_string(it); + request.labels = labels_ui.engine; + job := New(Background_Job); + job.* = .{kind = .SEARCH_EMBED, generation = generation, request = request}; + background_add(job); + return; + } + search_dispatch_chunks(request); +} + +// The chunks of a search whose like: vectors (if any) are made. +search_dispatch_chunks :: (request: *Search_Request) { + push_allocator(context.default_allocator); + chunks := max((request.items.count + BACKGROUND_SEARCH_CHUNK - 1) / BACKGROUND_SEARCH_CHUNK, 1); + request.chunk_matches = NewArray(chunks, [..] Search_Item); for chunk: 0..chunks - 1 { job := New(Background_Job); first := chunk * BACKGROUND_SEARCH_CHUNK; - job.* = .{kind = .SEARCH_CHUNK, generation = generation, request = request, chunk = chunk, - first = first, count = clamp(items.count - first, 0, BACKGROUND_SEARCH_CHUNK)}; + job.* = .{kind = .SEARCH_CHUNK, generation = request.generation, request = request, chunk = chunk, + first = first, count = clamp(request.items.count - first, 0, BACKGROUND_SEARCH_CHUNK)}; background_add(job); } } @@ -153,7 +183,7 @@ background_stats_start :: (items: [] Search_Item, model: string) { #scope_file -Background_Kind :: enum u8 { SEARCH_CHUNK; SEARCH_MERGE; STATS; } +Background_Kind :: enum u8 { SEARCH_EMBED; SEARCH_CHUNK; SEARCH_MERGE; STATS; } Background_Job :: struct { kind: Background_Kind; @@ -177,6 +207,9 @@ Search_Request :: struct { terms: [] Search_Term; items: [] Search_Item; by_date: bool; // newest first; else in the captures' own order + like_texts: [] string; // the like: phrases; the order is then best first + like_vectors: [] float; // their vectors, made by the SEARCH_EMBED job + labels: *Labels_Engine; chunk_matches: [] [..] Search_Item; // each chunk's, as they come back chunks_back: s64; } @@ -190,6 +223,9 @@ search_request_free :: (request: *Search_Request) { push_allocator(context.default_allocator); for request.terms free(it.text); array_free(request.terms); + for request.like_texts free(it); + array_free(request.like_texts); + array_free(request.like_vectors); array_free(request.items); for request.chunk_matches array_free(it); array_free(request.chunk_matches); @@ -201,6 +237,11 @@ background_proc :: (group: *Thread_Group, thread: *Thread, work: *void) -> Threa context.allocator = context.default_allocator; job := cast(*Background_Job) work; if job.kind == { + case .SEARCH_EMBED; + // A newer search may have started; the vectors are still made, as the run cannot be stopped. + vectors, error := labels_embed_texts(job.request.labels, job.request.like_texts,, context.default_allocator); + if error log_error("[search] like: could not be worked out: %", error); + job.request.like_vectors = vectors; case .SEARCH_CHUNK; search_chunk(job); case .SEARCH_MERGE; search_merge(job); case .STATS; job.stats = stats_compute(job.items, job.model); @@ -214,9 +255,35 @@ search_chunk :: (job: *Background_Job) { for i: job.first..job.first + job.count - 1 { if (i & 63) == 0 && atomic_read(*background.search_generation) != job.generation return; // a newer search: these are wasted item := request.items[i]; - if search_matches(item.entry, request.terms) array_add(*job.matches, item); + if !search_matches(item.entry, request.terms) continue; + if request.like_vectors.count { + item.score = search_like_score(request, item.entry); + } + array_add(*job.matches, item); } - if request.by_date quick_sort(job.matches, search_order_by_date); + if request.like_vectors.count quick_sort(job.matches, search_order_by_score); + else if request.by_date quick_sort(job.matches, search_order_by_date); +} + +// The mean cosine of the entry's image vector with the like: phrases' vectors (all unit vectors). +search_like_score :: (request: *Search_Request, entry: *Catalog_Entry) -> float { + half := entry.embedding; + dims := half.count; + phrases := request.like_vectors.count / max(dims, 1); + if !dims || phrases * dims != request.like_vectors.count return -1; + image := NewArray(dims, float,, temp); + for half image[it_index] = float16_to_float32(cast(Float16) it); + total: float; + for p: 0..phrases - 1 { + phrase := request.like_vectors.data + p * dims; + for 0..dims - 1 total += image[it] * phrase[it]; + } + return total / phrases; +} + +search_order_by_score :: (a: Search_Item, b: Search_Item) -> s64 { + if a.score != b.score return ifx a.score > b.score then -1 else 1; + return search_order_by_date(a, b); } // Every chunk's matches as one list. In date order, each chunk is sorted already, so they are merged @@ -226,12 +293,12 @@ search_merge :: (job: *Background_Job) { runs: [..] [] Search_Item; runs.allocator = temp; for request.chunk_matches array_add(*runs, it); - if request.by_date { + if request.by_date || request.like_vectors.count { while runs.count > 1 { next: [..] [] Search_Item; next.allocator = temp; for i: 0..(runs.count + 1) / 2 - 1 { - if 2 * i + 1 < runs.count array_add(*next, merge_two(runs[2 * i], runs[2 * i + 1])); + if 2 * i + 1 < runs.count array_add(*next, merge_two(runs[2 * i], runs[2 * i + 1], request.like_vectors.count > 0)); else array_add(*next, runs[2 * i]); } runs = next; @@ -243,13 +310,17 @@ search_merge :: (job: *Background_Job) { for run: runs for run array_add(*job.visible, it.index); } -// Two lists in date order as one. Temp. -merge_two :: (a: [] Search_Item, b: [] Search_Item) -> [] Search_Item { +// Two lists in order (by score, or by date) as one. Temp. +merge_two :: (a: [] Search_Item, b: [] Search_Item, by_score: bool) -> [] Search_Item { result := NewArray(a.count + b.count, Search_Item,, temp); i, j := 0, 0; for * result { - if j >= b.count || (i < a.count && search_order_by_date(a[i], b[j]) <= 0) { it.* = a[i]; i += 1; } - else { it.* = b[j]; j += 1; } + take_a: bool; + if j >= b.count take_a = true; + else if i >= a.count take_a = false; + else take_a = (ifx by_score then search_order_by_score(a[i], b[j]) else search_order_by_date(a[i], b[j])) <= 0; + if take_a { it.* = a[i]; i += 1; } + else { it.* = b[j]; j += 1; } } return result; } diff --git a/src/app/catalog.jai b/src/app/catalog.jai index 0cb76c3..8f7b107 100644 --- a/src/app/catalog.jai +++ b/src/app/catalog.jai @@ -9,9 +9,12 @@ Catalog_Entry :: struct { title, url, app: string; // WindowName, WebURL, AppName notes: string; // the annotations' text, one per line ocr: [..] Mcap_Ocr_Section; // one per model that read it + labels: [..] Mcap_Label; // the label model's (LABELS_MODEL) best labels, best first + label_list: string; // the hash of the label list they were scored against + embedding: [] u16; // the label model's image vector, 16-bit floats; empty when not labelled // Lowercased copies, for matching without case. - find_title, find_url, find_app, find_notes, find_ocr: string; + find_title, find_url, find_app, find_notes, find_ocr, find_labels: string; } // Everything the entry holds is from the default allocator. @@ -25,7 +28,11 @@ catalog_entry_free :: (entry: *Catalog_Entry) { array_free(it.lines); } array_free(entry.ocr); - free(entry.find_title); free(entry.find_url); free(entry.find_app); free(entry.find_notes); free(entry.find_ocr); + for entry.labels free(it.name); + array_free(entry.labels); + free(entry.label_list); + 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); } @@ -77,6 +84,22 @@ catalog_load :: (path: string) -> *Catalog_Entry, error: string { sections, ocr_ok := mcap_ocr_parse(front.ocr); if ocr_ok entry.ocr = sections; + // The label model's chunks (app/labels_ui.jai). + labels_chunk := mcap_find_chunk(front.front_chunks, mcap_tag(MCAP_LABELS_TAG), LABELS_MODEL); + if labels_chunk { + labels, labels_ok := mcap_labels_parse(cast(string) labels_chunk.data); + if labels_ok { + entry.labels = labels.labels; + entry.label_list = labels.list; + free(labels.engine); free(labels.date); + } + } + embedding_chunk := mcap_find_chunk(front.front_chunks, mcap_tag(MCAP_EMBEDDING_TAG), LABELS_MODEL); + if embedding_chunk { + half, half_ok := mcap_embedding_parse(embedding_chunk.data); + if half_ok entry.embedding = half; + } + entry.find_title = to_lower_copy(entry.title); entry.find_url = to_lower_copy(entry.url); entry.find_app = to_lower_copy(entry.app); @@ -85,6 +108,15 @@ catalog_load :: (path: string) -> *Catalog_Entry, error: string { all_ocr.allocator = temp; for entry.ocr for line: it.lines { append(*all_ocr, line.text); append(*all_ocr, "\n"); } entry.find_ocr = to_lower_copy(builder_to_string(*all_ocr,, temp)); + all_labels: String_Builder; + all_labels.allocator = temp; + // The labels search matches: the best, and any scoring at least half as much (a close second). + for entry.labels { + if it_index > 0 && it.score < entry.labels[0].score * 0.5 continue; + append(*all_labels, it.name); + append(*all_labels, "\n"); + } + entry.find_labels = to_lower_copy(builder_to_string(*all_labels,, temp)); return entry, ""; } @@ -96,14 +128,16 @@ catalog_section :: (entry: *Catalog_Entry, model: string) -> *Mcap_Ocr_Section { // ---------------------------------------------------------------- search -Search_Field :: enum u8 { ANY; OCR; TITLE; URL; APP; NOTE; } +// LIKE is by meaning: the captures whose image is nearest the phrase, ranked by the label model +// (app/background.jai does the ranking); here it only asks for a labelled capture. +Search_Field :: enum u8 { ANY; OCR; TITLE; URL; APP; NOTE; LABEL; LIKE; } Search_Term :: struct { field: Search_Field; text: string; // lowercased } -// Words, "quoted phrases", and either with a field prefix: ocr: title: url: app: note:. Temp. +// Words, "quoted phrases", and either with a field prefix: ocr: title: url: app: note: label: like:. Temp. search_parse :: (query: string) -> [..] Search_Term { terms: [..] Search_Term; terms.allocator = temp; @@ -121,6 +155,8 @@ search_parse :: (query: string) -> [..] Search_Term { case "url"; field = .URL; case "app"; field = .APP; case "note"; field = .NOTE; + case "label"; field = .LABEL; + case "like"; field = .LIKE; } if field != .ANY { term.field = field; rest = advance(rest, colon + 1); } // "http:..." and the like stay words } @@ -148,12 +184,14 @@ search_matches :: (entry: *Catalog_Entry, terms: [] Search_Term) -> bool { found: bool; if it.field == { case .ANY; found = contains(entry.find_ocr, it.text) || contains(entry.find_title, it.text) || contains(entry.find_url, it.text) || - contains(entry.find_app, it.text) || contains(entry.find_notes, it.text); + contains(entry.find_app, it.text) || contains(entry.find_notes, it.text) || contains(entry.find_labels, it.text); case .OCR; found = contains(entry.find_ocr, it.text); case .TITLE; found = contains(entry.find_title, it.text); case .URL; found = contains(entry.find_url, it.text); case .APP; found = contains(entry.find_app, it.text); case .NOTE; found = contains(entry.find_notes, it.text); + case .LABEL; found = contains(entry.find_labels, it.text); + case .LIKE; found = entry.embedding.count > 0; } if !found return false; } diff --git a/src/app/labels_ui.jai b/src/app/labels_ui.jai new file mode 100644 index 0000000..94f24dc --- /dev/null +++ b/src/app/labels_ui.jai @@ -0,0 +1,371 @@ +// Labels in the app (modules/Labels): the label list, the label engine, the Labels window, and the +// labels row of the Capture window. +// +// The label list is labels.config in the app's data folder: a name and a prompt per label, the +// prompt being what the model scores (a fuller description scores better than a word). It is read +// again when the file changes. Each version of it becomes a Label_Set: its names, its prompts' +// vectors and a hash of it, made on the OCR thread (app/ocr_ui.jai) and never changed after, so a +// scan can hold the set it started with while the list is edited. +// +// The engine is made on the OCR thread too, once, after the OCR engine; it is destroyed only at +// shutdown, after the background workers (app/background.jai), which use its text model for like: +// searches, have stopped. + +LABELS_WINDOW :: "Labels"; +LABELS_MODEL :: "siglip2-base-patch16-384"; // the one label model; its chunks are keyed by this +LABELS_TOP :: 5; // labels stored per capture, best first +LABELS_LIST_FILE :: "labels.config"; + +Label_Entry :: struct { name, prompt: string; } + +// A version of the label list with its prompts' vectors. Heap; never changed once made. +Label_Set :: struct { + hash: string; + names: [] string; + vectors: [] float; // one unit vector per label, engine.dimensions long, one after another +} + +labels_ui: struct { + // Owned by the OCR thread while ocr_ui.thread_busy, like ocr_ui.engine. + engine: *Labels_Engine; + engine_error: string; + engine_tried: bool; + set: *Label_Set; // the newest list's + retired_sets: [..] *Label_Set; // replaced while a scan might hold them; freed when the OCR thread is idle + + // Main thread. + list: [..] Label_Entry; // as last read; heap + list_hash: string; // heap + list_error: string; + list_modtime: Apollo_Time; + list_checked_at: Apollo_Time; + window_open: bool; + scan_labels := true; // the library scan labels too + search_ready: bool; // the engine was up when the shown captures were last worked out + + // The Labels window's counts, worked out again only when the catalog or the list changes. + counted_generation: s64 = -1; + counted_hash: string; // heap + labelled_count, current_count: s64; +}; + +labels_models_folder :: () -> string { + return tprint("%label-models", path_strip_filename(get_path_of_running_executable())); +} + +labels_list_path :: () -> string { + return tprint("%/%", app_program_data_dir(), LABELS_LIST_FILE); +} + +// ---------------------------------------------------------------- startup and shutdown + +labels_ui_startup :: () { + labels_ui.list.allocator = context.default_allocator; + labels_ui.retired_sets.allocator = context.default_allocator; + path := labels_list_path(); + if !file_exists(path) { + if write_entire_file(path, LABELS_STARTER_LIST) log("[Labels] Wrote a starter label list to %.", path); + else log_error("[Labels] Could not write %.", path); + } + labels_read_list(); +} + +// After ocr_ui_shutdown has stopped the OCR thread, and viewer_shutdown the background workers. +labels_ui_shutdown :: () { + labels_destroy(labels_ui.engine); + labels_ui.engine = null; + free(labels_ui.counted_hash,, context.default_allocator); +} + +// Whether a scan can start with labels: the engine has been tried and, when it is up, the list's +// set is made. +labels_ui_settled :: () -> bool { + if ocr_ui.thread_busy || !labels_ui.engine_tried return false; + return !labels_ui.engine || (labels_ui.set && labels_ui.set.hash == labels_ui.list_hash); +} + +// ---------------------------------------------------------------- the label list + +LABELS_STARTER_LIST :: #string END +@version 1 +# The labels each capture is scored against. One per line: its name, then the prompt the model scores, +# both quoted. The model compares an image with the prompt, and fuller descriptions score better than +# single words. Saved changes are picked up while the app runs; captures labelled with an earlier list +# are labelled again from their stored vectors by the next scan, without reading their images again. + +label: "Code editor" "a screenshot of a code editor showing source code" +label: "Terminal" "a screenshot of a terminal window with command line text" +label: "Web page" "a screenshot of a web page in a browser" +label: "Search results" "a screenshot of search engine results" +label: "Documentation" "a screenshot of technical documentation with headings and paragraphs" +label: "Forum or social" "a screenshot of a forum thread or social media posts" +label: "Chat" "a screenshot of a chat conversation with messages" +label: "Email" "a screenshot of an email" +label: "Video call" "a screenshot of a video call with people's faces" +label: "Video" "a screenshot of a video player" +label: "Schematic" "an electronic circuit schematic diagram" +label: "PCB layout" "a printed circuit board layout in a PCB design tool" +label: "CAD model" "a 3D CAD model of a mechanical part" +label: "Plot" "a plot of a signal or data on axes" +label: "Chart" "a bar chart or pie chart" +label: "Dashboard" "a dashboard with graphs and numbers" +label: "Spreadsheet" "a spreadsheet with rows and columns of numbers" +label: "Table" "a table of data" +label: "Diagram" "a block diagram or flowchart" +label: "Map" "a map" +label: "Settings" "a settings dialog with options and checkboxes" +label: "File browser" "a file explorer window listing files and folders" +label: "System monitor" "a system monitor showing CPU temperatures and usage" +label: "Error message" "an error message dialog" +label: "Installer" "a software installer window" +label: "Shopping" "an online store page with products and prices" +label: "Datasheet" "a page of an electronic component datasheet" +label: "Paper" "a page of a scientific paper" +label: "Slide" "a presentation slide" +label: "Photo" "a photograph of a real scene" +label: "Product photo" "a photo of an electronic device or circuit board" +label: "Person" "a photo of a person" +label: "Game" "a screenshot of a video game" +label: "Code review" "a screenshot of a code diff or pull request" +END + +// Reads labels.config into labels_ui.list, its hash into list_hash. A line that cannot be read is +// logged and skipped, as in the Ex4 formats. +labels_read_list :: () { + push_allocator(context.default_allocator); + path := labels_list_path(); + modtime, _, found := file_modtime_and_size(path); + labels_ui.list_modtime = modtime; + for labels_ui.list { free(it.name); free(it.prompt); } + array_reset_keeping_memory(*labels_ui.list); + free(labels_ui.list_error); + labels_ui.list_error = ""; + + text, ok := read_entire_file(path, log_errors = false,, temp); + if !ok { labels_ui.list_error = sprint("Cannot read %.", path); labels_set_list_hash(); return; } + line_number := 0; + for raw: split(text, "\n",, temp) { + line_number += 1; + line := trim(raw); + if !line.count || line[0] == #char "#" || begins_with(line, "@version") continue; + if !begins_with(line, "label:") { log_error("[Labels] %:%: expected label: \"name\" \"prompt\". Ignored.", path, line_number); continue; } + rest := trim(advance(line, 6)); + name, after, name_ok := labels_quoted(rest); + prompt, end, prompt_ok := labels_quoted(trim_left(after)); + if !name_ok || !prompt_ok || trim(end).count || !name.count || !prompt.count { + log_error("[Labels] %:%: expected label: \"name\" \"prompt\". Ignored.", path, line_number); + continue; + } + array_add(*labels_ui.list, .{copy_string(name), copy_string(prompt)}); + } + labels_set_list_hash(); + log("[Labels] % labels from % (list %).", labels_ui.list.count, path, labels_ui.list_hash); +} + +// Each frame: reads the list again when its file has changed (looked at every 2 s). +labels_check_list :: () { + now := current_time_monotonic(); + if to_float64_seconds(now - labels_ui.list_checked_at) < 2 return; + labels_ui.list_checked_at = now; + modtime, _, found := file_modtime_and_size(labels_list_path()); + if found && modtime != labels_ui.list_modtime labels_read_list(); +} + +// A Label_Set for the list given, with engine: its prompts' vectors. Heap. On the OCR thread. +labels_make_set :: (engine: *Labels_Engine, names: [] string, prompts: [] string, hash: string) -> *Label_Set, error: string { + push_allocator(context.default_allocator); + vectors, error := labels_embed_texts(engine, prompts); + if error return null, error; + set := New(Label_Set); + set.hash = copy_string(hash); + set.names = NewArray(names.count, string); + for names set.names[it_index] = copy_string(it); + set.vectors = vectors; + return set, ""; +} + +label_set_free :: (set: *Label_Set) { + if !set return; + push_allocator(context.default_allocator); + free(set.hash); + for set.names free(it); + array_free(set.names); + array_free(set.vectors); + free(set); +} + +// The LABELS_TOP best labels of set for an image vector, best first. Temp. +labels_for_vector :: (engine: *Labels_Engine, set: *Label_Set, vector: *float) -> [..] Mcap_Label { + scored: [..] Mcap_Label; + scored.allocator = temp; + dims := engine.dimensions; + for set.names array_add(*scored, .{labels_score(engine, vector, set.vectors.data + it_index * dims), it}); + quick_sort(scored, (a: Mcap_Label, b: Mcap_Label) -> s64 { return ifx a.score > b.score then -1 else ifx a.score < b.score then 1 else 0; }); + if scored.count > LABELS_TOP scored.count = LABELS_TOP; + return scored; +} + +// ---------------------------------------------------------------- each frame + +// The engine and the list's set, made on the OCR thread when it is free; the list file watched. +labels_ui_frame :: () { using ImGui; + labels_check_list(); + // like: searches need the engine: work them out again once it is up. + if labels_ui.engine && !labels_ui.search_ready && !ocr_ui.thread_busy { + labels_ui.search_ready = true; + for viewer.search_terms if it.field == .LIKE viewer.visible_dirty = true; + } + if !ocr_ui.thread_busy { + for labels_ui.retired_sets label_set_free(it); + array_reset_keeping_memory(*labels_ui.retired_sets); + if !labels_ui.engine_tried && ocr_ui.engine_tried { // after the OCR engine + ocr_ui_start(.{kind = .LABELS_CREATE}); + } else if labels_ui.engine && (!labels_ui.set || labels_ui.set.hash != labels_ui.list_hash) && labels_ui.list.count { + ocr_ui_start(.{kind = .LABEL_SET}); + } + } + + if !labels_ui.window_open return; + SetNextWindowSize(.{640, 560}, .FirstUseEver); + if Begin(LABELS_WINDOW, *labels_ui.window_open) { + labels_window_contents(); + } + End(); +} + +// On the OCR thread: makes the engine (LABELS_CREATE) or the newest list's set (LABEL_SET), with +// the list the task was given. +labels_ui_task :: (task: Ocr_Task) { + if task.kind == .LABELS_CREATE { + engine, ok, error := labels_create(labels_models_folder(), LABELS_MODEL); + free(labels_ui.engine_error); + labels_ui.engine_error = ifx ok then "" else copy_string(error); + if ok log("[Labels] % on %.%", LABELS_MODEL, engine.provider, ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); + else log_error("[Labels] Could not start %: %", LABELS_MODEL, error); + labels_ui.engine = engine; + labels_ui.engine_tried = true; + } + if labels_ui.engine && task.label_names.count { + start := current_time_monotonic(); + set, error := labels_make_set(labels_ui.engine, task.label_names, task.label_prompts, task.label_hash); + if !set { log_error("[Labels] The label list's prompts could not be scored: %", error); return; } + if labels_ui.set array_add(*labels_ui.retired_sets, labels_ui.set); + labels_ui.set = set; + log("[Labels] % labels ready (list %) in % ms.", set.names.count, set.hash, FF(to_float64_seconds(current_time_monotonic() - start) * 1000, 0)); + } +} + +// ---------------------------------------------------------------- the Labels window + +#scope_file + +labels_window_contents :: () { using ImGui; + busy := ocr_ui.thread_busy; + if busy && ocr_ui.task.kind == .LABELS_CREATE { + TextColored(app_theme.palette.dim, "Starting % ...", LABELS_MODEL); + } else if labels_ui.engine_error { + TextColored(app_theme.palette.warning, "Labelling could not start:"); + TextWrapped("%", labels_ui.engine_error); + TextWrapped("Run `jai fetch_models.jai` in modules/Labels, then rebuild."); + } else if labels_ui.engine { + in_use, explanation := labels_gpu_status(labels_ui.engine); + Text("%: images on %, text on the CPU", labels_ui.engine.model.title, ifx in_use then "the GPU (CUDA)" else "the CPU"); + if explanation text_with_links(explanation); + } else if !labels_ui.engine_tried { + TextColored(app_theme.palette.dim, "Starts after OCR."); + } + + // How much of the library is labelled, and with which list: counted again only when the catalog + // or the list has changed, not every frame. + if labels_ui.counted_generation != viewer.catalog_generation || labels_ui.counted_hash != labels_ui.list_hash { + labels_ui.labelled_count, labels_ui.current_count = 0, 0; + for viewer.captures { + if it.removed || !it.catalog || !it.catalog.embedding.count continue; + labels_ui.labelled_count += 1; + if it.catalog.label_list == labels_ui.list_hash labels_ui.current_count += 1; + } + labels_ui.counted_generation = viewer.catalog_generation; + free(labels_ui.counted_hash,, context.default_allocator); + labels_ui.counted_hash = copy_string(labels_ui.list_hash,, context.default_allocator); + } + SeparatorText("The library"); + Text("% of % captures labelled; % with this list.", labels_ui.labelled_count, viewer_capture_count(), labels_ui.current_count); + TextColored(app_theme.palette.dim, "The OCR window's scan labels the rest, and labels again from their vectors those scored with an earlier list."); + if SmallButton("OCR window") ocr_ui.window_open = true; + + SeparatorText("The label list"); + TextUnformatted(labels_list_path()); + SameLine(); + if SmallButton("Open") platform_open_file_in_default_program(labels_list_path(), app_program_data_dir()); + if labels_ui.list_error TextColored(app_theme.palette.warning, "%", labels_ui.list_error); + ready := labels_ui.set && labels_ui.set.hash == labels_ui.list_hash; + TextColored(app_theme.palette.dim, "% labels, list %%", labels_ui.list.count, labels_ui.list_hash, ifx ready || !labels_ui.engine then "" else "; scoring its prompts..."); + if BeginTable("##list", 2, .RowBg | .SizingStretchProp | .ScrollY, .{0, -FLOAT32_MIN}) { + TableSetupColumn("Name", .WidthFixed); + TableSetupColumn("Prompt"); + TableSetupScrollFreeze(0, 1); + TableHeadersRow(); + for labels_ui.list { + TableNextRow(); + TableNextColumn(); TextUnformatted(it.name); + TableNextColumn(); TextUnformatted(it.prompt); + } + EndTable(); + } +} + +labels_set_list_hash :: () { + // FNV-1a (the Hash module's, 64-bit) over every name and prompt, each ended by a 0xFF byte, which + // UTF-8 text never has. + hash := FNV_64_OFFSET_BIAS; + for labels_ui.list { + hash = fnv1a_hash(it.name.data, it.name.count, hash); + hash = fnv1a_hash(0xFF, hash); + hash = fnv1a_hash(it.prompt.data, it.prompt.count, hash); + hash = fnv1a_hash(0xFF, hash); + } + free(labels_ui.list_hash,, context.default_allocator); + labels_ui.list_hash = sprint("%", formatInt(hash, base = 16, minimum_digits = 16),, context.default_allocator); +} + +// A quoted string at the start of s, "" for a quote: its text (temp) and what follows. +labels_quoted :: (s: string) -> text: string, rest: string, ok: bool { + if !s.count || s[0] != #char "\"" return "", s, false; + builder: String_Builder; + builder.allocator = temp; + i := 1; + while i < s.count { + if s[i] == #char "\"" { + if i + 1 < s.count && s[i + 1] == #char "\"" { append(*builder, #char "\""); i += 2; continue; } + return builder_to_string(*builder,, temp), advance(s, i + 1), true; + } + append(*builder, s[i]); + i += 1; + } + return "", s, false; +} + +#scope_export + +// ---------------------------------------------------------------- the Capture window + +// The capture's labels, under its OCR row: dim when they come from an earlier list. +labels_ui_capture_line :: (index: s32) { using ImGui; + entry := viewer.captures[index].catalog; + if !entry || !entry.labels.count { + if entry && entry.embedding.count TextColored(app_theme.palette.dim, "Labels: none yet with this list (the next scan adds them)"); + return; + } + current := entry.label_list == labels_ui.list_hash; + builder: String_Builder; + builder.allocator = temp; + append(*builder, "Labels: "); + for entry.labels { + if it_index append(*builder, " "); + print_to_builder(*builder, "% %", it.name, FF(it.score, 3)); + } + if !current append(*builder, " (from an earlier list)"); + if current TextUnformatted(builder_to_string(*builder,, temp)); + else TextColored(app_theme.palette.dim, "%", builder_to_string(*builder,, temp)); + if IsItemHovered() SetTooltip("Scores are the model's own (0 to 1): small, but their order is what counts. Click a label in Stats, or search label:NAME."); +} diff --git a/src/app/main.jai b/src/app/main.jai index fa54df2..1ddd689 100644 --- a/src/app/main.jai +++ b/src/app/main.jai @@ -20,6 +20,7 @@ main :: () { for app.args if it == "-library" && it_index + 1 < app.args.count app.library_folder = app.args[it_index + 1]; app.scan_and_exit = array_find(app.args, "-scan"); 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]; defer report_memory_leaks(); diff --git a/src/app/ocr_ui.jai b/src/app/ocr_ui.jai index 71917a7..b5357da 100644 --- a/src/app/ocr_ui.jai +++ b/src/app/ocr_ui.jai @@ -3,11 +3,13 @@ // and in the Capture window the OCR button, the lines read drawn over the image, and the search's // matches among them. // -// Every reading is stored in its .mcap (mcap_store_ocr: one section per model) and its catalog entry -// is loaded again, so the overlay, the text shown and search all come from what is stored. +// Every reading is stored in its .mcap (one OCR section per model) and its catalog entry is loaded +// again, so the overlay, the text shown and search all come from what is stored. A capture read for +// labels too (app/labels_ui.jai) gets its labels and image vector stored in the same rewrite. // -// Making an engine, reading one capture and the scan all run on one background thread, one at a time; -// the UI only looks at what it has finished, and at the scan's progress counters. +// Making the engines (OCR, then labels), the label list's vectors, reading one capture and the scan +// all run on one background thread, one at a time; the UI only looks at what it has finished, and +// at the scan's progress counters. OCR_WINDOW :: "OCR"; OCR_DEFAULT_MODEL :: "ppocr-v5-mobile"; @@ -22,6 +24,7 @@ ocr_ui: struct { // Owned by the main thread while thread_busy is false; by the background thread while it is true. engine: *Ocr_Engine; engine_error: string; // why the last engine could not be made + engine_tried: bool; // the first engine has been made, or failed: labels start after it thread: Thread; thread_started: bool; thread_busy: bool; // set by the main thread, cleared by the background one @@ -45,13 +48,16 @@ ocr_ui: struct { toast_shown := false; } -Ocr_Task_Kind :: enum { NONE; CREATE; READ; SCAN; } +Ocr_Task_Kind :: enum { NONE; CREATE; READ; SCAN; LABELS_CREATE; LABEL_SET; } Ocr_Task :: struct { kind: Ocr_Task_Kind; model: string; // CREATE path: string; // READ index: s32; // READ + // LABELS_CREATE and LABEL_SET: the label list as it was when the task started (heap). + label_names, label_prompts: [] string; + label_hash: string; } // The library scan. The counters and flags are shared with the background thread. @@ -68,9 +74,17 @@ Ocr_Scan :: struct { paused_seconds: float64; // time spent paused, left out of the rate 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 } -Scan_Item :: struct { index: s32; path: string; } +// A capture to read, and what it needs. +Scan_Item :: struct { + index: s32; + path: string; + ocr: bool; // read with the OCR engine + embed: bool; // its image vector made, and its labels + relabel: bool; // only its labels, from its stored vector (labelled with an earlier list) +} ocr_models_folder :: () -> string { return tprint("%ocr-models", path_strip_filename(get_path_of_running_executable())); @@ -132,6 +146,15 @@ ocr_ui_start :: (task: Ocr_Task) -> bool { ocr_ui.task = task; ocr_ui.task.model = copy_string(task.model); ocr_ui.task.path = copy_string(task.path); + if task.kind == .LABELS_CREATE || task.kind == .LABEL_SET { + ocr_ui.task.label_names = NewArray(labels_ui.list.count, string); + ocr_ui.task.label_prompts = NewArray(labels_ui.list.count, string); + for labels_ui.list { + ocr_ui.task.label_names[it_index] = copy_string(it.name); + ocr_ui.task.label_prompts[it_index] = copy_string(it.prompt); + } + ocr_ui.task.label_hash = copy_string(labels_ui.list_hash); + } ocr_ui.thread_busy = true; if ocr_ui.thread_started thread_deinit(*ocr_ui.thread); ocr_ui.thread_started = thread_init(*ocr_ui.thread, ocr_ui_thread_proc, temporary_storage_size = 4 * 1024 * 1024, name = "ocr"); @@ -155,11 +178,17 @@ ocr_ui_thread_proc :: (thread: *Thread) -> s64 { ocr_ui.engine_error = ifx ok then "" else copy_string(error); if ok log("[OCR] % on %.%", task.model, engine.provider, ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else ""); else log_error("[OCR] Could not start %: %", task.model, error); + ocr_ui.engine_tried = true; + + case .LABELS_CREATE; #through; + case .LABEL_SET; + labels_ui_task(task); case .READ; start := current_time_monotonic(); - items := Scan_Item.[.{task.index, task.path}]; - errors := ocr_read_and_store(items); + set := ifx labels_ui.engine then labels_ui.set else null; + items := Scan_Item.[.{index = task.index, path = task.path, ocr = true, embed = set != null}]; + errors := ocr_read_and_store(items, set); free(ocr_ui.read_error); ocr_ui.read_error = ifx errors[0] then errors[0] else ""; ocr_ui.read_index = task.index; @@ -172,59 +201,110 @@ ocr_ui_thread_proc :: (thread: *Thread) -> s64 { while scan.paused && !scan.cancel sleep_milliseconds(100); if scan.cancel break; count := min(OCR_SCAN_BATCH, scan.items.count - position); - errors := ocr_read_and_store(array_view(scan.items, position, count)); + errors := ocr_read_and_store(array_view(scan.items, position, count), scan.label_set); for errors if it { atomic_add(*scan.failed, 1); log_error("[OCR 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("[OCR scan] % of % captures read with % in % (% failed)%.", scan.done, scan.total, scan.model, ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx scan.cancel then ", cancelled" else ""); + log("[OCR scan] % of % captures read with %% in % (% failed)%.", scan.done, scan.total, scan.model, ifx scan.label_set then tprint(" and labelled with %", LABELS_MODEL) else "", + ocr_duration_text(scan.elapsed_seconds), scan.failed, ifx scan.cancel then ", cancelled" else ""); atomic_swap(*scan.running, false); } free(task.model); free(task.path); + for task.label_names free(it); + for task.label_prompts free(it); + array_free(task.label_names); + array_free(task.label_prompts); + free(task.label_hash); atomic_swap(*ocr_ui.thread_busy, false); return 0; } -// Reads items with the engine, stores each reading in its file, and hands each new catalog entry to -// the viewer. An error message (heap) for each item that failed, "" for the others. Background thread. -ocr_read_and_store :: (items: [] Scan_Item) -> [] string { +// Where value is in list, or -1. +scan_index_of :: (list: [] s64, value: s64) -> s64 { + found, index := array_find(list, value); + return ifx found then index else -1; +} + +// Reads items, as each asks: with the OCR engine; with the label engine (labels scored against set); +// or only labels again, from the vector stored in the file. Stores what each got in its file in one +// rewrite, and hands each new catalog entry to the viewer. An error message (heap) for each item that +// failed, "" for the others. Background thread. +ocr_read_and_store :: (items: [] Scan_Item, set: *Label_Set) -> [] string { errors := NewArray(items.count, string); if !ocr_ui.engine { for * errors it.* = copy_string("No OCR engine is running."); return errors; } + labels := labels_ui.engine; // The image the Capture window shows: the rendered one, annotations and all, or the original. - Prepared :: struct { pixels: *u8; width, height: s32; image: string; } + Prepared :: struct { pixels: *u8; width, height: s32; image: string; stored_vector: [] u16; } prepared := NewArray(items.count, Prepared,, temp); - images: [..] Ocr_Image; - images.allocator = temp; - image_of: [..] s64; // which item each image is - image_of.allocator = temp; + ocr_images: [..] Ocr_Image; + ocr_images.allocator = temp; + ocr_of: [..] s64; // which item each OCR image is + ocr_of.allocator = temp; + label_images: [..] Labels_Image; + label_images.allocator = temp; + label_of: [..] s64; // which item each image to label is + label_of.allocator = temp; for item, n: items { data, read_ok := read_entire_file(item.path, log_errors = false,, temp); if !read_ok { errors[n] = copy_string("cannot be read"); continue; } m, ok, error := mcap_parse(cast([] u8) data); if !ok { errors[n] = copy_string(error); continue; } + if item.relabel { + chunk := mcap_find_chunk(m.front_chunks, mcap_tag(MCAP_EMBEDDING_TAG), LABELS_MODEL); + if chunk prepared[n].stored_vector = mcap_embedding_parse(chunk.data,, temp); + } + if !item.ocr && !item.embed continue; png := mcap_rendered(m); kind := "rendered"; if !png.count { png = m.original; kind = "original"; } w, h, channels: s32; pixels := stbi_load_from_memory(png.data, cast(s32) png.count, *w, *h, *channels, 4); if !pixels { errors[n] = copy_string("its image cannot be decoded"); continue; } - prepared[n] = .{pixels, w, h, kind}; - array_add(*images, .{pixels, w, h, w * 4}); - array_add(*image_of, n); + prepared[n] = .{pixels, w, h, kind, prepared[n].stored_vector}; + if item.ocr { array_add(*ocr_images, .{pixels, w, h, w * 4}); array_add(*ocr_of, n); } + if item.embed && labels && set { array_add(*label_images, .{pixels, w, h, w * 4}); array_add(*label_of, n); } } defer for prepared if it.pixels stbi_image_free(it.pixels); - results := ocr_process(ocr_ui.engine, images,, temp); - engine_text := tprint("onnxruntime %, %", ort.version, ifx ocr_ui.engine.provider == .CUDA then "CUDA" else "CPU"); - for result, r: results { - n := image_of[r]; - if result.failed { errors[n] = copy_string(result.error); continue; } + // Labels: vectors for the images to label, then each item's labels. + Labelled :: struct { done: bool; labels: Mcap_Labels; vector: [] u16; } + labelled := NewArray(items.count, Labelled,, temp); + if labels && set { + vectors, label_error := labels_embed_images(labels, label_images,, temp); + if label_error for label_of if !errors[it] errors[it] = sprint("labelling failed: %", label_error); + dims := labels.dimensions; + label_engine := tprint("onnxruntime %, %", ort.version, ifx labels.provider == .CUDA then "CUDA" else "CPU"); + for n: 0..items.count - 1 { + if errors[n] continue; + vector: [] float; + at := scan_index_of(label_of, n); + if at >= 0 vector = array_view(vectors, at * dims, dims); + else if prepared[n].stored_vector.count == dims vector = labels_from_half(prepared[n].stored_vector,, temp); + if !vector.count continue; + l := *labelled[n]; + l.done = true; + l.labels = .{engine = label_engine, date = mcap_ocr_now(), list = set.hash}; + l.labels.labels = labels_for_vector(labels, set, vector.data); + if at >= 0 l.vector = labels_to_half(vector); + } + } + ocr_results := ocr_process(ocr_ui.engine, ocr_images,, temp); + engine_text := tprint("onnxruntime %, %", ort.version, ifx ocr_ui.engine.provider == .CUDA then "CUDA" else "CPU"); + for n: 0..items.count - 1 { + if errors[n] continue; section: Mcap_Ocr_Section; + has_section := false; + r := scan_index_of(ocr_of, n); + if r >= 0 { + result := ocr_results[r]; + if result.failed { errors[n] = copy_string(result.error); continue; } + has_section = true; section.model = ocr_ui.engine.model.name; section.engine = engine_text; section.image = prepared[n].image; @@ -239,11 +319,20 @@ ocr_read_and_store :: (items: [] Scan_Item) -> [] string { line.corners[it_index * 2 + 1] = cast(s32) round(corner.y); } } + } + + chunks: [..] Mcap_Chunk; + chunks.allocator = temp; + if labelled[n].done { + array_add(*chunks, mcap_labels_chunk(LABELS_MODEL, labelled[n].labels)); + if labelled[n].vector.count array_add(*chunks, mcap_embedding_chunk(LABELS_MODEL, labelled[n].vector)); + } + if !has_section && !chunks.count { errors[n] = copy_string("nothing to store (the label list is not ready)"); continue; } stored: bool; error: string; for 1..OCR_STORE_ATTEMPTS { - stored, error = mcap_store_ocr(items[n].path, section); + stored, error = mcap_store(items[n].path, ifx has_section then *section else null, chunks); if stored break; sleep_milliseconds(250); } @@ -272,7 +361,7 @@ ocr_ui_scan_and_exit :: (busy: bool) { if ocr_ui.models_error { finish(1, "[OCR scan] %", ocr_ui.models_error); return; } if viewer.scan_error { finish(1, "[OCR scan] %", viewer.scan_error); return; } if !ocr_ui.scan_requested { - if busy || !ocr_ui.engine || viewer.catalog_loaded < viewer.captures.count return; + if busy || !ocr_ui.engine || viewer.catalog_loaded < viewer.captures.count || !labels_ui_settled() return; ocr_ui.scan_requested = true; ocr_ui.window_open = true; // its progress bar ocr_scan_start(); @@ -300,15 +389,24 @@ ocr_scan_start :: () { scan.* = .{}; scan.items.allocator = context.default_allocator; scan.model = copy_string(ocr_ui.engine.model.name,, context.default_allocator); - // Every capture without a reading by this model, in the order the grid shows them. + if labels_ui.scan_labels && labels_ui.engine && labels_ui.set && labels_ui.set.hash == labels_ui.list_hash scan.label_set = labels_ui.set; + // Every capture this model has not read, or not labelled with this list, in the order the grid shows them. for viewer.visible { capture := *viewer.captures[it]; if capture.removed continue; - if capture.catalog && catalog_section(capture.catalog, scan.model) continue; - array_add(*scan.items, .{it, copy_string(capture.path,, context.default_allocator)}); + entry := capture.catalog; + item := Scan_Item.{index = it}; + item.ocr = !entry || !catalog_section(entry, scan.model); + if scan.label_set { + item.embed = !entry || !entry.embedding.count; + item.relabel = !item.embed && entry.label_list != scan.label_set.hash; + } + if !item.ocr && !item.embed && !item.relabel continue; + item.path = copy_string(capture.path,, context.default_allocator); + array_add(*scan.items, item); } scan.total = scan.items.count; - if !scan.total { toast_show(.Success, "Every capture shown has been read with %", scan.model); return; } + if !scan.total { toast_show(.Success, "Every capture shown has been read with %%", scan.model, ifx scan.label_set then " and labelled" else ""); return; } scan.running = true; scan.start = current_time_monotonic(); ocr_ui_start(.{kind = .SCAN}); @@ -326,8 +424,13 @@ ocr_scan_ui :: () { using ImGui; if Button(temp_c_string(label)) ocr_scan_start(); if !can_scan EndDisabled(); SameLine(); + labels_ready := labels_ui.engine && labels_ui.set && labels_ui.set.hash == labels_ui.list_hash; + if !labels_ready BeginDisabled(); + Checkbox("Labels too", *labels_ui.scan_labels); + if !labels_ready EndDisabled(); + if IsItemHovered(.AllowWhenDisabled) SetTooltip(ifx labels_ready then tprint("Also labels each capture with %, storing its labels and image vector", LABELS_MODEL) else "Labelling is not ready: see the Labels window"); if !catalog_ready TextColored(app_theme.palette.dim, "(once the catalog has loaded)"); - else TextColored(app_theme.palette.dim, "Reads every capture without a reading by this model. Each is stored as it is read, so a scan can stop and start again."); + 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 { 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); } diff --git a/src/app/stats_ui.jai b/src/app/stats_ui.jai index 6bac00d..c1b2f55 100644 --- a/src/app/stats_ui.jai +++ b/src/app/stats_ui.jai @@ -28,6 +28,8 @@ Stats_Result :: struct { years: [..] Stats_Year; // oldest first most_words: [..] Stats_Capture; // most first common_words: [..] Stats_Word; // commonest first + labelled: s64; // captures with labels (app/labels_ui.jai) + top_labels: [..] Stats_Word; // how many captures have each label as their best, most first } stats_ui: struct { @@ -60,6 +62,8 @@ stats_result_free :: (s: *Stats_Result) { free(s.model); for s.years free(it.year); for s.common_words free(it.word); + for s.top_labels free(it.word); + array_free(s.top_labels); array_free(s.years); array_free(s.most_words); array_free(s.common_words); @@ -74,8 +78,10 @@ stats_model :: () -> string { stats_ui_frame :: () { using ImGui; // -stats: once the catalog and the search are settled, at once, on this thread. + waiting_for_labels := false; + for viewer.search_terms if it.field == .LIKE && !labels_ui.search_ready waiting_for_labels = true; if app.stats_and_exit && viewer.captures.count && viewer.catalog_loaded >= viewer.captures.count && !ocr_ui.thread_busy - && !viewer.visible_dirty && !background.searching { + && !viewer.visible_dirty && !background.searching && !waiting_for_labels { items := background_snapshot_visible(); stats_result_arrived(background.stats_generation, stats_compute(items, stats_model())); array_free(items,, context.default_allocator); @@ -121,6 +127,21 @@ stats_log :: () { builder.allocator = temp; for s.common_words print_to_builder(*builder, "% %, ", it.word, it.count); log("[stats] commonest: %", builder_to_string(*builder,, temp)); + labels: String_Builder; + labels.allocator = temp; + for s.top_labels print_to_builder(*labels, "% %, ", it.word, it.count); + log("[stats] % labelled; best labels: %", s.labelled, builder_to_string(*labels,, temp)); + if viewer.search_terms.count { + shown: String_Builder; + shown.allocator = temp; + for viewer.visible { + if it_index >= 10 break; + capture := *viewer.captures[it]; + print_to_builder(*shown, "%%", ifx it_index then ", " else "", capture.name); + if capture.catalog && capture.catalog.labels.count print_to_builder(*shown, " (%)", capture.catalog.labels[0].name); + } + log("[stats] search \"%\": % shown; first: %", to_string(viewer.search_buffer.data), viewer.visible.count, builder_to_string(*shown,, temp)); + } } stats_window_contents :: () { using ImGui; @@ -197,6 +218,26 @@ stats_window_contents :: () { using ImGui; if Selectable(temp_c_string(tprint("% words % %##most%", stats_number(it.words), date, capture.name, it_index))) viewer_request_full(it.index); } + SeparatorText("Top labels"); + if !s.labelled { + TextColored(app_theme.palette.dim, "None of these captures is labelled yet: the OCR window's scan labels them (Labels too)."); + } else { + TextColored(app_theme.palette.dim, "Each capture's best label, of % labelled. Click one to search for it.", stats_number(s.labelled)); + if BeginTable("##labels", 3, .SizingFixedFit) { + TableSetupColumn("label"); + TableSetupColumn("captures"); + TableSetupColumn("bar", .WidthStretch); + for s.top_labels { + TableNextRow(); + TableNextColumn(); + if SmallButton(temp_c_string(tprint("%##label%", it.word, it_index))) viewer_set_search(tprint("label:\"%\"", it.word)); + TableNextColumn(); TextUnformatted(stats_number(it.count)); + TableNextColumn(); ProgressBar(cast(float) it.count / max(s.labelled, 1), .{-FLOAT32_MIN, 0}, temp_c_string(stats_percent(it.count, s.labelled))); + } + EndTable(); + } + } + SeparatorText("Commonest words"); TextColored(app_theme.palette.dim, "Words of 3 letters or more, without the commonest English ones. Click one to search for it."); right_edge := GetCursorScreenPos().x + GetContentRegionAvail().x; @@ -224,6 +265,9 @@ stats_compute :: (items: [] Search_Item, model: string) -> *Stats_Result { s.years.allocator = context.default_allocator; s.most_words.allocator = context.default_allocator; s.common_words.allocator = context.default_allocator; + s.top_labels.allocator = context.default_allocator; + top: Table(string, s64); // keys borrow the entries' label names + top.allocator = temp; per_capture: [..] Stats_Capture; per_capture.allocator = temp; @@ -240,6 +284,11 @@ stats_compute :: (items: [] Search_Item, model: string) -> *Stats_Result { year := find_or_add(*years, year_text); year.year = year_text; year.captures += 1; + if entry && entry.labels.count { + s.labelled += 1; + count := find_or_add(*top, entry.labels[0].name); + count.* += 1; + } section := catalog_section(entry, model); if !section continue; @@ -286,6 +335,9 @@ stats_compute :: (items: [] Search_Item, model: string) -> *Stats_Result { array_add(*s.common_words, .{copy_string(it.word,, context.default_allocator), it.count}); } + for count, name: top array_add(*s.top_labels, .{copy_string(name,, context.default_allocator), count}); + quick_sort(s.top_labels, (a: Stats_Word, b: Stats_Word) -> s64 { return ifx a.count != b.count then b.count - a.count else compare(a.word, b.word); }); + s.compute_ms = to_float64_seconds(current_time_monotonic() - start) * 1000; return s; } diff --git a/src/app/viewer.jai b/src/app/viewer.jai index 500b89e..aa8851c 100644 --- a/src/app/viewer.jai +++ b/src/app/viewer.jai @@ -167,6 +167,7 @@ viewer_startup :: () { viewer.workers_started += 1; } background_startup(); + if app.start_search viewer_set_search(app.start_search); } // On the shutdown thread: stops the workers. The textures go with the device. @@ -398,18 +399,23 @@ viewer_apply_catalog_updates :: () { more_to_load := viewer.catalog_next < viewer.captures.count; unlock(*viewer.mutex); + // Whether the date order can change: a capture's first entry (one a refresh added), one lost, or a + // changed date. An entry read again after the scan stored a reading keeps its place, so a scan + // does not set off a search of the whole library for each capture it stores. + order_changes := false; for updates { capture := *viewer.captures[it.index]; if !capture.catalog viewer.catalog_loaded += 1; + if !capture.catalog || !it.entry || it.entry.date != capture.catalog.date order_changes = true; if it.entry || !capture.catalog { background_retire(capture.catalog); // a search or Stats job may be reading it capture.catalog = it.entry; } } if updates.count viewer.catalog_generation += 1; - // Under a search, or once sorted by date, entries arriving can move captures (after a refresh, - // the new ones get their dates). - if updates.count && (viewer.search_terms.count || viewer.sorted_by_date) viewer.visible_dirty = true; + // Under a search, any entry can start or stop matching (a scan adds OCR text and labels); + // background_search_start keeps those searches to a few a second. + if updates.count && (viewer.search_terms.count || (viewer.sorted_by_date && order_changes)) viewer.visible_dirty = true; if !viewer.sorted_by_date && viewer.catalog_loaded >= viewer.captures.count { viewer.sorted_by_date = true; viewer.visible_dirty = true; @@ -766,6 +772,8 @@ viewer_ui :: () { using ImGui; SameLine(); if Button("Stats") stats_ui.window_open = !stats_ui.window_open; // app/stats_ui.jai SameLine(); + if Button("Labels") labels_ui.window_open = !labels_ui.window_open; // app/labels_ui.jai + SameLine(); busy := ocr_ui.thread_busy; // a scan or a reading refers to captures by index if busy BeginDisabled(); if Button("Refresh") viewer_refresh(); @@ -935,6 +943,7 @@ viewer_full_window :: () { using ImGui; SameLine(); if SmallButton("100%") { viewer.full_fit = .OFF; viewer.full_zoom = 100; } ocr_ui_capture_controls(viewer.full_index, capture.path); // app/ocr_ui.jai + labels_ui_capture_line(viewer.full_index); // app/labels_ui.jai // Left and Right step through the captures in grid order while this window has focus. if IsWindowFocused(.RootAndChildWindows) { diff --git a/src/test/catalog_tests.jai b/src/test/catalog_tests.jai index d09ffc6..ec33330 100644 --- a/src/test/catalog_tests.jai +++ b/src/test/catalog_tests.jai @@ -58,6 +58,10 @@ query_tests :: () { expect(terms[3].field == .NOTE && terms[3].text == "late", "note: %", terms[3]); expect(terms[4].field == .ANY && terms[4].text == "http://x.y", "an unknown prefix stays part of the word: %", terms[4]); + labels := search_parse("label:schematic like:\"a circuit board\""); + expect(labels.count == 2 && labels[0].field == .LABEL && labels[0].text == "schematic" && labels[1].field == .LIKE && labels[1].text == "a circuit board", + "label: and like: prefixes: %", labels); + phrase := search_parse("title:\"two words\""); expect(phrase.count == 1 && phrase[0].field == .TITLE && phrase[0].text == "two words", "a prefixed phrase: %", phrase); expect(search_parse(" ").count == 0, "blank is no terms"); @@ -81,6 +85,16 @@ match_tests :: () { expect(search_matches(*entry, search_parse("")), "no terms match everything"); expect(!search_matches(null, search_parse("x")), "an entry not loaded matches no search"); + entry.find_labels = "schematic\npcb layout\n"; + expect(search_matches(*entry, search_parse("label:schematic")), "a label word"); + expect(search_matches(*entry, search_parse("schematic")), "a plain word matches labels too"); + expect(!search_matches(*entry, search_parse("ocr:schematic")), "an OCR word is not a label's"); + expect(!search_matches(*entry, search_parse("like:anything")), "like: needs an image vector"); + half := u16.[0x3C00]; + entry.embedding = half; + expect(search_matches(*entry, search_parse("like:anything")), "like: takes a labelled capture (the ranking is the background's)"); + entry.embedding = .[]; + expect(search_line_matches("0.1uF : Capacitor_SMD", search_parse("capacitor")), "a line holding a word"); expect(!search_line_matches("0.1uF : Capacitor_SMD", search_parse("title:capacitor")), "a title word is not an OCR line's"); } diff --git a/src/test/labels_tests.jai b/src/test/labels_tests.jai new file mode 100644 index 0000000..07a3f69 --- /dev/null +++ b/src/test/labels_tests.jai @@ -0,0 +1,133 @@ +// Labels (modules/Labels): the tokenizer's ids for known phrases (from Hugging Face's tokenizers on +// the same tokenizer.json), unit vectors from both encoders, and the model ranking a drawn text image +// as text rather than a photo. Needs the label models beside the program (the build copies them); +// without them, the suite says so and passes. + +#scope_file + +checks_run := 0; +checks_failed := 0; + +MAX_GROUPS :: 8; +group_names : [MAX_GROUPS] string; +group_checks : [MAX_GROUPS] int; +group_fails : [MAX_GROUPS] int; +group_count := 0; + +output_builder: *String_Builder; + +emit :: (format_string: string, args: ..Any) { + if output_builder print_to_builder(output_builder, format_string, ..args); + else print(format_string, ..args); +} @PrintLike + +group :: (name: string) { + assert(group_count < MAX_GROUPS, "raise MAX_GROUPS"); + group_names[group_count] = name; + group_count += 1; + emit("\n-- % --\n", name); +} + +expect :: (condition: bool, format: string, args: ..Any, loc := #caller_location) { + checks_run += 1; + if group_count group_checks[group_count-1] += 1; + if condition return; + + checks_failed += 1; + if group_count group_fails[group_count-1] += 1; + emit(" FAIL [%:%] %\n", group_names[max(group_count-1, 0)], loc.line_number, tprint(format, ..args)); +} + +report :: () -> bool { + emit("\n"); + for 0..group_count-1 { + if group_fails[it] emit(" % -- % checks, % FAILED\n", group_names[it], group_checks[it], group_fails[it]); + else emit(" % -- % checks\n", group_names[it], group_checks[it]); + } + if checks_failed emit("\n% OF % LABELS CHECKS FAILED\n", checks_failed, checks_run); + else emit("\nALL LABELS TESTS PASSED (% checks over % groups)\n", checks_run, group_count); + return checks_failed == 0; +} + +// From Hugging Face's tokenizers 0.23.2 on the same tokenizer.json, lowercased, cut and padded to 64. +Known_Tokens :: struct { text: string; ids: [] s64; } +KNOWN :: Known_Tokens.[ + .{"code editor", .[2564, 10476, 1]}, + .{"A screenshot of a terminal", .[235250, 13703, 576, 476, 17145, 1]}, + .{"", .[1]}, +]; + +// labels_float_to_half is a copy of the Float16 module's float_to_half_fast3 marked #no_aoc: the +// same bits for every value (here in a debug build, where the module's own is not stopped). +half_tests :: () { + group("half floats"); + values: [..] float; + values.allocator = temp; + array_add(*values, ..float.[0, -0.0, 1, -1, 0.5, 65504, 65520, 1000000, -1000000, cast(float) 0.00006103515625, cast(float) 0.0000001, FLOAT32_INFINITY, -FLOAT32_INFINITY]); + x: float = 0.00001; + for 1..10_000 { array_add(*values, x, -x * 0.37); x *= 1.0027; } + differ := 0; + for values if labels_float_to_half(it) != cast(u16) float32_to_float16(it) differ += 1; + expect(differ == 0, "% of % values convert to different bits", differ, values.count); +} + +labels_tests :: () { + half_tests(); + folder := labels_models_folder(); + if !file_exists(tprint("%/models.config", folder)) { + emit(" (skipped: no label models in %)\n", folder); + return; + } + group("engine"); + engine, created, error := labels_create(folder, LABELS_MODEL); + expect(created, "the label engine starts: %", error); + if !created return; + defer labels_destroy(engine); + emit(" images on %; % dimensions\n", ifx engine.provider == .CUDA then "CUDA" else "the CPU", engine.dimensions); + + group("tokenizer"); + for KNOWN { + ids := tokenizer_encode(engine.tokenizer, it.text, 64); + same := ids.count == 64; + for known, i: it.ids if !same || ids[i] != known same = false; + for it.ids.count..63 if same && ids[it] != 0 same = false; // + expect(same, "\"%\" tokenizes as Hugging Face's does: % (got %)", it.text, it.ids, array_view(ids, 0, it.ids.count + 1)); + } + + group("vectors"); + prompts := string.["a screenshot of text on a white background", "a photo of a cat", "a photo of a mountain landscape"]; + texts, text_error := labels_embed_texts(engine, prompts,, temp); + expect(!text_error && texts.count == prompts.count * engine.dimensions, "the prompts' vectors are made (%)", text_error); + if text_error || texts.count != prompts.count * engine.dimensions return; + for 0..prompts.count - 1 expect(abs(vector_length(array_view(texts, it * engine.dimensions, engine.dimensions)) - 1) < 0.001, "the vector of \"%\" is a unit vector", prompts[it]); + + drawn, drawn_ok := test_drawn_text(); + if !drawn_ok { + emit(" (the text image is drawn with GDI, on Windows only: skipped)\n"); + return; + } + image := Labels_Image.{drawn.pixels, drawn.width, drawn.height, drawn.stride}; + vector, image_error := labels_embed_images(engine, .[image],, temp); + expect(!image_error && vector.count == engine.dimensions, "the image's vector is made (%)", image_error); + if image_error || vector.count != engine.dimensions return; + expect(abs(vector_length(vector) - 1) < 0.001, "the image's vector is a unit vector"); + scores: [3] float; + for 0..2 scores[it] = labels_score(engine, vector.data, texts.data + it * engine.dimensions); + expect(scores[0] > scores[1] && scores[0] > scores[2], "drawn text scores as text (%) above a cat (%) and a landscape (%)", scores[0], scores[1], scores[2]); +} + +vector_length :: (v: [] float) -> float { + sum: float; + for v sum += it * it; + return sqrt(sum); +} + +#scope_export + +run_labels_tests :: (builder: *String_Builder = null) -> bool { + output_builder = builder; + emit("=== labels test suite ===\n"); + push_allocator(context.default_allocator); + labels_tests(); + return report(); +} diff --git a/src/test/ocr_tests.jai b/src/test/ocr_tests.jai index 25e7aec..ed36589 100644 --- a/src/test/ocr_tests.jai +++ b/src/test/ocr_tests.jai @@ -204,6 +204,16 @@ turn_tests :: () { #scope_export +// LINES drawn black on white, as these tests read them; for other suites (labels_tests.jai). Temp. +test_drawn_text :: () -> Ocr_Image, ok: bool { + #if OS == .WINDOWS { + image, ok := drawn_lines(); + return image, ok; + } else { + return .{}, false; + } +} + run_ocr_tests :: (builder: *String_Builder = null) -> bool { output_builder = builder; emit("=== OCR test suite ===\n"); diff --git a/src/test/testing.jai b/src/test/testing.jai index 0e826f9..69152b1 100644 --- a/src/test/testing.jai +++ b/src/test/testing.jai @@ -40,6 +40,7 @@ Test_Type :: enum s32 { logger_suite; catalog_suite; ocr_suite; + labels_suite; } testing_started :: () -> bool { @@ -82,6 +83,7 @@ run_all_tests :: () { add_suite_test(.logger_suite, "logger"); add_suite_test(.catalog_suite, "catalog"); add_suite_test(.ocr_suite, "ocr"); + add_suite_test(.labels_suite, "labels"); } // One work item per suite. Each suite keeps its tallies in file-scope globals, so the whole @@ -190,6 +192,8 @@ test_thread_group_proc :: (group: *Thread_Group, thread: *Thread, work: *void) - task.pass = run_catalog_tests(task.output); case .ocr_suite; task.pass = run_ocr_tests(task.output); + case .labels_suite; + task.pass = run_labels_tests(task.output); } return .CONTINUE;