233 lines
12 KiB
Plaintext
233 lines
12 KiB
Plaintext
// models.config, and running a model: creating its onnxruntime session on the GPU or the CPU, and
|
|
// running it on a float tensor.
|
|
|
|
Ocr_Model_Info :: struct {
|
|
name: string;
|
|
title: string;
|
|
detector, recognizer, dictionary: string; // file names in the models folder
|
|
textline_ori: string; // the orientation classifier's file name; "" for none
|
|
det_max_side: s32 = 2560;
|
|
det_tile_height: s32 = 1536;
|
|
det_threshold: float = 0.3;
|
|
box_threshold: float = 0.6;
|
|
unclip_ratio: float = 1.5;
|
|
rec_height: s32 = 48;
|
|
rec_batch_gpu: s32 = 16;
|
|
rec_batch_cpu: s32 = 1;
|
|
|
|
// macOS: the detector and recognizer as Core ML packages (tools/export_ppocr_coreml.py), "" for none,
|
|
// and the sizes they take, which inputs are padded up to: the detector's height and width multiples
|
|
// of these steps, and the recognizer coreml_rec_batch lines at a multiple of coreml_rec_width_step.
|
|
detector_coreml, recognizer_coreml: string;
|
|
coreml_det_height_step: s32 = 128;
|
|
coreml_det_width_step: s32 = 256;
|
|
coreml_rec_batch: s32 = 16;
|
|
coreml_rec_width_step: s32 = 64;
|
|
}
|
|
|
|
OCR_MODELS_CONFIG :: "models.config";
|
|
|
|
// The models models.config lists, in its order; allocated with the context allocator. A line that
|
|
// cannot be read is logged and skipped, as in the Ex4 formats.
|
|
ocr_list_models :: (models_folder: string) -> [] Ocr_Model_Info, error: string {
|
|
models: [..] Ocr_Model_Info;
|
|
path := tprint("%/%", models_folder, OCR_MODELS_CONFIG);
|
|
text, ok := read_entire_file(path, log_errors = false,, temp);
|
|
if !ok return models, tprint("Cannot read %.", path);
|
|
|
|
model: *Ocr_Model_Info;
|
|
version_seen := false;
|
|
line_number := 0;
|
|
for raw_line: split(text, "\n",, temp) {
|
|
line_number += 1;
|
|
line := trim(ocr_strip_comment(raw_line));
|
|
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("[OCR] %:%: 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_text := value;
|
|
number, number_ok := string_to_float(number_text);
|
|
if key == {
|
|
case "title"; model.title = copy_string(value);
|
|
case "detector"; model.detector = copy_string(value);
|
|
case "recognizer"; model.recognizer = copy_string(value);
|
|
case "dictionary"; model.dictionary = copy_string(value);
|
|
case "textline_ori"; model.textline_ori = copy_string(value);
|
|
case "det_max_side"; if number_ok model.det_max_side = cast(s32) number;
|
|
case "det_tile_height"; if number_ok model.det_tile_height = cast(s32) number;
|
|
case "det_threshold"; if number_ok model.det_threshold = number;
|
|
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_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 "detector_coreml"; model.detector_coreml = copy_string(value);
|
|
case "recognizer_coreml"; model.recognizer_coreml = copy_string(value);
|
|
case "coreml_det_height_step"; if number_ok model.coreml_det_height_step = cast(s32) number;
|
|
case "coreml_det_width_step"; if number_ok model.coreml_det_width_step = cast(s32) number;
|
|
case "coreml_rec_batch"; if number_ok model.coreml_rec_batch = cast(s32) number;
|
|
case "coreml_rec_width_step"; if number_ok model.coreml_rec_width_step = cast(s32) number;
|
|
case; log_error("[OCR] %:%: unknown key %. Ignored.", path, line_number, key);
|
|
}
|
|
}
|
|
return models, "";
|
|
}
|
|
|
|
#scope_module
|
|
|
|
// "#" opens a comment at the start of a line or after whitespace, outside quotes.
|
|
ocr_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;
|
|
}
|
|
|
|
// 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 {
|
|
model := *engine.model;
|
|
ok: bool;
|
|
error: string;
|
|
// The detector's and recognizer's inputs change size with nearly every run: their arenas give back
|
|
// what no run is using once per call (ocr_shrink_arena), or they keep a block of every size.
|
|
engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads, engine.options, shrink_arena = true);
|
|
if !ok return false, error;
|
|
engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads, engine.options, shrink_arena = true);
|
|
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.options);
|
|
if !ok return false, error;
|
|
ocr_orientation_size(engine);
|
|
}
|
|
return true, "";
|
|
}
|
|
|
|
// PP-LCNet text line orientation takes 160 by 80; a model that says otherwise is taken at its word.
|
|
ocr_orientation_size :: (engine: *Ocr_Engine) {
|
|
shape := engine.orientation.input_shape;
|
|
engine.orientation_height = ifx engine.orientation.input_rank == 4 && shape[2] > 0 then cast(s32) shape[2] else 80;
|
|
engine.orientation_width = ifx engine.orientation.input_rank == 4 && shape[3] > 0 then cast(s32) shape[3] else 160;
|
|
}
|
|
|
|
// macOS: the detector and recognizer through Core ML (modules/Core_ML) on the GPU, each run once on
|
|
// its smallest size, and the orientation classifier, if any, through onnxruntime on the CPU. On
|
|
// failure, why; what was made is left for ocr_release_coreml and ocr_release_sessions.
|
|
ocr_create_coreml :: (engine: *Ocr_Engine, models_folder: string, cores: s32) -> ok: bool, error: string {
|
|
#if OS != .MACOS {
|
|
return false, "Core ML is macOS's.";
|
|
} else {
|
|
model := *engine.model;
|
|
if !model.detector_coreml || !model.recognizer_coreml return false, tprint("there are no Core ML models for %; tools/export_ppocr_coreml.py makes them.", model.name);
|
|
for string.[model.detector_coreml, model.recognizer_coreml] {
|
|
if !file_exists(tprint("%/%", models_folder, it)) return false, tprint("% is not there; tools/export_ppocr_coreml.py makes it.", it);
|
|
}
|
|
ok: bool;
|
|
error: string;
|
|
engine.det_coreml, ok, error = coreml_load(tprint("%/%", models_folder, model.detector_coreml), "x", "y");
|
|
if !ok return false, error;
|
|
engine.rec_coreml, ok, error = coreml_load(tprint("%/%", models_folder, model.recognizer_coreml), "x", "y");
|
|
if !ok return false, error;
|
|
if model.textline_ori {
|
|
engine.orientation, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.textline_ori), false, max(cores / 4, 1), engine.options);
|
|
if !ok return false, error;
|
|
ocr_orientation_size(engine);
|
|
}
|
|
|
|
// The first runs, at the smallest sizes the packages take.
|
|
Warm_Up :: struct { model: *Coreml_Model; shape: [4] s64; }
|
|
runs := Warm_Up.[
|
|
.{engine.det_coreml, .[1, 3, model.coreml_det_height_step, model.coreml_det_width_step]},
|
|
.{engine.rec_coreml, .[model.coreml_rec_batch, 3, model.rec_height, model.coreml_rec_width_step]},
|
|
];
|
|
for runs {
|
|
input := NewArray(it.shape[0] * it.shape[1] * it.shape[2] * it.shape[3], float, initialized = true,, context.default_allocator);
|
|
defer array_free(input,, context.default_allocator);
|
|
out, run_ok, run_error := coreml_predict(it.model, input, it.shape);
|
|
if !run_ok return false, tprint("Core ML could not run the % on the GPU: %", ifx it_index == 0 then "detector" else "recognizer", run_error);
|
|
coreml_release_output(*out);
|
|
}
|
|
return true, "";
|
|
}
|
|
}
|
|
|
|
ocr_release_coreml :: (engine: *Ocr_Engine) {
|
|
#if OS == .MACOS {
|
|
coreml_free(engine.det_coreml);
|
|
coreml_free(engine.rec_coreml);
|
|
engine.det_coreml, engine.rec_coreml = null;
|
|
}
|
|
}
|
|
|
|
ocr_release_sessions :: (engine: *Ocr_Engine) {
|
|
ocr_release_session(*engine.det);
|
|
ocr_release_session(*engine.rec);
|
|
ocr_release_session(*engine.orientation);
|
|
}
|
|
|
|
// Runs each of a CUDA engine's sessions once, on a small input, before it is used or destroyed.
|
|
// With cuBLASLt 13.8 (CUDA 13, onnxruntime 1.30.0, RTX 5080 laptop), the process's first CUDA
|
|
// engine destroyed without ever having run left the next engine to crash in cublasLtCtxInit, with a
|
|
// DEP access violation, on its first run; an engine that has run once leaves nothing broken. It also
|
|
// moves the first run's setup into creation. On failure, the reason.
|
|
// On CUDA, the session's arena (the detector's or the recognizer's) gives back every block no run is
|
|
// using: called by ocr_process once a stage's runs are all done, as a run on a small input, since
|
|
// onnxruntime shrinks an arena only at the end of a run. Once per run instead cost a third to a half
|
|
// of OCR's speed.
|
|
ocr_shrink_arena :: (engine: *Ocr_Engine, session: *Ocr_Session) {
|
|
if engine.provider != .CUDA || !session.session return;
|
|
side := ifx session == *engine.rec then cast(s64) engine.model.rec_height else 32;
|
|
shape := s64.[1, 3, side, side];
|
|
input := NewArray(3 * side * side, float, initialized = true,, temp);
|
|
out, ok := ocr_run(session, engine.memory_info, input, shape, shrink = true);
|
|
if ok ocr_release_output(*out);
|
|
}
|
|
|
|
ocr_warm_up :: (engine: *Ocr_Engine) -> ok: bool, error: string {
|
|
Warm_Up :: struct { session: *Ocr_Session; height, width: s64; }
|
|
runs := Warm_Up.[
|
|
.{*engine.det, 64, 64},
|
|
.{*engine.rec, engine.model.rec_height, engine.model.rec_height * 2},
|
|
.{*engine.orientation, engine.orientation_height, engine.orientation_width},
|
|
];
|
|
for runs {
|
|
if !it.session.session continue; // no orientation classifier
|
|
shape := s64.[1, 3, it.height, it.width];
|
|
input := NewArray(3 * it.height * it.width, float, initialized = true,, context.default_allocator);
|
|
defer array_free(input,, context.default_allocator);
|
|
out, ok, error := ocr_run(it.session, engine.memory_info, input, shape);
|
|
if !ok return false, tprint("The first run on the GPU failed: %", error);
|
|
ocr_release_output(*out);
|
|
}
|
|
return true, "";
|
|
}
|
|
|
|
// 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, shrink_arena := false) -> Ocr_Session, ok: bool, error: string {
|
|
settings := Ort_Gpu_Settings.{arena_extend_strategy = ocr_options.cuda_arena_extend_strategy, cudnn_max_workspace = ocr_options.cudnn_max_workspace, shrink_arena = shrink_arena};
|
|
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;
|
|
}
|
|
|
|
// 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;
|