MusaMahmood/mcap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

[0.0.0] OCR on the GPU: measured defaults for memory and batching; scan progress bars

Browse Source
This commit is contained in:
Musa Mahmood 2026-10-07 22:27:29 +01:00
parent ea9f055abc
commit e2334e8be5
8 changed files with 117 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
modules/OCR/README.md
View File

@ -51,8 +51,8 @@ as converted to ONNX by RapidOCR, and read Chinese, English and Japanese:
| Name | Files | On the 20 sample captures: RTX 5080 laptop / CPU |
| --- | --- | --- |
| `ppocr-v5-mobile` | detector 4.8 MB, recognizer 16.6 MB | 0.13 s / 0.5 to 0.65 s per capture |
| `ppocr-v5-server` | detector 88 MB, recognizer 85 MB | 2.3 s / 2.5 s per capture (0.36 s on the GPU with `rec_batch: 4`) |
| `ppocr-v5-mobile` | detector 4.8 MB, recognizer 16.6 MB | about 0.08 s / 0.26 s per capture |
| `ppocr-v5-server` | detector 88 MB, recognizer 85 MB | about 0.4 s / 2.4 s per capture |
Both also use the text line orientation classifier `ch_PP-LCNet_x0_25_textline_ori_cls_mobile.onnx`
(1 MB, `textline_ori:` in `models.config`).
@ -72,6 +72,13 @@ onnxruntime's CUDA build needs, on PATH or beside the executable:
when cuDNN is not on it, the newest `C:\Program Files\NVIDIA\CUDNN\v9.*\bin\13.*\x64` is added to the
program's own PATH, so an install in the usual place needs nothing more.
`Ocr_Options` holds how it runs on the GPU. Its defaults came from measuring: onnxruntime's
`kNextPowerOfTwo` arena, cuDNN's workspace capped, recognizer batches rounded to a width of 64, and 4
worker threads (with 8, the server model filled GPU memory and ran 5 times slower again). With the exact-size arena GPU memory filled up, as nearly every image and batch is a
new size, and the server model ran 5 times slower. Several copies of the sessions, each on its own
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
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.

10
modules/OCR/models.config
View File

@ -13,7 +13,9 @@
# box_threshold a text region is kept when its mean probability is at least this
# unclip_ratio how far each region is grown, as DB detection networks expect
# rec_height each text line is scaled to this height for recognition (px)
# rec_batch text lines recognized together in one run
# rec_batch_gpu text lines recognized together in one run, on the GPU
# rec_batch_cpu the same on the CPU, where one at a time is fastest: a batch is padded to its
# widest line, and the padding costs time and changes what is read a little
[model ppocr-v5-mobile]
title: "PP-OCRv5 mobile (light)"
@ -27,7 +29,8 @@ det_threshold: 0.3
box_threshold: 0.6
unclip_ratio: 1.5
rec_height: 48
rec_batch: 16
rec_batch_gpu: 16
rec_batch_cpu: 1
[model ppocr-v5-server]
title: "PP-OCRv5 server (heavy)"
@ -41,4 +44,5 @@ det_threshold: 0.3
box_threshold: 0.6
unclip_ratio: 1.5
rec_height: 48
rec_batch: 16
rec_batch_gpu: 8
rec_batch_cpu: 1

18
modules/OCR/models.jai
View File

@ -12,7 +12,8 @@ Ocr_Model_Info :: struct {
box_threshold: float = 0.6;
unclip_ratio: float = 1.5;
rec_height: s32 = 48;
rec_batch: s32 = 16;
rec_batch_gpu: s32 = 16;
rec_batch_cpu: s32 = 1;
}
OCR_MODELS_CONFIG :: "models.config";
@ -61,7 +62,8 @@ ocr_list_models :: (models_folder: string) -> [] Ocr_Model_Info, error: string {
case "box_threshold"; if number_ok model.box_threshold = number;
case "unclip_ratio"; if number_ok model.unclip_ratio = number;
case "rec_height"; if number_ok model.rec_height = cast(s32) number;
case "rec_batch"; if number_ok model.rec_batch = cast(s32) number;
case "rec_batch_gpu"; if number_ok model.rec_batch_gpu = cast(s32) number;
case "rec_batch_cpu"; if number_ok model.rec_batch_cpu = cast(s32) number;
case; log_error("[OCR] %:%: unknown key %. Ignored.", path, line_number, key);
}
}
@ -94,12 +96,12 @@ ocr_create_sessions :: (engine: *Ocr_Engine, models_folder: string, gpu: bool, i
model := *engine.model;
ok: bool;
error: string;
engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads);
engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads, engine.options);
if !ok return false, error;
engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads);
engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads, engine.options);
if !ok return false, error;
if model.textline_ori {
engine.orientation, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.textline_ori), gpu, intra_op_threads);
engine.orientation, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.textline_ori), gpu, intra_op_threads, engine.options);
if !ok return false, error;
// PP-LCNet text line orientation takes 160 by 80; a model that says otherwise is taken at its word.
shape := engine.orientation.input_shape;
@ -116,7 +118,7 @@ ocr_release_sessions :: (engine: *Ocr_Engine) {
}
// A session for the model file; with CUDA when gpu is set. On failure, the reason.
ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32) -> Ocr_Session, ok: bool, error: string {
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;
@ -134,8 +136,8 @@ ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads:
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"];
values := (*u8).["0", "HEURISTIC", "kSameAsRequested"];
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));

