173 lines
8.0 KiB
Plaintext
173 lines
8.0 KiB
Plaintext
// ocr-bench: runs OCR over a folder of .mcap captures with each model, on the GPU when it can and
|
|
// on the CPU, and reports the time per image and per stage. Each capture's lines are written to
|
|
// OUT/<model>-<provider>/<name>.txt, to read against the images.
|
|
//
|
|
// ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER]
|
|
// [-workers N] [-rec-batch N] [-warmup N]
|
|
//
|
|
// -models the models to run (default: every model in ocr-models/models.config)
|
|
// -cpu the CPU only; by default the GPU too, when it is available
|
|
// -gpu-only the GPU only
|
|
// -workers on the GPU, threads preparing images and decoding results (default Ocr_Options')
|
|
// -rec-batch text lines per recognizer run, in place of models.config's
|
|
// -warmup untimed passes over every image before the timed one (default 0)
|
|
// -batch images per ocr_process call (default 8)
|
|
// -limit at most this many captures (default all)
|
|
// -out where the text goes (default ocr-out beside this program)
|
|
//
|
|
// The models are read from ocr-models beside this program, where the build copies them.
|
|
|
|
USAGE :: "Usage: ocr-bench MCAP_FOLDER [-models a,b] [-cpu | -gpu-only] [-batch N] [-limit N] [-out FOLDER] [-workers N] [-rec-batch N] [-warmup N]\n";
|
|
|
|
main :: () {
|
|
runtime_support_headless = true; // a command-line tool: a failure is reported on the console and exits with 1, no message box
|
|
args := get_command_line_arguments();
|
|
if args.count < 2 { print(USAGE); exit(2); }
|
|
folder := args[1];
|
|
model_names: [..] string;
|
|
cpu_only := false;
|
|
gpu_only := false;
|
|
options: Ocr_Options;
|
|
rec_batch: s32 = 0; // 0: as models.config says
|
|
warmup := 0;
|
|
batch := 8;
|
|
limit := S64_MAX_VALUE;
|
|
exe_folder := path_strip_filename(get_path_of_running_executable());
|
|
// Lasting copies: temporary storage is reset after every batch.
|
|
out_folder := sprint("%ocr-out", exe_folder);
|
|
i := 2;
|
|
while i < args.count {
|
|
arg := args[i];
|
|
value := ifx i + 1 < args.count then args[i + 1] else "";
|
|
if arg == {
|
|
case "-models"; for split(value, ",") array_add(*model_names, trim(it)); i += 2;
|
|
case "-cpu"; cpu_only = true; i += 1;
|
|
case "-gpu-only"; gpu_only = true; i += 1;
|
|
case "-workers"; options.gpu_workers = cast(s32) max(1, string_to_int(value)); i += 2;
|
|
case "-rec-batch"; rec_batch = cast(s32) max(1, string_to_int(value)); i += 2;
|
|
case "-warmup"; warmup = max(0, string_to_int(value)); i += 2;
|
|
case "-batch"; batch = max(1, string_to_int(value)); i += 2;
|
|
case "-limit"; limit = max(1, string_to_int(value)); i += 2;
|
|
case "-out"; out_folder = value; i += 2;
|
|
case; print("Unknown option %.\n%", arg, USAGE); exit(2);
|
|
}
|
|
}
|
|
models_folder := sprint("%ocr-models", exe_folder);
|
|
if !model_names.count {
|
|
models, error := ocr_list_models(models_folder);
|
|
if error { print("%\n", error); exit(1); }
|
|
for models array_add(*model_names, it.name);
|
|
}
|
|
|
|
// The captures, decoded once: the original image of each, as the camera took it.
|
|
Capture :: struct { name: string; pixels: *u8; width, height: s32; }
|
|
captures: [..] Capture;
|
|
decode_start := current_time_monotonic();
|
|
for file_list(folder) {
|
|
if captures.count >= limit break;
|
|
if !ends_with_nocase(it, ".mcap") continue;
|
|
data, ok := read_entire_file(it);
|
|
if !ok continue;
|
|
m, parsed, error := mcap_parse(cast([] u8) data);
|
|
if !parsed { print("Skipping %: %\n", it, error); continue; }
|
|
w, h, channels: s32;
|
|
pixels := stbi_load_from_memory(m.original.data, cast(s32) m.original.count, *w, *h, *channels, 4);
|
|
if !pixels { print("Skipping %: its image cannot be decoded.\n", it); continue; }
|
|
array_add(*captures, .{copy_string(path_strip_extension(path_filename(it))), pixels, w, h});
|
|
free(data);
|
|
}
|
|
if !captures.count { print("No .mcap files in %.\n", folder); exit(1); }
|
|
quick_sort(captures, (a: Capture, b: Capture) -> s64 { return compare(a.name, b.name); });
|
|
megapixels: float64;
|
|
for captures megapixels += cast(float64) it.width * it.height / 1_000_000;
|
|
print("% captures (% megapixels) decoded in % s.\n\n", captures.count, FF(megapixels, 1), FF(seconds_since(decode_start), 2));
|
|
|
|
images := NewArray(captures.count, Ocr_Image);
|
|
for captures images[it_index] = .{it.pixels, it.width, it.height, it.width * 4};
|
|
|
|
for model_name: model_names {
|
|
for prefer_gpu: bool.[true, false] {
|
|
if prefer_gpu && cpu_only continue;
|
|
if !prefer_gpu && gpu_only continue;
|
|
create_start := current_time_monotonic();
|
|
engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu, options);
|
|
if !ok { print("% : %\n\n", model_name, error); break; }
|
|
if rec_batch engine.rec_batch = rec_batch;
|
|
create_seconds := seconds_since(create_start);
|
|
in_use, explanation := ocr_gpu_status(engine);
|
|
provider := ocr_provider(engine);
|
|
if prefer_gpu && !in_use {
|
|
// Said once; the CPU run that follows is the same thing.
|
|
print("%\n\n", explanation);
|
|
ocr_destroy(engine);
|
|
continue;
|
|
}
|
|
print("== % on % (sessions created in % s, % worker threads, % lines per recognizer run)\n", model_name, provider, FF(create_seconds, 2), engine.workers, engine.rec_batch);
|
|
|
|
text_folder := sprint("%/%-%", out_folder, model_name, to_lower_copy(tprint("%", provider),, temp));
|
|
make_directory_if_it_does_not_exist(text_folder, recursive = true);
|
|
|
|
// The first image alone: the first run pays for setting up kernels and memory.
|
|
first_start := current_time_monotonic();
|
|
ocr_process(engine, array_view(images, 0, 1),, temp);
|
|
first_seconds := seconds_since(first_start);
|
|
// Untimed passes over everything, so the timed one does not pay for setting up each new shape.
|
|
for 1..warmup {
|
|
for start: 0..images.count - 1 {
|
|
if start % batch continue;
|
|
ocr_process(engine, array_view(images, start, min(batch, images.count - start)),, temp);
|
|
reset_temporary_storage();
|
|
}
|
|
}
|
|
ocr_reset_stats(engine);
|
|
|
|
start := current_time_monotonic();
|
|
lines, characters, failures: s64;
|
|
done := 0;
|
|
while done < images.count {
|
|
count := min(batch, images.count - done);
|
|
results := ocr_process(engine, array_view(images, done, count),, temp);
|
|
for result, n: results {
|
|
name := captures[done + n].name;
|
|
if result.failed { failures += 1; print(" % failed: %\n", name, result.error); continue; }
|
|
builder: String_Builder;
|
|
builder.allocator = temp;
|
|
for result.lines {
|
|
print_to_builder(*builder, "% % % % % %\n", it.x, it.y, it.width, it.height, FF(it.confidence, 3), it.text);
|
|
characters += it.text.count;
|
|
}
|
|
lines += result.lines.count;
|
|
write_entire_file(tprint("%/%.txt", text_folder, name), *builder);
|
|
}
|
|
done += count;
|
|
reset_temporary_storage();
|
|
}
|
|
seconds := seconds_since(start);
|
|
stats := ocr_stats(engine);
|
|
print(" first image alone: % s (warm-up)\n", FF(first_seconds, 3));
|
|
print(" all % images: % s, % ms per image, % megapixels/s, % failed\n", images.count, FF(seconds, 2), FF(1000 * seconds / images.count, 1), FF(megapixels / seconds, 2), failures);
|
|
print(" found: % lines, % bytes of text\n", lines, characters);
|
|
print(" stage time (summed over threads): detection % s, orientation % s, recognition % s\n", FF(stats.detect_us / 1_000_000.0, 2), FF(stats.orientation_us / 1_000_000.0, 2), FF(stats.recognize_us / 1_000_000.0, 2));
|
|
print(" lines read better turned: %\n", stats.turned);
|
|
print(" text in %\n\n", text_folder);
|
|
ocr_destroy(engine);
|
|
}
|
|
}
|
|
}
|
|
|
|
seconds_since :: (start: Apollo_Time) -> float64 {
|
|
return to_float64_seconds(current_time_monotonic() - start);
|
|
}
|
|
|
|
S64_MAX_VALUE :: 0x7FFF_FFFF_FFFF_FFFF;
|
|
|
|
#import "Basic";
|
|
#import "String";
|
|
#import "File";
|
|
#import "File_Utilities";
|
|
#import "Sort";
|
|
#import "System";
|
|
#import "stb_image";
|
|
#import "Capture_File";
|
|
#import "OCR";
|