mcap/modules/OCR/models.jai

// models.config, and running a model: creating its onnxruntime session on the GPU or the CPU, and
// running it on a float tensor.

Ocr_Model_Info :: struct {
  name: string;
  title: string;
  detector, recognizer, dictionary: string;   // file names in the models folder
  textline_ori: string;                        // the orientation classifier's file name; "" for none
  det_max_side:    s32   = 2560;
  det_tile_height: s32   = 1536;
  det_threshold:   float = 0.3;
  box_threshold:   float = 0.6;
  unclip_ratio:    float = 1.5;
  rec_height:      s32   = 48;
  rec_batch_gpu:   s32   = 16;
  rec_batch_cpu:   s32   = 1;
}

OCR_MODELS_CONFIG :: "models.config";

// The models models.config lists, in its order; allocated with the context allocator. A line that
// cannot be read is logged and skipped, as in the Ex4 formats.
ocr_list_models :: (models_folder: string) -> [] Ocr_Model_Info, error: string {
  models: [..] Ocr_Model_Info;
  path := tprint("%/%", models_folder, OCR_MODELS_CONFIG);
  text, ok := read_entire_file(path, log_errors = false,, temp);
  if !ok  return models, tprint("Cannot read %.", path);

  model: *Ocr_Model_Info;
  version_seen := false;
  line_number := 0;
  for raw_line: split(text, "\n",, temp) {
    line_number += 1;
    line := trim(ocr_strip_comment(raw_line));
    if !line.count  continue;
    if begins_with(line, "@version") { version_seen = true; continue; }
    if !version_seen  return models, tprint("% does not start with @version 1.", path);

    if begins_with(line, "[model ") && ends_with(line, "]") {
      model = array_add(*models);
      model.* = .{};
      model.name = copy_string(trim(slice(line, 7, line.count - 8)));
      continue;
    }
    colon := find_index_from_left(line, #char ":");
    if colon < 0 || !model { log_error("[OCR] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; }
    key   := trim(slice(line, 0, colon));
    value := trim(advance(line, colon + 1));
    if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\""  value = slice(value, 1, value.count - 2);

    number_text := value;
    number, number_ok := string_to_float(number_text);
    if key == {
      case "title";           model.title      = copy_string(value);
      case "detector";        model.detector   = copy_string(value);
      case "recognizer";      model.recognizer = copy_string(value);
      case "dictionary";      model.dictionary = copy_string(value);
      case "textline_ori";    model.textline_ori = copy_string(value);
      case "det_max_side";    if number_ok  model.det_max_side    = cast(s32) number;
      case "det_tile_height"; if number_ok  model.det_tile_height = cast(s32) number;
      case "det_threshold";   if number_ok  model.det_threshold   = number;
      case "box_threshold";   if number_ok  model.box_threshold   = number;
      case "unclip_ratio";    if number_ok  model.unclip_ratio    = number;
      case "rec_height";      if number_ok  model.rec_height      = cast(s32) number;
      case "rec_batch_gpu";   if number_ok  model.rec_batch_gpu   = cast(s32) number;
      case "rec_batch_cpu";   if number_ok  model.rec_batch_cpu   = cast(s32) number;
      case;  log_error("[OCR] %:%: unknown key %. Ignored.", path, line_number, key);
    }
  }
  return models, "";
}

#scope_module

// "#" opens a comment at the start of a line or after whitespace, outside quotes.
ocr_strip_comment :: (line: string) -> string {
  in_quotes := false;
  for i: 0..line.count - 1 {
    c := line[i];
    if c == #char "\""  in_quotes = !in_quotes;
    else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t")  return slice(line, 0, i);
  }
  return line;
}

// The model's sessions on the GPU or the CPU: detector, recognizer and, when it has one, orientation
// classifier. On failure, the reason; the sessions made so far are left for ocr_release_sessions.
ocr_create_sessions :: (engine: *Ocr_Engine, models_folder: string, gpu: bool, intra_op_threads: s32) -> ok: bool, error: string {
  model := *engine.model;
  ok: bool;
  error: string;
  // The detector's and recognizer's inputs change size with nearly every run: their arenas give back
  // what no run is using once per call (ocr_shrink_arena), or they keep a block of every size.
  engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads, engine.options, shrink_arena = true);
  if !ok  return false, error;
  engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads, engine.options, shrink_arena = true);
  if !ok  return false, error;
  if model.textline_ori {
    engine.orientation, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.textline_ori), gpu, intra_op_threads, engine.options);
    if !ok  return false, error;
    // PP-LCNet text line orientation takes 160 by 80; a model that says otherwise is taken at its word.
    shape := engine.orientation.input_shape;
    engine.orientation_height = ifx engine.orientation.input_rank == 4 && shape[2] > 0 then cast(s32) shape[2] else 80;
    engine.orientation_width  = ifx engine.orientation.input_rank == 4 && shape[3] > 0 then cast(s32) shape[3] else 160;
  }
  return true, "";
}

