219 lines
9.8 KiB
Plaintext
219 lines
9.8 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;
|
|
}
|
|
|
|
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; 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;
|
|
}
|
|
|
|
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 {
|
|
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.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.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.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;
|
|
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;
|
|
}
|
|
return true, "";
|
|
}
|
|
|
|
ocr_release_sessions :: (engine: *Ocr_Engine) {
|
|
ocr_release_session(*engine.det);
|
|
ocr_release_session(*engine.rec);
|
|
ocr_release_session(*engine.orientation);
|
|
}
|
|
|
|
// 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_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, "";
|
|
}
|
|
|
|
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.* = .{};
|
|
}
|