diff --git a/README.md b/README.md index 37eeb28..580651d 100644 --- a/README.md +++ b/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. diff --git a/check_targets.jai b/check_targets.jai index c70023d..ee72276 100644 --- a/check_targets.jai +++ b/check_targets.jai @@ -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; diff --git a/first.jai b/first.jai index b43f186..906426e 100644 --- a/first.jai +++ b/first.jai @@ -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."); } diff --git a/modules/Core_ML/module.jai b/modules/Core_ML/module.jai index 3127ba7..08798e7 100644 --- a/modules/Core_ML/module.jai +++ b/modules/Core_ML/module.jai @@ -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); diff --git a/modules/OCR/detect.jai b/modules/OCR/detect.jai index b5fb184..42cfa5f 100644 --- a/modules/OCR/detect.jai +++ b/modules/OCR/detect.jai @@ -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; diff --git a/modules/OCR/models.config b/modules/OCR/models.config index 28b58f6..a2e5214 100644 --- a/modules/OCR/models.config +++ b/modules/OCR/models.config @@ -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)" diff --git a/modules/OCR/models.jai b/modules/OCR/models.jai index 0f6131d..0ab7c73 100644 --- a/modules/OCR/models.jai +++ b/modules/OCR/models.jai @@ -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,14 +113,67 @@ 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; - // 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; + 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); diff --git a/modules/OCR/module.jai b/modules/OCR/module.jai index 203b5cc..b22d46b 100644 --- a/modules/OCR/module.jai +++ b/modules/OCR/module.jai @@ -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"; diff --git a/modules/OCR/recognize.jai b/modules/OCR/recognize.jai index 90fd218..0a888b5 100644 --- a/modules/OCR/recognize.jai +++ b/modules/OCR/recognize.jai @@ -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]; } diff --git a/modules/onnxruntime/cuda.jai b/modules/onnxruntime/cuda.jai index 14ddd87..60f4ccf 100644 --- a/modules/onnxruntime/cuda.jai +++ b/modules/onnxruntime/cuda.jai @@ -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); } diff --git a/src/app/ocr_ui.jai b/src/app/ocr_ui.jai index fdc6091..248fd99 100644 --- a/src/app/ocr_ui.jai +++ b/src/app/ocr_ui.jai @@ -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(); diff --git a/src/app/viewer.jai b/src/app/viewer.jai index 9767465..9ca7ed4 100644 --- a/src/app/viewer.jai +++ b/src/app/viewer.jai @@ -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"); } diff --git a/src/test/ocr_tests.jai b/src/test/ocr_tests.jai index ed36589..fb8a5b2 100644 --- a/src/test/ocr_tests.jai +++ b/src/test/ocr_tests.jai @@ -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"); diff --git a/tools/export_ppocr_coreml.py b/tools/export_ppocr_coreml.py new file mode 100644 index 0000000..0e68ebe --- /dev/null +++ b/tools/export_ppocr_coreml.py @@ -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()