ocr_release_sessions :: (engine: *Ocr_Engine) {
  ocr_release_session(*engine.det);
  ocr_release_session(*engine.rec);
  ocr_release_session(*engine.orientation);
}

// Runs each of a CUDA engine's sessions once, on a small input, before it is used or destroyed.
// With cuBLASLt 13.8 (CUDA 13, onnxruntime 1.30.0, RTX 5080 laptop), the process's first CUDA
// engine destroyed without ever having run left the next engine to crash in cublasLtCtxInit, with a
// DEP access violation, on its first run; an engine that has run once leaves nothing broken. It also
// moves the first run's setup into creation. On failure, the reason.
// On CUDA, the session's arena (the detector's or the recognizer's) gives back every block no run is
// using: called by ocr_process once a stage's runs are all done, as a run on a small input, since
// onnxruntime shrinks an arena only at the end of a run. Once per run instead cost a third to a half
// of OCR's speed.
ocr_shrink_arena :: (engine: *Ocr_Engine, session: *Ocr_Session) {
  if engine.provider != .CUDA || !session.session  return;
  side := ifx session == *engine.rec then cast(s64) engine.model.rec_height else 32;
  shape := s64.[1, 3, side, side];
  input := NewArray(3 * side * side, float, initialized = true,, temp);
  out, ok := ocr_run(session, engine.memory_info, input, shape, shrink = true);
  if ok  ocr_release_output(*out);
}

ocr_warm_up :: (engine: *Ocr_Engine) -> ok: bool, error: string {
  Warm_Up :: struct { session: *Ocr_Session; height, width: s64; }
  runs := Warm_Up.[
    .{*engine.det, 64, 64},
    .{*engine.rec, engine.model.rec_height, engine.model.rec_height * 2},
    .{*engine.orientation, engine.orientation_height, engine.orientation_width},
  ];
  for runs {
    if !it.session.session  continue;   // no orientation classifier
    shape := s64.[1, 3, it.height, it.width];
    input := NewArray(3 * it.height * it.width, float, initialized = true,, context.default_allocator);
    defer array_free(input,, context.default_allocator);
    out, ok, error := ocr_run(it.session, engine.memory_info, input, shape);
    if !ok  return false, tprint("The first run on the GPU failed: %", error);
    ocr_release_output(*out);
  }
  return true, "";
}

// A session for the model file, set up as Ocr_Options says (modules/onnxruntime/session.jai).
ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, ocr_options: Ocr_Options, shrink_arena := false) -> Ocr_Session, ok: bool, error: string {
  settings := Ort_Gpu_Settings.{arena_extend_strategy = ocr_options.cuda_arena_extend_strategy, cudnn_max_workspace = ocr_options.cudnn_max_workspace, shrink_arena = shrink_arena};
  session, ok, error := ort_create_session(env, path, gpu, intra_op_threads, settings, missing_hint = "run `jai fetch_models.jai` in modules/OCR and rebuild");
  return session, ok, error;
}

// onnxruntime's session and run, under the names the OCR code has always used.
Ocr_Session        :: Ort_Session;
Ocr_Output         :: Ort_Output;
ocr_run            :: ort_run_float;
ocr_release_output :: ort_release_output;
ocr_release_session :: ort_release_session;