mcap/modules/OCR/detect.jai

// Detection: where the text lines are. The detector (a DB network) gives each pixel a probability
// of being text. Pixels above det_threshold are grouped into connected regions, each kept region
// becomes a box, and the box is grown by the DB unclip distance, since the network is trained to
// mark a shrunken core of each line.
//
// Images wider than det_max_side are scaled down for detection. Images taller than det_tile_height
// (after scaling) are detected in overlapping strips, so a long scrolling capture keeps its text at
// full size; each box is kept by the strip whose own band holds its centre.
//
// Version 1 finds horizontal text only, as screenshots have: boxes are axis-aligned.

Detected_Box :: struct {
  x0, y0, x1, y1: s32;   // in the image, x1 and y1 exclusive
  score: float;
}

DET_TILE_OVERLAP : s32 : 128;   // px of each strip shared with the next, after scaling

// The image's text boxes, from the default allocator.
ocr_detect :: (engine: *Ocr_Engine, image: Ocr_Image) -> [..] Detected_Box, ok: bool, error: string {
  boxes: [..] Detected_Box;
  boxes.allocator = context.default_allocator;
  model := *engine.model;
  if image.width <= 0 || image.height <= 0  return boxes, true, "";

  scale := min(1.0, cast(float) model.det_max_side / image.width);
  scaled_width  := max(cast(s32) 1, cast(s32) (image.width  * scale + 0.5));
  scaled_height := max(cast(s32) 1, cast(s32) (image.height * scale + 0.5));
  tile_height   := max(model.det_tile_height, cast(s32) DET_TILE_OVERLAP * 2);

  tile_top: s32 = 0;
  while tile_top < scaled_height {
    tile_bottom := min(tile_top + tile_height, scaled_height);
    // This strip's own band: what is not left to the strips above and below it.
    band_top    := ifx tile_top == 0 then 0 else tile_top + DET_TILE_OVERLAP / 2;
    band_bottom := ifx tile_bottom == scaled_height then scaled_height else tile_bottom - DET_TILE_OVERLAP / 2;

    ok, error := ocr_detect_tile(engine, image, scale, scaled_width, tile_top, tile_bottom, band_top, band_bottom, *boxes);
    if !ok  return boxes, false, error;
    if tile_bottom == scaled_height  break;
    tile_top = tile_bottom - DET_TILE_OVERLAP;
  }
  return boxes, true, "";
}

#scope_module

// Normalization the PP-OCR detectors were trained with: ImageNet mean and deviation, applied to the
// channels in BGR order, as PaddleOCR reads images with OpenCV.
DET_MEAN :: float.[0.485, 0.456, 0.406];
DET_STD  :: float.[0.229, 0.224, 0.225];

ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_width: s32, tile_top: s32, tile_bottom: s32, band_top: s32, band_bottom: s32, boxes: *[..] Detected_Box) -> ok: bool, error: string {
  model := *engine.model;
  valid_w := scaled_width;
  valid_h := tile_bottom - tile_top;
  w := align_up_32(valid_w);
  h := align_up_32(valid_h);

  // Planar BGR, normalized; the padding to a multiple of 32 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);
  inverse := 1.0 / scale;
  for y: 0..valid_h - 1 {
    source_y := min(cast(s32) ((y + tile_top + 0.5) * inverse), image.height - 1);
    row := image.pixels + cast(s64) source_y * image.stride;
    for x: 0..valid_w - 1 {
      source_x := min(cast(s32) ((x + 0.5) * inverse), image.width - 1);
      p := row + source_x * 4;
      i := cast(s64) y * w + x;
      input[i]             = (p[2] / 255.0 - DET_MEAN[0]) / DET_STD[0];   // B
      input[i + plane]     = (p[1] / 255.0 - DET_MEAN[1]) / DET_STD[1];   // G
      input[i + plane * 2] = (p[0] / 255.0 - DET_MEAN[2]) / DET_STD[2];   // R
    }
  }

  shape := s64.[1, 3, h, w];
  out, ok, error := ocr_run(*engine.det, engine.memory_info, input, shape);
  if !ok  return false, error;
  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);
  }

  regions := ocr_find_regions(out.data, w, valid_w, valid_h, model.det_threshold);
  defer array_free(regions,, context.default_allocator);
  for regions {
    if it.pixels == 0  continue;
    score := cast(float) (it.probability_sum / it.pixels);
    if score < model.box_threshold  continue;
    bw := it.x1 - it.x0;
    bh := it.y1 - it.y0;
    if min(bw, bh) < 3  continue;   // specks

    // Grown as DB's unclip does: by area * ratio / perimeter on every side.
    distance := (cast(float) bw * bh * model.unclip_ratio) / (2.0 * (bw + bh));
    gx0 := it.x0 - distance;
    gy0 := it.y0 - distance;
    gx1 := it.x1 + distance;
    gy1 := it.y1 + distance;

    centre_y := tile_top + (it.y0 + it.y1) / 2.0;
    if centre_y < cast(float) band_top || centre_y >= cast(float) band_bottom  continue;   // the neighbouring strip's

    box: Detected_Box;
    box.x0 = clamp(cast(s32) (gx0 * inverse), 0, image.width);
    box.x1 = clamp(cast(s32) (gx1 * inverse + 0.999), 0, image.width);
    box.y0 = clamp(cast(s32) ((gy0 + tile_top) * inverse), 0, image.height);
    box.y1 = clamp(cast(s32) ((gy1 + tile_top) * inverse + 0.999), 0, image.height);
    box.score = score;
    if box.x1 - box.x0 >= 2 && box.y1 - box.y0 >= 2  array_add(boxes, box);
  }
  return true, "";
}

