mcap/modules/OCR/recognize.jai

// Recognition: the text of each detected line. Each box is sampled along its own sides into a level
// crop rec_height high, keeping its shape, and normalized to -1..1 in BGR order, as PaddleOCR's
// recognizers expect. Lines of similar width are batched, padded to the widest in the batch.
// The output is a probability per character class per step; greedy CTC decoding takes the best
// class at each step, merges repeats and drops blanks.
//
// A box's corners only say which way its text runs up to a half turn. The orientation classifier,
// when the model has one, says which lines look upside down; such a line is read both ways, as is a
// box taken as a line read downwards, and the more confident reading is kept.

Line_Job :: struct {
  image: *Ocr_Image;
  corners: [4] Vector2;  // as Detected_Box's
  width: s32;            // after scaling to rec_height
  primary := -1;         // for another way of reading a line: the index of the line's first job
  // Filled in by recognition:
  text: string;          // default allocator
  confidence: float;
}

// Recognizes jobs (one batch) in place.
ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, error: string {
  if !jobs.count  return true, "";
  height := engine.model.rec_height;
  width: s32 = 0;
  for jobs  width = max(width, it.width);
  width = max(width, height);   // the network needs a few steps even for one character
  if engine.provider == .CUDA {   // fewer shapes on the GPU; on the CPU, padding only costs
    step := max(engine.options.rec_width_step, 1);
    width = (width + step - 1) / step * step;
  }

  plane := cast(s64) height * width;
  input := NewArray(plane * 3 * jobs.count, 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];
  reserved := ocr_budget_recognizer(engine, jobs.count, width);
  ocr_budget_reserve(engine, reserved);   // budget.jai: waits while the GPU's memory is taken
  out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
  ocr_budget_release(engine, reserved);
  if !ok  return false, error;
  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 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);
  }

  for job, n: jobs {
    builder: String_Builder;
    builder.allocator = temp;
    previous := -1;
    probability_sum: float64;
    characters := 0;
    for t: 0..steps - 1 {
      row := out.data + (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]; }
      if best != 0 && best != previous {
        // Class 0 is the CTC blank, 1..count the dictionary, and the last one a space.
        if best <= engine.dictionary.count  append(*builder, engine.dictionary[best - 1]);
        else                               append(*builder, " ");
        probability_sum += best_p;
        characters += 1;
      }
      previous = best;
    }
    job.text = copy_string(trim(builder_to_string(*builder,, temp)),, context.default_allocator);
    job.confidence = ifx characters then cast(float) (probability_sum / characters) else 0;
  }
  return true, "";
}

// The width corners are recognized at: rec_height high, keeping their shape.
ocr_line_width :: (corners: [4] Vector2, rec_height: s32) -> s32 {
  w := length(corners[1] - corners[0]);
  h := max(length(corners[3] - corners[0]), 1);
  return clamp(cast(s32) (w * rec_height / h + 0.5), 1, REC_MAX_WIDTH);
}

REC_MAX_WIDTH :: 4096;   // a line longer than this, at rec_height, is squeezed to fit

// The normalization the PP-OCR recognizers were trained with.
REC_MEAN :: float.[0.5, 0.5, 0.5];
REC_STD  :: float.[0.5, 0.5, 0.5];

// Samples the rectangle the corners make (clockwise from its top left, at any angle) into a level
// crop width by height, bilinearly: crops are mostly enlarged (small UI text up to 48 px), where
// nearest would be blocky. Written planar in BGR order at dest, row_stride floats from one row to the
// next and plane from one channel to the next, as (value / 255 - mean) / deviation.
ocr_sample_crop :: (image: *Ocr_Image, corners: [4] Vector2, width: s32, height: s32, dest: *float, row_stride: s64, plane: s64, mean: [3] float, deviation: [3] float) {
  step_x := (corners[1] - corners[0]) / cast(float) width;    // one crop pixel along the line
  step_y := (corners[3] - corners[0]) / cast(float) height;   // and across it
  // Pixel centres: an image pixel's centre is at +0.5.
  origin := corners[0] + step_x * 0.5 + step_y * 0.5 - Vector2.{0.5, 0.5};
  for y: 0..height - 1 {
    row_start := origin + step_y * cast(float) y;
    for x: 0..width - 1 {
      s := row_start + step_x * cast(float) x;
      x0 := clamp(cast(s32) floor(s.x), 0, image.width - 1);
      y0 := clamp(cast(s32) floor(s.y), 0, image.height - 1);
      x1 := min(x0 + 1, image.width - 1);
      y1 := min(y0 + 1, image.height - 1);
      fx := clamp(s.x - x0, 0, 1);
      fy := clamp(s.y - y0, 0, 1);
      row0 := image.pixels + cast(s64) y0 * image.stride;
      row1 := image.pixels + cast(s64) y1 * image.stride;
      i := cast(s64) y * row_stride + x;
      for c: 0..2 {
        channel := 2 - c;   // B, G, R from RGBA
        top    := row0[x0 * 4 + channel] * (1 - fx) + row0[x1 * 4 + channel] * fx;
        bottom := row1[x0 * 4 + channel] * (1 - fx) + row1[x1 * 4 + channel] * fx;
        value  := top * (1 - fy) + bottom * fy;
        dest[i + plane * c] = (value / 255.0 - mean[c]) / deviation[c];
      }
    }
  }
}

// The orientation classifier (PP-LCNet text line orientation) on one batch: sets upside_down[n] when
// it finds jobs[n] more likely upside down than not. Its input is a fixed size, the line stretched to
// fit, normalized as the detector's input is.
ocr_classify_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job, upside_down: [] bool) -> ok: bool, error: string {
  if !jobs.count  return true, "";
  width, height := engine.orientation_width, engine.orientation_height;
  plane := cast(s64) width * height;
  input := NewArray(plane * 3 * jobs.count, float,, context.default_allocator);
  defer array_free(input,, context.default_allocator);
  for job, n: jobs  ocr_sample_crop(job.image, job.corners, width, height, input.data + plane * 3 * n, width, plane, DET_MEAN, DET_STD);

  shape := s64.[jobs.count, 3, height, width];
  out, ok, error := ocr_run(*engine.orientation, engine.memory_info, input, shape);
  if !ok  return false, error;
  defer ocr_release_output(*out);
  if out.dim_count != 2 || out.dims[0] != jobs.count || out.dims[1] != 2  return false, tprint("The orientation classifier gave an output of shape %, not [%, 2].", array_view(out.dims, 0, out.dim_count), jobs.count);
  for 0..jobs.count - 1  upside_down[it] = out.data[it * 2 + 1] > out.data[it * 2];   // classes: 0°, 180°
  return true, "";
}