257 lines
10 KiB
Plaintext
257 lines
10 KiB
Plaintext
// OCR in the app: an engine made on a background thread at startup (modules/OCR), an OCR window that
|
|
// says whether it runs on the GPU and, if not, what to install, and the OCR button in the Capture
|
|
// window, which reads the capture shown and draws its lines over it.
|
|
//
|
|
// Engines take a moment to make and OCR takes a moment to run, so both happen on one background
|
|
// thread; the UI only looks at what it has finished.
|
|
|
|
OCR_WINDOW :: "OCR";
|
|
OCR_DEFAULT_MODEL :: "ppocr-v5-mobile";
|
|
|
|
ocr_ui: struct {
|
|
models: [] Ocr_Model_Info;
|
|
models_error: string;
|
|
model_index: s32; // the model chosen in the window
|
|
|
|
// Owned by the main thread while thread_busy is false; by the background thread while it is true.
|
|
engine: *Ocr_Engine;
|
|
engine_model: string; // the model engine runs
|
|
engine_error: string; // why the last engine could not be made
|
|
thread: Thread;
|
|
thread_started: bool;
|
|
thread_busy: bool; // atomic: set by the main thread, cleared by the background one
|
|
task: Ocr_Task;
|
|
|
|
// The last capture read, for the Capture window.
|
|
lines: [] Ocr_Line;
|
|
lines_index: s32 = -1; // the capture they are for
|
|
lines_error: string;
|
|
lines_seconds: float64;
|
|
show_boxes := true;
|
|
|
|
window_open := false;
|
|
toast_shown := false;
|
|
}
|
|
|
|
Ocr_Task_Kind :: enum { NONE; CREATE; READ; }
|
|
|
|
Ocr_Task :: struct {
|
|
kind: Ocr_Task_Kind;
|
|
model: string;
|
|
path: string; // READ: the .mcap
|
|
index: s32; // READ: its capture index
|
|
}
|
|
|
|
ocr_models_folder :: () -> string {
|
|
return tprint("%ocr-models", path_strip_filename(get_path_of_running_executable()));
|
|
}
|
|
|
|
ocr_ui_startup :: () {
|
|
push_allocator(context.default_allocator);
|
|
models, error := ocr_list_models(ocr_models_folder());
|
|
ocr_ui.models = models;
|
|
if error ocr_ui.models_error = copy_string(error);
|
|
for ocr_ui.models if it.name == OCR_DEFAULT_MODEL ocr_ui.model_index = cast(s32) it_index;
|
|
if ocr_ui.models.count ocr_ui_start(.{kind = .CREATE, model = ocr_ui.models[ocr_ui.model_index].name});
|
|
}
|
|
|
|
ocr_ui_shutdown :: () {
|
|
if ocr_ui.thread_started thread_is_done(*ocr_ui.thread, -1); // finishes what it is doing
|
|
ocr_destroy(ocr_ui.engine);
|
|
}
|
|
|
|
// Starts task on the background thread; false while it is still busy with the last one.
|
|
ocr_ui_start :: (task: Ocr_Task) -> bool {
|
|
if ocr_ui.thread_busy return false;
|
|
if ocr_ui.thread_started thread_is_done(*ocr_ui.thread, -1); // finished: thread_busy is clear
|
|
push_allocator(context.default_allocator);
|
|
ocr_ui.task = task;
|
|
ocr_ui.task.model = copy_string(task.model);
|
|
ocr_ui.task.path = copy_string(task.path);
|
|
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");
|
|
if !ocr_ui.thread_started { ocr_ui.thread_busy = false; return false; }
|
|
thread_start(*ocr_ui.thread);
|
|
return true;
|
|
}
|
|
|
|
ocr_ui_thread_proc :: (thread: *Thread) -> s64 {
|
|
context.allocator = context.default_allocator;
|
|
task := ocr_ui.task;
|
|
if task.kind == {
|
|
case .CREATE;
|
|
ocr_destroy(ocr_ui.engine);
|
|
ocr_ui.engine = null;
|
|
engine, ok, error := ocr_create(ocr_models_folder(), task.model);
|
|
ocr_ui.engine = engine;
|
|
ocr_ui.engine_model = task.model;
|
|
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);
|
|
|
|
case .READ;
|
|
start := current_time_monotonic();
|
|
lines, error := ocr_ui_read_capture(task.path);
|
|
ocr_ui.lines = lines;
|
|
ocr_ui.lines_error = error;
|
|
ocr_ui.lines_index = task.index;
|
|
ocr_ui.lines_seconds = to_float64_seconds(current_time_monotonic() - start);
|
|
}
|
|
free(task.model);
|
|
free(task.path);
|
|
atomic_swap(*ocr_ui.thread_busy, false);
|
|
return 0;
|
|
}
|
|
|
|
// The lines in a capture's image as the Capture window shows it: the rendered one, annotations and
|
|
// all, or the original when there is none. Heap.
|
|
ocr_ui_read_capture :: (path: string) -> [] Ocr_Line, error: string {
|
|
if !ocr_ui.engine return .[], copy_string("No OCR engine is running.");
|
|
data, read_ok := read_entire_file(path, log_errors = false,, temp);
|
|
if !read_ok return .[], sprint("Cannot read %.", path);
|
|
m, ok, error := mcap_parse(cast([] u8) data);
|
|
if !ok return .[], copy_string(error);
|
|
png := mcap_rendered(m);
|
|
if !png.count png = m.original;
|
|
w, h, channels: s32;
|
|
pixels := stbi_load_from_memory(png.data, cast(s32) png.count, *w, *h, *channels, 4);
|
|
if !pixels return .[], copy_string("Its image cannot be decoded.");
|
|
defer stbi_image_free(pixels);
|
|
images := Ocr_Image.[.{pixels, w, h, w * 4}];
|
|
results := ocr_process(ocr_ui.engine, images);
|
|
defer array_free(results);
|
|
if results[0].failed return .[], results[0].error;
|
|
return results[0].lines, "";
|
|
}
|
|
|
|
ocr_ui_free_lines :: () {
|
|
for ocr_ui.lines free(it.text);
|
|
array_free(ocr_ui.lines);
|
|
free(ocr_ui.lines_error);
|
|
ocr_ui.lines = .[];
|
|
ocr_ui.lines_error = "";
|
|
ocr_ui.lines_index = -1;
|
|
}
|
|
|
|
// ---------------------------------------------------------------- drawing
|
|
|
|
// Each frame: the toast once the engine is up on the CPU though a GPU could run it, and the OCR window.
|
|
ocr_ui_frame :: () { using ImGui;
|
|
busy := ocr_ui.thread_busy;
|
|
if !busy && ocr_ui.engine && !ocr_ui.toast_shown {
|
|
ocr_ui.toast_shown = true;
|
|
in_use, explanation := ocr_gpu_status(ocr_ui.engine);
|
|
if !in_use && contains(explanation, "can run it once") {
|
|
toast_show(.Error, "OCR runs on the CPU: CUDA libraries are missing. See the OCR window for what to install.");
|
|
ocr_ui.window_open = true;
|
|
}
|
|
}
|
|
|
|
if !ocr_ui.window_open return;
|
|
SetNextWindowSize(.{640, 360}, .FirstUseEver);
|
|
if Begin(OCR_WINDOW, *ocr_ui.window_open) {
|
|
if ocr_ui.models_error {
|
|
TextColored(app_theme.palette.warning, "%", ocr_ui.models_error);
|
|
TextWrapped("Run `jai fetch_models.jai` in modules/OCR and `jai fetch.jai` in modules/onnxruntime, then rebuild.");
|
|
} else {
|
|
// The model: changing it makes a new engine on the background thread.
|
|
preview := ifx ocr_ui.models.count then ocr_ui.models[ocr_ui.model_index].title else "";
|
|
SetNextItemWidth(320);
|
|
if BeginCombo("Model", temp_c_string(preview)) {
|
|
for ocr_ui.models {
|
|
if Selectable(temp_c_string(it.title), it_index == ocr_ui.model_index) && it_index != ocr_ui.model_index && !busy {
|
|
ocr_ui.model_index = cast(s32) it_index;
|
|
ocr_ui.toast_shown = true; // said once already, or not needed
|
|
ocr_ui_free_lines();
|
|
ocr_ui_start(.{kind = .CREATE, model = it.name});
|
|
}
|
|
}
|
|
EndCombo();
|
|
}
|
|
|
|
if busy && ocr_ui.task.kind == .CREATE {
|
|
TextColored(app_theme.palette.dim, "Starting % ...", ocr_ui.task.model);
|
|
} else if ocr_ui.engine_error {
|
|
TextColored(app_theme.palette.warning, "OCR could not start:");
|
|
TextWrapped("%", ocr_ui.engine_error);
|
|
} else if ocr_ui.engine && !busy {
|
|
in_use, explanation := ocr_gpu_status(ocr_ui.engine);
|
|
Text("Running on: %", ifx in_use then "the GPU (CUDA)" else "the CPU");
|
|
SameLine();
|
|
TextColored(app_theme.palette.dim, "onnxruntime %, % worker threads", ort.version, ocr_ui.engine.workers);
|
|
if explanation { Separator(); text_with_links(explanation); }
|
|
}
|
|
|
|
// The text of the capture last read, in reading order.
|
|
if !busy && ocr_ui.lines_index >= 0 && ocr_ui.lines_index < viewer.captures.count {
|
|
Separator();
|
|
Text("Text of %:", viewer.captures[ocr_ui.lines_index].name);
|
|
BeginChild("ocr text");
|
|
ocr_ui_capture_text(ocr_ui.lines_index);
|
|
EndChild();
|
|
}
|
|
}
|
|
}
|
|
End();
|
|
}
|
|
|
|
// The OCR row of the Capture window, above its image: the button, and how the last read went.
|
|
ocr_ui_capture_controls :: (index: s32, path: string) { using ImGui;
|
|
busy := ocr_ui.thread_busy;
|
|
can_read := ocr_ui.engine && !busy;
|
|
if !can_read BeginDisabled();
|
|
if SmallButton("OCR") {
|
|
ocr_ui_free_lines();
|
|
ocr_ui_start(.{kind = .READ, path = path, index = index});
|
|
}
|
|
if !can_read EndDisabled();
|
|
SameLine();
|
|
if busy && ocr_ui.task.kind == .READ && ocr_ui.task.index == index {
|
|
TextColored(app_theme.palette.dim, "Reading...");
|
|
} else if !busy && ocr_ui.lines_index == index {
|
|
if ocr_ui.lines_error TextColored(app_theme.palette.warning, "%", ocr_ui.lines_error);
|
|
else TextColored(app_theme.palette.dim, "% lines in % s", ocr_ui.lines.count, FF(ocr_ui.lines_seconds, 2));
|
|
SameLine();
|
|
Checkbox("Boxes", *ocr_ui.show_boxes);
|
|
} else if !ocr_ui.engine {
|
|
TextColored(app_theme.palette.dim, "(OCR is not running; see the OCR window)");
|
|
}
|
|
SameLine();
|
|
if SmallButton("OCR window") ocr_ui.window_open = true;
|
|
}
|
|
|
|
// The lines read from this capture, as boxes over its image, drawn at image_min with scale; and their
|
|
// text in a list below the image area.
|
|
ocr_ui_capture_overlay :: (index: s32, image_min: ImGui.ImVec2, scale: float) { using ImGui;
|
|
if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.show_boxes return;
|
|
draw := GetWindowDrawList();
|
|
colour := GetColorU32(.PlotHistogram);
|
|
for ocr_ui.lines {
|
|
p0 := ImVec2.{image_min.x + it.x * scale, image_min.y + it.y * scale};
|
|
p1 := ImVec2.{p0.x + it.width * scale, p0.y + it.height * scale};
|
|
draw.AddRect(draw, p0, p1, colour, 0, 1.5);
|
|
if IsMouseHoveringRect(p0, p1) SetTooltip("% (%)", it.text, FF(it.confidence, 2));
|
|
}
|
|
}
|
|
|
|
// Text a line at a time, each https:// address in it drawn as a link that opens in the browser.
|
|
text_with_links :: (text: string) { using ImGui;
|
|
for line: split(text, "\n",, temp) {
|
|
start := find_index_from_left(line, "https://");
|
|
if start < 0 { TextWrapped("%", line); continue; }
|
|
end := start;
|
|
while end < line.count && line[end] != #char " " && line[end] != #char ")" && line[end] != #char "," end += 1;
|
|
url := slice(line, start, end - start);
|
|
if start > 0 { TextUnformatted(slice(line, 0, start)); SameLine(0, 0); }
|
|
TextLinkOpenURL(temp_c_string(url));
|
|
if IsItemHovered() SetTooltip("Open % in the browser", url);
|
|
if end < line.count { SameLine(0, 0); TextUnformatted(advance(line, end)); }
|
|
}
|
|
}
|
|
|
|
ocr_ui_capture_text :: (index: s32) { using ImGui;
|
|
if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.lines.count return;
|
|
for ocr_ui.lines TextUnformatted(it.text);
|
|
}
|