align_up_32 :: (n: s32) -> s32 {
  return (n + 31) & cast(s32) ~31;
}

Region :: struct {
  x0, y0, x1, y1: s32;   // x1 and y1 exclusive
  pixels: s64;
  probability_sum: float64;
}

// Groups the pixels above threshold into 8-connected regions: one pass assigning provisional labels
// and recording which touch (union-find), then one summing each region's extent and probability.
ocr_find_regions :: (probability: *float, stride: s32, width: s32, height: s32, threshold: float) -> [..] Region {
  regions: [..] Region;
  regions.allocator = context.default_allocator;

  labels := NewArray(cast(s64) width * height, s32, initialized = true,, context.default_allocator);
  defer array_free(labels,, context.default_allocator);
  parent: [..] s32;
  parent.allocator = context.default_allocator;
  defer array_free(parent);
  array_add(*parent, 0);   // label 0 is background

  find :: (parent: [] s32, label: s32) -> s32 {
    root := label;
    while parent[root] != root  root = parent[root];
    while parent[label] != root { next := parent[label]; parent[label] = root; label = next; }   // path compression
    return root;
  }
  unite :: (parent: [] s32, a: s32, b: s32) {
    ra, rb := find(parent, a), find(parent, b);
    if ra < rb  parent[rb] = ra;
    else if rb < ra  parent[ra] = rb;
  }

  for y: 0..height - 1 {
    for x: 0..width - 1 {
      if probability[cast(s64) y * stride + x] <= threshold  continue;
      i := cast(s64) y * width + x;
      label: s32 = 0;
      // Neighbours already visited: left, and the three above.
      neighbours: [4] s32;
      neighbours[0] = ifx x > 0                       then labels[i - 1]         else 0;
      neighbours[1] = ifx y > 0 && x > 0              then labels[i - width - 1] else 0;
      neighbours[2] = ifx y > 0                       then labels[i - width]     else 0;
      neighbours[3] = ifx y > 0 && x < width - 1      then labels[i - width + 1] else 0;
      for neighbours {
        if !it  continue;
        if !label  label = it;
        else if it != label  unite(parent, label, it);
      }
      if !label {
        label = cast(s32) parent.count;
        array_add(*parent, label);
      }
      labels[i] = label;
    }
  }

  // Each root gets a region.
  region_of := NewArray(parent.count, s32, initialized = true,, context.default_allocator);
  defer array_free(region_of,, context.default_allocator);
  for y: 0..height - 1 {
    for x: 0..width - 1 {
      i := cast(s64) y * width + x;
      if !labels[i]  continue;
      root := find(parent, labels[i]);
      index := region_of[root] - 1;
      if index < 0 {
        region := array_add(*regions);
        region.* = .{x0 = x, y0 = y, x1 = x + 1, y1 = y + 1};
        index = cast(s32) regions.count - 1;
        region_of[root] = index + 1;
      }
      r := *regions[index];
      r.x0 = min(r.x0, x);  r.x1 = max(r.x1, x + 1);
      r.y0 = min(r.y0, y);  r.y1 = max(r.y1, y + 1);
      r.pixels += 1;
      r.probability_sum += probability[cast(s64) y * stride + x];
    }
  }
  return regions;
}