mcap/tools/ocr_bench.jai

// 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] [-batch N] [-limit N] [-out FOLDER]
//
//   -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
//   -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] [-batch N] [-limit N] [-out FOLDER]\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;
  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 "-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;
      create_start := current_time_monotonic();
      engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu);
      if !ok { print("% : %\n\n", model_name, error); break; }
      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)\n", model_name, provider, FF(create_seconds, 2), engine.workers);

      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);
      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";