24
modules/OCR/module.jai
View File

@ -44,6 +44,8 @@ Ocr_Engine :: struct {
orientation_width, orientation_height: s32; // its input size
dictionary: [] string; // class n is dictionary[n - 1]
workers: s32; // threads ocr_process uses
rec_batch: s32; // text lines per recognizer run: the model's rec_batch_gpu or rec_batch_cpu
options: Ocr_Options;
stats: Ocr_Stats; // since creation, or the last ocr_reset_stats
}
@ -60,9 +62,20 @@ Ocr_Stats :: struct {
ocr_stats :: (engine: *Ocr_Engine) -> Ocr_Stats { return engine.stats; }
ocr_reset_stats :: (engine: *Ocr_Engine) { engine.stats = .{}; }
// How an engine runs on the GPU, beyond its model's models.config entry. The defaults are what ran
// fastest on the 20 sample captures (RTX 5080 laptop, onnxruntime 1.30.0): nearly every image and
// batch is a different size, and with onnxruntime's exact-size arena and cuDNN's largest workspaces
// GPU memory filled up and the server model ran several times slower.
Ocr_Options :: struct {
gpu_workers: s32 = 4; // threads preparing images and decoding results while the GPU runs the networks; more let the server model fill GPU memory
cuda_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)
rec_width_step: s32 = 64; // a recognizer batch's width is rounded up to a multiple of this, so shapes recur
}
// An engine for the model named in models_folder/models.config. With prefer_gpu it runs on CUDA when
// every library that needs is present; otherwise, or when CUDA fails to start, on the CPU.
ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Ocr_Engine, ok: bool, error: string {
ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, options := Ocr_Options.{}) -> *Ocr_Engine, ok: bool, error: string {
push_allocator(context.default_allocator);
loaded, load_error := ort_load();
@ -76,6 +89,7 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) ->
engine := New(Ocr_Engine);
engine.model = model.*;
engine.options = options;
dictionary_text, dictionary_ok := read_entire_file(tprint("%/%", models_folder, model.dictionary), log_errors = false);
if !dictionary_ok { ocr_destroy(engine); return null, false, tprint("Cannot read the dictionary %/%.", models_folder, model.dictionary); }
@ -102,10 +116,11 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) ->
if use_gpu {
// A few CPU threads prepare images and decode results while the GPU runs the networks.
engine.workers = 4;
engine.workers = max(options.gpu_workers, 1);
sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, true, 2);
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);
@ -119,6 +134,7 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) ->
sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, false, threads);
if !sessions_ok { ocr_destroy(engine); return null, false, sessions_error; }
engine.provider = .CPU;
engine.rec_batch = max(engine.model.rec_batch_cpu, 1);
}
return engine, true, "";
}
@ -188,7 +204,7 @@ ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.
upside_down := NewArray(first_count, bool,, context.default_allocator);
defer array_free(upside_down,, context.default_allocator);
Orientation_Context :: struct { jobs: [] Line_Job; upside_down: [] bool; size: s64; failed: bool; error: string; }
classify := Orientation_Context.{jobs = jobs, upside_down = upside_down, size = max(engine.model.rec_batch, 1)};
classify := Orientation_Context.{jobs = jobs, upside_down = upside_down, size = max(engine.rec_batch, 1)};
if engine.orientation.session && first_count {
run_on_workers(engine, (first_count + classify.size - 1) / classify.size, (engine: *Ocr_Engine, index: s64, data: *void) {
c := cast(*Orientation_Context) data;
@ -217,7 +233,7 @@ ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.
for * jobs order[it_index] = it;
quick_sort(order, (a: *Line_Job, b: *Line_Job) -> s64 { return a.width - b.width; });
batch_size := max(engine.model.rec_batch, 1);
batch_size := max(engine.rec_batch, 1);
batch_count := (order.count + batch_size - 1) / batch_size;
Batch_Context :: struct { order: [] *Line_Job; size: s64; failed: bool; error: string; }
batches := Batch_Context.{order = order, size = batch_size};

4
modules/OCR/recognize.jai
View File

@ -25,6 +25,10 @@ ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, er
width: s32 = 0;
for jobs width = max(width, it.width);
width = max(width, height); // the network needs a few steps even for one character
if engine.provider == .CUDA { // fewer shapes on the GPU; on the CPU, padding only costs
step := max(engine.options.rec_width_step, 1);
width = (width + step - 1) / step * step;
}
plane := cast(s64) height * width;
input := NewArray(plane * 3 * jobs.count, float, initialized = true,, context.default_allocator); // 0 is the padding

38
src/app/ocr_ui.jai
View File

@ -33,6 +33,7 @@ ocr_ui: struct {
scan: Ocr_Scan;
scan_requested: bool; // -scan has started its scan
scan_logged_step: s64; // -scan's last progress line, in twentieths of the scan
// The Capture window.
overlay_model: string; // the section shown over the image; "" for the chosen model's
@ -271,10 +272,20 @@ ocr_ui_scan_and_exit :: (busy: bool) {
if !ocr_ui.scan_requested {
if busy || !ocr_ui.engine || viewer.catalog_loaded < viewer.captures.count return;
ocr_ui.scan_requested = true;
ocr_ui.window_open = true; // its progress bar
ocr_scan_start();
return;
}
if busy || ocr_ui.scan.running return;
if ocr_ui.scan.running {
done, total, failed, rate, left := ocr_scan_progress();
step := max(total / 20, 1);
if done / step > ocr_ui.scan_logged_step {
ocr_ui.scan_logged_step = done / step;
log("[OCR scan] % read, % failed, % captures/s, % left.", ocr_scan_progress_text(done, total), failed, FF(rate, 1), ocr_duration_text(left));
}
return;
}
if busy return;
finish(ifx ocr_ui.scan.failed then cast(s32) 1 else 0, "[OCR scan] -scan is done: % read, % failed.", ocr_ui.scan.done, ocr_ui.scan.failed);
}
@ -320,12 +331,8 @@ ocr_scan_ui :: () { using ImGui;
return;
}
done, total, failed := scan.done, scan.total, scan.failed;
elapsed := to_float64_seconds(current_time_monotonic() - scan.start) - scan.paused_seconds;
if scan.paused elapsed -= to_float64_seconds(current_time_monotonic() - scan.pause_start);
rate := ifx elapsed > 0 then done / elapsed else 0;
left := ifx rate > 0 then (total - done) / rate else 0;
ProgressBar(cast(float) done / max(total, 1), .{-1, 0}, temp_c_string(tprint("% / %", done, total)));
done, total, failed, rate, left := ocr_scan_progress();
ProgressBar(cast(float) done / max(total, 1), .{-FLOAT32_MIN, 0}, temp_c_string(ocr_scan_progress_text(done, total)));
Text("% captures/s, % left, % failed, on %", FF(rate, 1), ocr_duration_text(left), failed, ifx ocr_ui.engine && ocr_ui.engine.provider == .CUDA then "the GPU" else "the CPU");
if scan.paused {
if Button("Resume") { scan.paused_seconds += to_float64_seconds(current_time_monotonic() - scan.pause_start); atomic_swap(*scan.paused, false); }
@ -336,6 +343,23 @@ ocr_scan_ui :: () { using ImGui;
if Button("Cancel") { atomic_swap(*scan.cancel, true); atomic_swap(*scan.paused, false); }
}
// Where a running scan is: captures read (or failed) of the total, the rate over the time not paused,
// and the time left at that rate.
ocr_scan_progress :: () -> done: s64, total: s64, failed: s64, rate: float64, seconds_left: float64 {
scan := *ocr_ui.scan;
done, total, failed := scan.done, scan.total, scan.failed;
elapsed := to_float64_seconds(current_time_monotonic() - scan.start) - scan.paused_seconds;
if scan.paused elapsed -= to_float64_seconds(current_time_monotonic() - scan.pause_start);
rate := ifx elapsed > 0 then done / elapsed else 0;
left := ifx rate > 0 then (total - done) / rate else 0;
return done, total, failed, rate, left;
}
// "1234 of 8637 (14%)". Temp.
ocr_scan_progress_text :: (done: s64, total: s64) -> string {
return join(tprint("% of % (%", done, total, done * 100 / max(total, 1)), "%)",, temp); // print has no escape for %
}
ocr_duration_text :: (seconds: float64) -> string {
s := cast(s64) seconds;
if s >= 3600 return tprint("% h % min", s / 3600, (s % 3600) / 60);

7
src/app/viewer.jai
View File

@ -545,6 +545,13 @@ viewer_ui :: () { using ImGui;
SetNextItemWidth(160);
SliderFloat("Size", *viewer.thumb_size, 64, MCAP_PREVIEW_MAX_SIDE, "%.0f px");
// A running OCR scan, as a thin bar under the search, so it shows with the OCR window closed.
if ocr_ui.scan.running {
done, total, failed, rate, left := ocr_scan_progress();
ProgressBar(cast(float) done / max(total, 1), .{-FLOAT32_MIN, max(2.0, GetFontSize() * 0.15)}, "");
if IsItemHovered() SetTooltip("OCR scan with %: %, % left, % failed%", ocr_ui.scan.model, ocr_scan_progress_text(done, total), ocr_duration_text(left), failed, ifx ocr_ui.scan.paused then " (paused)" else "");
}
BeginChild("grid");
viewer_grid();
EndChild();

35
tools/ocr_bench.jai
View File

@ -2,17 +2,22 @@
// on the CPU, and reports the time per image and per stage. Each capture's lines are written to
// OUT/<model>-<provider>/<name>.txt, to read against the images.
//
// ocr-bench MCAP_FOLDER [-models a,b] [-cpu] [-batch N] [-limit N] [-out FOLDER]
// ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER]
// [-workers N] [-rec-batch N] [-warmup N]
//
// -models the models to run (default: every model in ocr-models/models.config)
// -cpu the CPU only; by default the GPU too, when it is available
// -models the models to run (default: every model in ocr-models/models.config)
// -cpu the CPU only; by default the GPU too, when it is available
// -gpu-only the GPU only
// -workers on the GPU, threads preparing images and decoding results (default Ocr_Options')
// -rec-batch text lines per recognizer run, in place of models.config's
// -warmup untimed passes over every image before the timed one (default 0)
// -batch images per ocr_process call (default 8)
// -limit at most this many captures (default all)
// -out where the text goes (default ocr-out beside this program)
//
// The models are read from ocr-models beside this program, where the build copies them.
USAGE :: "Usage: ocr-bench MCAP_FOLDER [-models a,b] [-cpu] [-batch N] [-limit N] [-out FOLDER]\n";
USAGE :: "Usage: ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER] [-workers N] [-rec-batch N] [-warmup N]\n";
main :: () {
runtime_support_headless = true; // a command-line tool: a failure is reported on the console and exits with 1, no message box
@ -21,6 +26,10 @@ main :: () {
folder := args[1];
model_names: [..] string;
cpu_only := false;
gpu_only := false;
options: Ocr_Options;
rec_batch: s32 = 0; // 0: as models.config says
warmup := 0;
batch := 8;
limit := S64_MAX_VALUE;
exe_folder := path_strip_filename(get_path_of_running_executable());
@ -33,6 +42,10 @@ main :: () {
if arg == {
case "-models"; for split(value, ",") array_add(*model_names, trim(it)); i += 2;
case "-cpu"; cpu_only = true; i += 1;
case "-gpu-only"; gpu_only = true; i += 1;
case "-workers"; options.gpu_workers = cast(s32) max(1, string_to_int(value)); i += 2;
case "-rec-batch"; rec_batch = cast(s32) max(1, string_to_int(value)); i += 2;
case "-warmup"; warmup = max(0, string_to_int(value)); i += 2;
case "-batch"; batch = max(1, string_to_int(value)); i += 2;
case "-limit"; limit = max(1, string_to_int(value)); i += 2;
case "-out"; out_folder = value; i += 2;
@ -75,9 +88,11 @@ main :: () {
for model_name: model_names {
for prefer_gpu: bool.[true, false] {
if prefer_gpu && cpu_only continue;
if !prefer_gpu && gpu_only continue;
create_start := current_time_monotonic();
engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu);
engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu, options);
if !ok { print("% : %\n\n", model_name, error); break; }
if rec_batch engine.rec_batch = rec_batch;
create_seconds := seconds_since(create_start);
in_use, explanation := ocr_gpu_status(engine);
provider := ocr_provider(engine);
@ -87,7 +102,7 @@ main :: () {
ocr_destroy(engine);
continue;
}
print("== % on % (sessions created in % s, % worker threads)\n", model_name, provider, FF(create_seconds, 2), engine.workers);
print("== % on % (sessions created in % s, % worker threads, % lines per recognizer run)\n", model_name, provider, FF(create_seconds, 2), engine.workers, engine.rec_batch);
text_folder := sprint("%/%-%", out_folder, model_name, to_lower_copy(tprint("%", provider),, temp));
make_directory_if_it_does_not_exist(text_folder, recursive = true);
@ -96,6 +111,14 @@ main :: () {
first_start := current_time_monotonic();
ocr_process(engine, array_view(images, 0, 1),, temp);
first_seconds := seconds_since(first_start);
// Untimed passes over everything, so the timed one does not pay for setting up each new shape.
for 1..warmup {
for start: 0..images.count - 1 {
if start % batch continue;
ocr_process(engine, array_view(images, start, min(batch, images.count - start)),, temp);
reset_temporary_storage();
}
}
ocr_reset_stats(engine);
start := current_time_monotonic();