[0.1.30] macOS phase 5: PP-OCRv5 mobile OCR on the GPU through native Core ML
This commit is contained in:
parent
2d0fb855b1
commit
30cd00b6f2
10
README.md
10
README.md
@ -28,10 +28,12 @@ see [Screen capture](#screen-capture).
|
||||
place (`C:\Program Files\NVIDIA\CUDNN`) is found without changing PATH. The OCR window says
|
||||
what, if anything, is missing.
|
||||
|
||||
**On macOS**, labelling runs on the GPU through Core ML, about ten times as fast as the CPU, from
|
||||
a Core ML copy of SigLIP 2's image model made once with `tools/export_siglip_coreml.py` (its first
|
||||
lines say how to set up the Python it needs). Without it, labels are made on the CPU. OCR runs on
|
||||
the CPU: Core ML through onnxruntime was measured slower there.
|
||||
**On macOS**, labelling and OCR run on the GPU through Core ML, from Core ML copies of the models
|
||||
made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten times as fast as
|
||||
the CPU) and `tools/export_ppocr_coreml.py` (PP-OCRv5 mobile's detector and recognizer, about twice
|
||||
as fast; their first lines say how to set up the Python they need). Without them, each runs on the
|
||||
CPU. The first scan after making them is slower: Core ML prepares each image size the first time it
|
||||
meets it, and macOS keeps that for later runs.
|
||||
|
||||
3. **Build and run:** `jai first.jai - release`, then `bin\mcap.exe`. The build copies onnxruntime
|
||||
and the models beside the program.
|
||||
|
||||
@ -49,7 +49,7 @@ APP_NAME :: "mcap";
|
||||
APP_FOLDER_NAME :: "mcap";
|
||||
GIT_BRANCH :: "check";
|
||||
GIT_HASH :: "-";
|
||||
VERSION :: "0.1.29";
|
||||
VERSION :: "0.1.30";
|
||||
BUILD_DATE :: "-";
|
||||
DEBUG :: true;
|
||||
COMPILE_TESTS :: true;
|
||||
|
||||
@ -8,7 +8,7 @@ APP_NAME :: "mcap";
|
||||
APP_FOLDER_NAME :: "mcap";
|
||||
APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
|
||||
|
||||
VERSION :: "0.1.29";
|
||||
VERSION :: "0.1.30";
|
||||
// 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";
|
||||
@ -379,6 +379,7 @@ copy_assets :: () {
|
||||
copy_if_changed("modules/OCR/models.config", tprint("%/ocr-models/models.config", build_directory));
|
||||
if file_exists("modules/OCR/models") {
|
||||
for file_list("modules/OCR/models") copy_if_changed(it, tprint("%/ocr-models/%", build_directory, path_filename(it)));
|
||||
#if OS == .MACOS clone_packages("modules/OCR/models", tprint("%/ocr-models", build_directory));
|
||||
} else {
|
||||
log("OCR: the models have not been fetched; run jai fetch_models.jai in modules/OCR for OCR.");
|
||||
}
|
||||
|
||||
@ -7,9 +7,11 @@
|
||||
// coreml_free(model);
|
||||
//
|
||||
// Loading compiles the package (Core ML's compileModelAtURL, a fraction of a second) into a temporary
|
||||
// .mlmodelc, removed by coreml_free. The model runs on the CPU and GPU, never the Neural Engine
|
||||
// .mlmodelc, removed by coreml_free. What Core ML makes of a model for the GPU, the first time each
|
||||
// size runs, is kept by macOS across runs (the very first run of the OCR models took 5 s more, later
|
||||
// ones none). The model runs on the CPU and GPU, never the Neural Engine
|
||||
// unless asked: with float16 weights it can be badly wrong there (SigLIP 2: cosine 0.72 to the CPU).
|
||||
// A model runs one call at a time (a mutex), from any thread.
|
||||
// A model runs one prediction at a time (a mutex), from any thread; outputs are read outside it.
|
||||
|
||||
Coreml_Units :: enum s64 { // MLComputeUnits
|
||||
CPU_ONLY :: 0;
|
||||
@ -72,32 +74,35 @@ coreml_free :: (model: *Coreml_Model) {
|
||||
free(model,, context.default_allocator);
|
||||
}
|
||||
|
||||
// Runs the model on input, a float32 tensor of shape, and copies its output (float32 or float16) into
|
||||
// output as float32. ok is false, with why, when it fails or gives a different number of values.
|
||||
coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
|
||||
lock(*model.mutex);
|
||||
defer unlock(*model.mutex);
|
||||
// A run's output, as float32 values one after another (C order): a view of Core ML's own array when it
|
||||
// is float32 and packed, else a packed copy. Valid until coreml_release_output.
|
||||
Coreml_Output :: struct {
|
||||
data: *float;
|
||||
dims: [8] s64;
|
||||
dim_count: s64;
|
||||
count: s64;
|
||||
array: *NSObject; // the MLMultiArray, retained
|
||||
copy: *float; // when data is a copy: the default allocator's
|
||||
}
|
||||
|
||||
// Runs the model on input, a float32 tensor of shape. ok is false, with why, when it fails.
|
||||
coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Coreml_Output, ok: bool, error: string {
|
||||
result: Coreml_Output;
|
||||
init_objective_c();
|
||||
pool := objc_new(NSAutoreleasePool);
|
||||
defer NSAutoreleasePool.drain(pool);
|
||||
|
||||
// The input, over the caller's floats: packed, so each stride is the product of the sizes after it.
|
||||
numbers :: (values: [] s64) -> *NSObject {
|
||||
objects := NewArray(values.count, *NSObject,, temp);
|
||||
func_number :: #type (*void, Selector, s64) -> *NSObject #c_call;
|
||||
for values objects[it_index] = (cast(func_number) objc_msgSend)(objc_getClass("NSNumber"), selector("numberWithLongLong:"), it);
|
||||
func_array :: #type (*void, Selector, **NSObject, u64) -> *NSObject #c_call;
|
||||
return (cast(func_array) objc_msgSend)(objc_getClass("NSArray"), selector("arrayWithObjects:count:"), objects.data, xx objects.count);
|
||||
}
|
||||
strides := NewArray(shape.count, s64,, temp);
|
||||
step: s64 = 1;
|
||||
for < shape { strides[it_index] = step; step *= it; }
|
||||
if step != input.count return false, tprint("the input has % values, not % for its shape %", input.count, step, shape);
|
||||
if step != input.count return result, false, tprint("the input has % values, not % for its shape %", input.count, step, shape);
|
||||
|
||||
error: *NSError;
|
||||
func_array_init :: #type (*void, Selector, *void, *NSObject, s64, *NSObject, *void, **NSError) -> *NSObject #c_call;
|
||||
array := (cast(func_array_init) objc_msgSend)(objc_alloc(objc_getClass("MLMultiArray")), selector("initWithDataPointer:shape:dataType:strides:deallocator:error:"),
|
||||
input.data, numbers(shape), MULTIARRAY_FLOAT32, numbers(strides), null, *error);
|
||||
if !array return false, tprint("Core ML could not take the input: %", error_text(error));
|
||||
if !array return result, false, tprint("Core ML could not take the input: %", error_text(error));
|
||||
autorelease(array);
|
||||
|
||||
func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call;
|
||||
@ -107,30 +112,66 @@ coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: []
|
||||
dictionary := (cast(func_dictionary) objc_msgSend)(objc_getClass("NSDictionary"), selector("dictionaryWithObject:forKey:"), value, key);
|
||||
func_provider :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
|
||||
features := (cast(func_provider) objc_msgSend)(objc_alloc(objc_getClass("MLDictionaryFeatureProvider")), selector("initWithDictionary:error:"), dictionary, *error);
|
||||
if !features return false, tprint("Core ML could not take the input: %", error_text(error));
|
||||
if !features return result, false, tprint("Core ML could not take the input: %", error_text(error));
|
||||
autorelease(features);
|
||||
|
||||
func_predict :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
|
||||
result := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error);
|
||||
if !result return false, tprint("Core ML could not run the model: %", error_text(error));
|
||||
lock(*model.mutex);
|
||||
prediction := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error);
|
||||
unlock(*model.mutex);
|
||||
if !prediction return result, false, tprint("Core ML could not run the model: %", error_text(error));
|
||||
|
||||
out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(model.output_name));
|
||||
out_array := ifx out_value then cast(*NSObject) objc_msgSend_typed(out_value, selector("multiArrayValue")) else null;
|
||||
if !out_array return result, false, tprint("the model gives no %", model.output_name);
|
||||
|
||||
out_shape := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("shape")));
|
||||
out_strides := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("strides")));
|
||||
if out_shape.count > result.dims.count || out_shape.count != out_strides.count return result, false, tprint("the model's output has an unexpected shape %", out_shape);
|
||||
result.dim_count = out_shape.count;
|
||||
result.count = 1;
|
||||
for out_shape { result.dims[it_index] = it; result.count *= it; }
|
||||
|
||||
out_value := (cast(func_value) objc_msgSend)(result, selector("featureValueForName:"), NSString.getTempString(model.output_name));
|
||||
out_array := ifx out_value then objc_msgSend_typed(out_value, selector("multiArrayValue")) else null;
|
||||
if !out_array return false, tprint("the model gives no %", model.output_name);
|
||||
count := cast(s64) objc_msgSend_typed(out_array, selector("count"));
|
||||
if count != output.count return false, tprint("the model gives % values, not %", count, output.count);
|
||||
func_type :: #type (*void, Selector) -> s64 #c_call;
|
||||
type := (cast(func_type) objc_msgSend)(out_array, selector("dataType"));
|
||||
// The output is packed (one row of the pooled vector here); dataPointer is read before the pool drains.
|
||||
if type != MULTIARRAY_FLOAT32 && type != MULTIARRAY_FLOAT16 return result, false, tprint("the model's output is of type 0x%, not float32 or float16", formatInt(type, base = 16));
|
||||
packed := true;
|
||||
expected: s64 = 1;
|
||||
for < out_shape { if out_strides[it_index] != expected packed = false; expected *= it; }
|
||||
|
||||
data := objc_msgSend_typed(out_array, selector("dataPointer"));
|
||||
if type == MULTIARRAY_FLOAT32 {
|
||||
memcpy(output.data, data, count * size_of(float));
|
||||
} else if type == MULTIARRAY_FLOAT16 {
|
||||
half := cast(*u16) data;
|
||||
for 0..count - 1 output[it] = float16_to_float32(half[it]);
|
||||
} else {
|
||||
return false, tprint("the model's output is of type 0x%, not float32 or float16", formatInt(type, base = 16));
|
||||
result.array = xx retain(out_array); // outlives the pool, and with it the data
|
||||
if type == MULTIARRAY_FLOAT32 && packed {
|
||||
result.data = cast(*float) data;
|
||||
return result, true, "";
|
||||
}
|
||||
// A copy, packed: float16 widened, and rows Core ML padded (GPU outputs may be) closed up.
|
||||
result.copy = cast(*float) alloc(result.count * size_of(float),, context.default_allocator);
|
||||
result.data = result.copy;
|
||||
index: [8] s64;
|
||||
for i: 0..result.count - 1 {
|
||||
offset: s64 = 0;
|
||||
for d: 0..result.dim_count - 1 offset += index[d] * out_strides[d];
|
||||
result.copy[i] = ifx type == MULTIARRAY_FLOAT32 then (cast(*float) data)[offset] else float16_to_float32((cast(*u16) data)[offset]);
|
||||
d := result.dim_count - 1;
|
||||
while d >= 0 { index[d] += 1; if index[d] < out_shape[d] break; index[d] = 0; d -= 1; }
|
||||
}
|
||||
return result, true, "";
|
||||
}
|
||||
|
||||
coreml_release_output :: (out: *Coreml_Output) {
|
||||
if out.array release(out.array);
|
||||
if out.copy free(out.copy,, context.default_allocator);
|
||||
out.* = .{};
|
||||
}
|
||||
|
||||
// coreml_predict, its output copied into output (which must hold exactly as many values).
|
||||
coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
|
||||
out, ok, error := coreml_predict(model, input, shape);
|
||||
if !ok return false, error;
|
||||
defer coreml_release_output(*out);
|
||||
if out.count != output.count return false, tprint("the model gives % values, not %", out.count, output.count);
|
||||
memcpy(output.data, out.data, out.count * size_of(float));
|
||||
return true, "";
|
||||
}
|
||||
|
||||
@ -139,6 +180,28 @@ coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: []
|
||||
MULTIARRAY_FLOAT32 :: 0x10000 | 32; // MLMultiArrayDataTypeFloat32
|
||||
MULTIARRAY_FLOAT16 :: 0x10000 | 16; // MLMultiArrayDataTypeFloat16
|
||||
|
||||
numbers :: (values: [] s64) -> *NSObject {
|
||||
objects := NewArray(values.count, *NSObject,, temp);
|
||||
func_number :: #type (*void, Selector, s64) -> *NSObject #c_call;
|
||||
for values objects[it_index] = (cast(func_number) objc_msgSend)(objc_getClass("NSNumber"), selector("numberWithLongLong:"), it);
|
||||
func_array :: #type (*void, Selector, **NSObject, u64) -> *NSObject #c_call;
|
||||
return (cast(func_array) objc_msgSend)(objc_getClass("NSArray"), selector("arrayWithObjects:count:"), objects.data, xx objects.count);
|
||||
}
|
||||
|
||||
// An NSArray of NSNumbers, as integers. Temp.
|
||||
integers :: (array: *NSObject) -> [] s64 {
|
||||
if !array return .[];
|
||||
count := cast(s64) objc_msgSend_typed(array, selector("count"));
|
||||
result := NewArray(count, s64,, temp);
|
||||
func_at :: #type (*void, Selector, u64) -> *NSObject #c_call;
|
||||
func_value :: #type (*void, Selector) -> s64 #c_call;
|
||||
for 0..count - 1 {
|
||||
number := (cast(func_at) objc_msgSend)(array, selector("objectAtIndex:"), xx it);
|
||||
result[it] = (cast(func_value) objc_msgSend)(number, selector("longLongValue"));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
error_text :: (error: *NSError) -> string {
|
||||
if !error return "no reason given";
|
||||
return copy_string(to_string(NSError.localizedDescription(error)),, temp);
|
||||
|
||||
@ -63,8 +63,14 @@ ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_
|
||||
valid_h := tile_bottom - tile_top;
|
||||
w := align_up_32(valid_w);
|
||||
h := align_up_32(valid_h);
|
||||
#if OS == .MACOS if engine.det_coreml {
|
||||
// Up to a size the Core ML package takes: its steps, which det_max_side and det_tile_height stay within.
|
||||
w = align_up(valid_w, model.coreml_det_width_step);
|
||||
h = align_up(valid_h, model.coreml_det_height_step);
|
||||
}
|
||||
|
||||
// Planar BGR, normalized; the padding to a multiple of 32 is the normalized zero (the mean colour).
|
||||
// Planar BGR, normalized; the padding to a multiple of 32 (or the Core ML step) is the normalized
|
||||
// zero (the mean colour).
|
||||
plane := cast(s64) w * h;
|
||||
input := NewArray(plane * 3, float, initialized = true,, context.default_allocator);
|
||||
defer array_free(input,, context.default_allocator);
|
||||
@ -83,14 +89,31 @@ ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_
|
||||
}
|
||||
|
||||
shape := s64.[1, 3, h, w];
|
||||
out, ok, error := ocr_run(*engine.det, engine.memory_info, input, shape);
|
||||
if !ok return false, error;
|
||||
probabilities: *float;
|
||||
dims: [] s64;
|
||||
out: Ocr_Output;
|
||||
defer ocr_release_output(*out);
|
||||
if out.dim_count != 4 || out.dims[2] != h || out.dims[3] != w {
|
||||
return false, tprint("The detector gave an output of shape %, not [1, 1, %, %].", array_view(out.dims, 0, out.dim_count), h, w);
|
||||
#if OS == .MACOS {
|
||||
coreml_out: Coreml_Output;
|
||||
defer coreml_release_output(*coreml_out);
|
||||
}
|
||||
#if OS == .MACOS if engine.det_coreml {
|
||||
ok: bool; error: string;
|
||||
coreml_out, ok, error = coreml_predict(engine.det_coreml, input, shape);
|
||||
if !ok return false, error;
|
||||
probabilities, dims = coreml_out.data, array_view(coreml_out.dims, 0, coreml_out.dim_count);
|
||||
}
|
||||
if !probabilities {
|
||||
ok: bool; error: string;
|
||||
out, ok, error = ocr_run(*engine.det, engine.memory_info, input, shape);
|
||||
if !ok return false, error;
|
||||
probabilities, dims = out.data, array_view(out.dims, 0, out.dim_count);
|
||||
}
|
||||
if dims.count != 4 || dims[2] != h || dims[3] != w {
|
||||
return false, tprint("The detector gave an output of shape %, not [1, 1, %, %].", dims, h, w);
|
||||
}
|
||||
|
||||
regions := ocr_find_regions(out.data, w, valid_w, valid_h, model.det_threshold);
|
||||
regions := ocr_find_regions(probabilities, w, valid_w, valid_h, model.det_threshold);
|
||||
defer {
|
||||
for regions array_free(it.outline);
|
||||
array_free(regions,, context.default_allocator);
|
||||
@ -254,6 +277,12 @@ align_up_32 :: (n: s32) -> s32 {
|
||||
return (n + 31) & cast(s32) ~31;
|
||||
}
|
||||
|
||||
// n rounded up to a multiple of step (at least one step).
|
||||
align_up :: (n: s32, step: s32) -> s32 {
|
||||
step = max(step, 1);
|
||||
return max((n + step - 1) / step, 1) * step;
|
||||
}
|
||||
|
||||
Region :: struct {
|
||||
pixels: s64;
|
||||
probability_sum: float64;
|
||||
|
||||
@ -16,6 +16,12 @@
|
||||
# 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
|
||||
# detector_coreml, recognizer_coreml
|
||||
# macOS: the detector and recognizer as Core ML packages, quoted, run on the GPU
|
||||
# (tools/export_ppocr_coreml.py makes them); without both, onnxruntime runs OCR
|
||||
# coreml_det_height_step, coreml_det_width_step, coreml_rec_batch, coreml_rec_width_step
|
||||
# the sizes those packages take, as the script made them (the defaults: 128, 256,
|
||||
# 16, 64): detection strips are padded up to the steps, lines to 16 at the step
|
||||
|
||||
[model ppocr-v5-mobile]
|
||||
title: "PP-OCRv5 mobile (light)"
|
||||
@ -31,6 +37,8 @@ unclip_ratio: 1.5
|
||||
rec_height: 48
|
||||
rec_batch_gpu: 16
|
||||
rec_batch_cpu: 1
|
||||
detector_coreml: "ch_PP-OCRv5_det_mobile.mlpackage"
|
||||
recognizer_coreml: "ch_PP-OCRv5_rec_mobile.mlpackage"
|
||||
|
||||
[model ppocr-v5-server]
|
||||
title: "PP-OCRv5 server (heavy)"
|
||||
|
||||
@ -14,6 +14,15 @@ Ocr_Model_Info :: struct {
|
||||
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";
|
||||
@ -64,6 +73,12 @@ ocr_list_models :: (models_folder: string) -> [] Ocr_Model_Info, error: string {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@ -98,13 +113,66 @@ ocr_create_sessions :: (engine: *Ocr_Engine, models_folder: string, gpu: bool, i
|
||||
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);
|
||||
|
||||
@ -12,7 +12,7 @@
|
||||
// then recognition (recognize.jai) reads it level. When the model lists an orientation classifier,
|
||||
// lines that look upside down are read both ways; lines read downwards or upwards are read too.
|
||||
|
||||
Ocr_Provider :: enum { CPU; CUDA; }
|
||||
Ocr_Provider :: enum { CPU; CUDA; COREML; }
|
||||
|
||||
// Tightly packed or not: stride is the bytes from one row to the next.
|
||||
Ocr_Image :: struct {
|
||||
@ -40,6 +40,7 @@ Ocr_Engine :: struct {
|
||||
env: *OrtEnv;
|
||||
memory_info: *OrtMemoryInfo;
|
||||
det, rec: Ocr_Session;
|
||||
#if OS == .MACOS det_coreml, rec_coreml: *Coreml_Model; // the detector and recognizer, when Core ML runs them
|
||||
orientation: Ocr_Session; // the text line orientation classifier; none when its session is null
|
||||
orientation_width, orientation_height: s32; // its input size
|
||||
dictionary: [] string; // class n is dictionary[n - 1]
|
||||
@ -117,11 +118,24 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op
|
||||
cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1); // efficiency cores included
|
||||
|
||||
use_gpu := false;
|
||||
if prefer_gpu {
|
||||
#if OS == .MACOS if prefer_gpu {
|
||||
// Core ML on the GPU: the detector and recognizer, the orientation classifier on the CPU. On any
|
||||
// failure, onnxruntime on the CPU below, with the reason.
|
||||
coreml_ok, coreml_error := ocr_create_coreml(engine, models_folder, cores);
|
||||
if coreml_ok {
|
||||
engine.provider = .COREML;
|
||||
engine.workers = max(options.gpu_workers, 1);
|
||||
engine.rec_batch = max(engine.model.coreml_rec_batch, 1);
|
||||
return engine, true, "";
|
||||
}
|
||||
ocr_release_coreml(engine);
|
||||
engine.gpu_explanation = sprint("OCR runs on the CPU: %", coreml_error);
|
||||
}
|
||||
if prefer_gpu && OS != .MACOS {
|
||||
available, explanation := ort_cuda_check("OCR"); // modules/onnxruntime/cuda.jai
|
||||
use_gpu = available;
|
||||
if !available engine.gpu_explanation = copy_string(explanation);
|
||||
} else {
|
||||
} else if !prefer_gpu {
|
||||
engine.gpu_explanation = copy_string("OCR runs on the CPU, as it was asked to.");
|
||||
}
|
||||
|
||||
@ -160,7 +174,16 @@ ocr_provider :: (engine: *Ocr_Engine) -> Ocr_Provider {
|
||||
|
||||
// Whether the GPU is in use, and when not, why, and what to install so it can be.
|
||||
ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string {
|
||||
return engine.provider == .CUDA, engine.gpu_explanation;
|
||||
return engine.provider == .CUDA || engine.provider == .COREML, engine.gpu_explanation;
|
||||
}
|
||||
|
||||
// What runs the detector and recognizer, for people and the files it reads: "CUDA", "Core ML" or "CPU".
|
||||
ocr_provider_name :: (engine: *Ocr_Engine) -> string {
|
||||
if engine.provider == {
|
||||
case .CUDA; return "CUDA";
|
||||
case .COREML; return "Core ML";
|
||||
}
|
||||
return "CPU";
|
||||
}
|
||||
|
||||
ocr_destroy :: (engine: *Ocr_Engine) {
|
||||
@ -169,6 +192,7 @@ ocr_destroy :: (engine: *Ocr_Engine) {
|
||||
if engine.provider == .CUDA ort_gpu_lock();
|
||||
ocr_release_sessions(engine);
|
||||
if engine.provider == .CUDA ort_gpu_unlock();
|
||||
ocr_release_coreml(engine);
|
||||
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);
|
||||
@ -458,6 +482,7 @@ run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void
|
||||
#import "Atomics";
|
||||
#import "onnxruntime";
|
||||
Default_Allocator :: #import "Default_Allocator"; // rpmalloc (modules/Default_Allocator), for thread_idle where the workers wait
|
||||
#if OS == .MACOS #import "Core_ML";
|
||||
#if OS == .WINDOWS {
|
||||
#import "Windows";
|
||||
#import "Windows_Utf8";
|
||||
|
||||
@ -29,21 +29,44 @@ ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, er
|
||||
step := max(engine.options.rec_width_step, 1);
|
||||
width = (width + step - 1) / step * step;
|
||||
}
|
||||
rows := cast(s64) jobs.count;
|
||||
#if OS == .MACOS if engine.rec_coreml {
|
||||
// The sizes the Core ML package takes: coreml_rec_batch lines (the rest padding), at its width step.
|
||||
width = align_up(width, engine.model.coreml_rec_width_step);
|
||||
rows = max(rows, cast(s64) engine.model.coreml_rec_batch);
|
||||
}
|
||||
|
||||
plane := cast(s64) height * width;
|
||||
input := NewArray(plane * 3 * jobs.count, float, initialized = true,, context.default_allocator); // 0 is the padding
|
||||
input := NewArray(plane * 3 * rows, float, initialized = true,, context.default_allocator); // 0 is the padding
|
||||
defer array_free(input,, context.default_allocator);
|
||||
for job, n: jobs {
|
||||
ocr_sample_crop(job.image, job.corners, job.width, height, input.data + plane * 3 * n, width, plane, REC_MEAN, REC_STD);
|
||||
}
|
||||
|
||||
shape := s64.[jobs.count, 3, height, width];
|
||||
out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
|
||||
if !ok return false, error;
|
||||
shape := s64.[rows, 3, height, width];
|
||||
scores: *float;
|
||||
dims: [] s64;
|
||||
out: Ocr_Output;
|
||||
defer ocr_release_output(*out);
|
||||
if out.dim_count != 3 || out.dims[0] != jobs.count return false, tprint("The recognizer gave an output of shape %.", array_view(out.dims, 0, out.dim_count));
|
||||
steps := out.dims[1];
|
||||
classes := out.dims[2];
|
||||
#if OS == .MACOS {
|
||||
coreml_out: Coreml_Output;
|
||||
defer coreml_release_output(*coreml_out);
|
||||
}
|
||||
#if OS == .MACOS if engine.rec_coreml {
|
||||
ok: bool; error: string;
|
||||
coreml_out, ok, error = coreml_predict(engine.rec_coreml, input, shape);
|
||||
if !ok return false, error;
|
||||
scores, dims = coreml_out.data, array_view(coreml_out.dims, 0, coreml_out.dim_count);
|
||||
}
|
||||
if !scores {
|
||||
ok: bool; error: string;
|
||||
out, ok, error = ocr_run(*engine.rec, engine.memory_info, input, shape);
|
||||
if !ok return false, error;
|
||||
scores, dims = out.data, array_view(out.dims, 0, out.dim_count);
|
||||
}
|
||||
if dims.count != 3 || dims[0] != rows return false, tprint("The recognizer gave an output of shape %.", dims);
|
||||
steps := dims[1];
|
||||
classes := dims[2];
|
||||
if classes != engine.dictionary.count + 2 {
|
||||
return false, tprint("The recognizer has % classes, but its dictionary % characters (plus blank and space, % expected).", classes, engine.dictionary.count, engine.dictionary.count + 2);
|
||||
}
|
||||
@ -55,7 +78,7 @@ ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, er
|
||||
probability_sum: float64;
|
||||
characters := 0;
|
||||
for t: 0..steps - 1 {
|
||||
row := out.data + (n * steps + t) * classes;
|
||||
row := scores + (n * steps + t) * classes;
|
||||
best := 0;
|
||||
best_p := row[0];
|
||||
for c: 1..classes - 1 if row[c] > best_p { best = c; best_p = row[c]; }
|
||||
|
||||
@ -53,7 +53,8 @@ ort_cuda_check :: (what: string) -> available: bool, explanation: string {
|
||||
// (onnxruntime 1.30.0) and is slower than the CPU for these models: OCR's take a new image size
|
||||
// nearly every run, which CoreML compiles for each time, and SigLIP 2 converts only with every
|
||||
// dimension fixed (its ML Program conversion fails on the patch embedding), then runs at about
|
||||
// half the CPU's speed. So on macOS the models run on the CPU.
|
||||
// half the CPU's speed. So onnxruntime runs them on the CPU on macOS; OCR and labels use native
|
||||
// Core ML instead when their converted models are there (modules/Core_ML).
|
||||
#if OS == .MACOS {
|
||||
return false, tprint("% runs on the CPU: onnxruntime has no CUDA on macOS, and its CoreML provider was measured slower than the CPU on Apple silicon.", what);
|
||||
}
|
||||
|
||||
@ -283,7 +283,7 @@ ocr_ui_thread_proc :: (thread: *Thread) -> s64 {
|
||||
ocr_ui.engine = engine;
|
||||
free(ocr_ui.engine_error);
|
||||
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 "");
|
||||
if ok log("[OCR] % on %.%", task.model, ocr_provider_name(engine), 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;
|
||||
|
||||
@ -463,7 +463,7 @@ ocr_read_and_store :: (items: [] Scan_Item, set: *Label_Set) -> [] string {
|
||||
}
|
||||
|
||||
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");
|
||||
engine_text := ifx ocr_ui.engine.provider == .COREML then "Core ML, GPU" else tprint("onnxruntime %, %", ort.version, ocr_provider_name(ocr_ui.engine));
|
||||
for n: 0..items.count - 1 {
|
||||
if errors[n] continue;
|
||||
section: Mcap_Ocr_Section;
|
||||
@ -715,7 +715,7 @@ ocr_scan_progress_ui :: () { using ImGui;
|
||||
ProgressBar(cast(float) done / max(total, 1), .{-FLOAT32_MIN, 0}, temp_c_string(ocr_scan_progress_text(done, total)));
|
||||
if scan.erase_ocr || scan.erase_labels || scan.erase_segments Text("% elapsed, % left, % captures/s, % failed", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed);
|
||||
else if scan.segmenting Text("% elapsed, % left, % images/s, % failed", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed); // the engine is the OCR thread's while it runs
|
||||
else Text("% elapsed, % left, % captures/s, % failed, on %", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed, ifx ocr_ui.engine && ocr_ui.engine.provider == .CUDA then "the GPU" else "the CPU");
|
||||
else Text("% elapsed, % left, % captures/s, % failed, on %", ocr_duration_text(elapsed), ocr_duration_text(left), FF(rate, 1), failed, ifx ocr_ui.engine && ocr_gpu_status(ocr_ui.engine) 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); }
|
||||
} else {
|
||||
@ -811,7 +811,7 @@ ocr_ui_frame :: () { using ImGui;
|
||||
ocr_scan_ui();
|
||||
} else {
|
||||
in_use, explanation := ocr_gpu_status(ocr_ui.engine);
|
||||
Text("Running on: %", ifx in_use then "the GPU (CUDA)" else "the CPU");
|
||||
Text("Running on: %", ifx in_use then tprint("the GPU (%)", ocr_provider_name(ocr_ui.engine)) else "the CPU");
|
||||
SameLine();
|
||||
TextColored(app_theme.palette.dim, "onnxruntime %, % worker threads", ort.version, ocr_ui.engine.workers);
|
||||
Separator();
|
||||
|
||||
@ -1163,7 +1163,7 @@ viewer_ui :: () { using ImGui;
|
||||
if IsItemHovered() {
|
||||
engine: *Ocr_Engine;
|
||||
if !ocr_ui.thread_busy engine = ocr_ui.engine; // the background thread owns it while busy
|
||||
if engine SetTooltip("OCR with %, on %", engine.model.name, ifx engine.provider == .CUDA then "the GPU" else "the CPU");
|
||||
if engine SetTooltip("OCR with %, on %", engine.model.name, ifx ocr_gpu_status(engine) then tprint("the GPU (%)", ocr_provider_name(engine)) else "the CPU");
|
||||
else SetTooltip("OCR is not loaded, is loading, or could not start; the OCR window says which");
|
||||
}
|
||||
|
||||
|
||||
@ -169,7 +169,7 @@ turn_tests :: () {
|
||||
if !created return;
|
||||
defer ocr_destroy(engine);
|
||||
expect(engine.orientation.session != null, "% has an orientation classifier", OCR_DEFAULT_MODEL);
|
||||
emit(" running on %; classifier input % x %\n", ifx engine.provider == .CUDA then "CUDA" else "the CPU", engine.orientation_width, engine.orientation_height);
|
||||
emit(" running on %; classifier input % x %\n", ocr_provider_name(engine), engine.orientation_width, engine.orientation_height);
|
||||
|
||||
level, drawn := drawn_lines();
|
||||
expect(drawn, "the test text is drawn");
|
||||
|
||||
172
tools/export_ppocr_coreml.py
Normal file
172
tools/export_ppocr_coreml.py
Normal file
@ -0,0 +1,172 @@
|
||||
# Converts a PP-OCRv5 detector and recognizer to Core ML models for modules/OCR on macOS, then checks
|
||||
# each against the ONNX model on the CPU. Run once, on a Mac, after modules/OCR/fetch_models.jai;
|
||||
# mcap only uses the .mlpackage files it writes (into modules/OCR/models, beside the ONNX ones).
|
||||
#
|
||||
# python3 -m venv ~/mcap-coreml
|
||||
# ~/mcap-coreml/bin/pip install torch==2.7.0 coremltools==9.0 onnx onnxruntime pillow
|
||||
# ~/mcap-coreml/bin/pip install --no-deps onnx2torch torchvision==0.22.0
|
||||
# ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py [--model mobile|server]
|
||||
#
|
||||
# There is no PyTorch source for PP-OCR (it is PaddlePaddle's), and coremltools no longer reads ONNX, so
|
||||
# each model goes ONNX -> PyTorch (onnx2torch) -> Core ML, as an ML Program in float32: float16 moves
|
||||
# the detector's text pixels and changes recognized text, and the Neural Engine gets it wrong.
|
||||
#
|
||||
# The shapes. Core ML on the GPU cannot take a flexible height and width for the detector (its
|
||||
# depthwise convolutions fail to compile: "invalid axis" in mps.expand_dims), so both models take a
|
||||
# list of sizes (at most 128 a model), which mcap pads its inputs up to; models.config's coreml_ keys
|
||||
# say the same steps, and must be kept in step with these:
|
||||
# detector: 1 x 3 x H x W, H a multiple of 128 up to 1536 (det_tile_height), W of 256 up to 2560 (det_max_side)
|
||||
# recognizer: 16 x 3 x 48 x W, W a multiple of 64 up to 4096 (REC_MAX_WIDTH); fewer lines are padded to 16
|
||||
# The recognizer's Reshapes compute their target shape from tensor sizes, which tracing freezes at the
|
||||
# traced width; they are rewritten first as constant shapes of 0 (copy the dimension) and -1 (infer it),
|
||||
# worked out from the shapes onnxruntime gives at several sizes, and checked to give the same output.
|
||||
#
|
||||
# Measured on an M5 Pro against onnxruntime on the CPU, each at the same (padded) size: the detector
|
||||
# about 5 times as fast, the recognizer about 4 times for 16 lines, with the same text pixels and text.
|
||||
|
||||
import argparse, itertools, os, sys, time
|
||||
import numpy as np
|
||||
import onnx, onnxruntime as ort
|
||||
from onnx import helper, numpy_helper
|
||||
import torch, onnx2torch
|
||||
import coremltools as ct
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
MODELS = os.path.join(HERE, "..", "modules", "OCR", "models")
|
||||
|
||||
DET_HEIGHTS = list(range(128, 1536 + 1, 128))
|
||||
DET_WIDTHS = list(range(256, 2560 + 1, 256))
|
||||
REC_BATCH = 16
|
||||
REC_HEIGHT = 48
|
||||
REC_WIDTHS = list(range(64, 4096 + 1, 64))
|
||||
|
||||
# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only
|
||||
# 0-dimensional arrays can be converted"); the same cast, with the element taken out first.
|
||||
from coremltools.converters.mil.frontend.torch import ops as ct_ops
|
||||
_original_cast = ct_ops._cast
|
||||
def _cast_one_element(context, node, dtype, dtype_name):
|
||||
x = ct_ops._get_inputs(context, node, expected=1)[0]
|
||||
if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
|
||||
context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
|
||||
return
|
||||
_original_cast(context, node, dtype, dtype_name)
|
||||
ct_ops._cast = _cast_one_element
|
||||
|
||||
|
||||
def constant_reshapes(model, input_shapes):
|
||||
"""The model with each computed Reshape target made a constant of 0s and -1s."""
|
||||
graph = model.graph
|
||||
initializers = {i.name for i in graph.initializer}
|
||||
producers = {o: n for n in graph.node for o in n.output}
|
||||
computed = [n for n in graph.node if n.op_type == "Reshape" and n.input[1] not in initializers
|
||||
and not (n.input[1] in producers and producers[n.input[1]].op_type == "Constant")]
|
||||
if not computed:
|
||||
return model
|
||||
probe = onnx.ModelProto(); probe.CopyFrom(model)
|
||||
names = []
|
||||
for n in computed:
|
||||
for t in (n.input[0], n.output[0]):
|
||||
if t not in names:
|
||||
names.append(t)
|
||||
probe.graph.output.append(helper.make_tensor_value_info(t, onnx.TensorProto.FLOAT, None))
|
||||
session = ort.InferenceSession(probe.SerializeToString(), providers=["CPUExecutionProvider"])
|
||||
name = session.get_inputs()[0].name
|
||||
shapes = []
|
||||
for shape in input_shapes:
|
||||
outs = session.run(names, {name: np.zeros(shape, np.float32)})
|
||||
shapes.append({t: o.shape for t, o in zip(names, outs)})
|
||||
for n in computed:
|
||||
ins = [s[n.input[0]] for s in shapes]
|
||||
outs = [s[n.output[0]] for s in shapes]
|
||||
target = []
|
||||
for d in range(len(outs[0])):
|
||||
sizes = {o[d] for o in outs}
|
||||
if len(sizes) == 1:
|
||||
target.append(sizes.pop())
|
||||
elif all(len(i) > d and i[d] == o[d] for i, o in zip(ins, outs)):
|
||||
target.append(0)
|
||||
else:
|
||||
target.append(-1)
|
||||
if target.count(-1) > 1:
|
||||
sys.exit(f"{n.name}: more than one dimension to infer ({target})")
|
||||
constant = n.name + "_constant_shape"
|
||||
graph.initializer.append(numpy_helper.from_array(np.array(target, dtype=np.int64), constant))
|
||||
n.input[1] = constant
|
||||
print(f" {n.name}: target shape {target}")
|
||||
return model
|
||||
|
||||
|
||||
def convert(onnx_model, example, shape, path):
|
||||
module = onnx2torch.convert(onnx_model).eval()
|
||||
with torch.no_grad():
|
||||
traced = torch.jit.trace(module, torch.from_numpy(example))
|
||||
ml = ct.convert(traced, inputs=[ct.TensorType(name="x", shape=shape)], outputs=[ct.TensorType(name="y")],
|
||||
convert_to="mlprogram", compute_precision=ct.precision.FLOAT32,
|
||||
minimum_deployment_target=ct.target.macOS14)
|
||||
ml.save(path)
|
||||
return ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)
|
||||
|
||||
|
||||
def ctc(logits):
|
||||
texts = []
|
||||
for row in logits.argmax(-1):
|
||||
seq, last = [], -1
|
||||
for i in row:
|
||||
if i != last and i != 0:
|
||||
seq.append(int(i))
|
||||
last = i
|
||||
texts.append(seq)
|
||||
return texts
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", choices=["mobile", "server"], default="mobile")
|
||||
args = parser.parse_args()
|
||||
rng = np.random.default_rng(0)
|
||||
|
||||
# Detector.
|
||||
stem = f"ch_PP-OCRv5_det_{args.model}"
|
||||
onnx_path = os.path.join(MODELS, stem + ".onnx")
|
||||
out_path = os.path.join(MODELS, stem + ".mlpackage")
|
||||
shapes = [(1, 3, h, w) for h in DET_HEIGHTS for w in DET_WIDTHS]
|
||||
start = time.perf_counter()
|
||||
det = convert(onnx.load(onnx_path), np.zeros((1, 3, 768, 1280), np.float32),
|
||||
ct.EnumeratedShapes(shapes=shapes, default=(1, 3, 768, 1280)), out_path)
|
||||
print(f"Wrote {out_path} ({len(shapes)} sizes) in {time.perf_counter() - start:.1f} s.")
|
||||
reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
|
||||
worst = 0.0
|
||||
for (h, w) in [(768, 1280), (1536, 2560), (128, 256)]:
|
||||
x = rng.standard_normal((1, 3, h, w)).astype(np.float32)
|
||||
a = reference.run(None, {"x": x})[0]
|
||||
b = det.predict({"x": x})["y"]
|
||||
flips = float(np.mean((a > 0.3) != (b > 0.3)))
|
||||
worst = max(worst, flips)
|
||||
print(f" {h} x {w}: max difference {np.abs(a - b).max():.2e}, pixels across 0.3: {flips * 100:.4f}%")
|
||||
if worst > 0.0001:
|
||||
sys.exit("The Core ML detector is too far from the ONNX one.")
|
||||
|
||||
# Recognizer.
|
||||
stem = f"ch_PP-OCRv5_rec_{args.model}"
|
||||
onnx_path = os.path.join(MODELS, stem + ".onnx")
|
||||
out_path = os.path.join(MODELS, stem + ".mlpackage")
|
||||
print("Recognizer: making its reshapes constant")
|
||||
model = constant_reshapes(onnx.load(onnx_path), [(n, 3, REC_HEIGHT, w) for n in (1, 2) for w in (320, 640)])
|
||||
shapes = [(REC_BATCH, 3, REC_HEIGHT, w) for w in REC_WIDTHS]
|
||||
start = time.perf_counter()
|
||||
rec = convert(model, np.zeros((REC_BATCH, 3, REC_HEIGHT, 640), np.float32),
|
||||
ct.EnumeratedShapes(shapes=shapes, default=(REC_BATCH, 3, REC_HEIGHT, 640)), out_path)
|
||||
print(f"Wrote {out_path} ({len(shapes)} widths) in {time.perf_counter() - start:.1f} s.")
|
||||
reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
|
||||
for w in (64, 640, 4096):
|
||||
x = rng.standard_normal((REC_BATCH, 3, REC_HEIGHT, w)).astype(np.float32) * 0.5
|
||||
a = reference.run(None, {"x": x})[0]
|
||||
b = rec.predict({"x": x})["y"]
|
||||
same = np.mean([p == q for p, q in zip(ctc(a), ctc(b))])
|
||||
print(f" width {w}: max difference {np.abs(a - b).max():.2e}, decoded the same {same * 100:.0f}%")
|
||||
if same < 1:
|
||||
sys.exit("The Core ML recognizer reads differently from the ONNX one.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user