381 lines
16 KiB
Plaintext
381 lines
16 KiB
Plaintext
// 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, and each kept region
|
|
// becomes the smallest rectangle around it, at whatever angle, 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.
|
|
|
|
Detected_Box :: struct {
|
|
// In the image's pixels, clockwise: the text's top left, top right, bottom right, bottom left, as
|
|
// far as its shape tells. A box at least TURN_RATIO times as tall as it is wide is taken as a line
|
|
// read downwards, so its first corner is its top right. Recognition may still turn it by 180°.
|
|
corners: [4] Vector2;
|
|
turned: bool; // taken as a line read downwards: started from its top right
|
|
score: float;
|
|
}
|
|
|
|
TURN_RATIO :: 1.5;
|
|
|
|
// Screenshots are mostly level: a region keeps a level box unless a turned one is this much smaller.
|
|
LEVEL_PREFERENCE :: 1.2; // a long line tilted by 1° already has a level box ~30% bigger
|
|
|
|
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);
|
|
#if OS == .MACOS if engine.det_coreml {
|
|
// Up to a size the Core ML package takes: its steps, which det_max_side and det_tile_height stay within.
|
|
w = align_up(valid_w, model.coreml_det_width_step);
|
|
h = align_up(valid_h, model.coreml_det_height_step);
|
|
}
|
|
|
|
// Planar BGR, normalized; the padding to a multiple of 32 (or the Core ML step) 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];
|
|
probabilities: *float;
|
|
dims: [] s64;
|
|
out: Ocr_Output;
|
|
defer ocr_release_output(*out);
|
|
#if OS == .MACOS {
|
|
coreml_out: Coreml_Output;
|
|
defer coreml_release_output(*coreml_out);
|
|
}
|
|
#if OS == .MACOS if engine.det_coreml {
|
|
ok: bool; error: string;
|
|
coreml_out, ok, error = coreml_predict(engine.det_coreml, input, shape);
|
|
if !ok return false, error;
|
|
probabilities, dims = coreml_out.data, array_view(coreml_out.dims, 0, coreml_out.dim_count);
|
|
}
|
|
if !probabilities {
|
|
ok: bool; error: string;
|
|
out, ok, error = ocr_run(*engine.det, engine.memory_info, input, shape);
|
|
if !ok return false, error;
|
|
probabilities, dims = out.data, array_view(out.dims, 0, out.dim_count);
|
|
}
|
|
if dims.count != 4 || dims[2] != h || dims[3] != w {
|
|
return false, tprint("The detector gave an output of shape %, not [1, 1, %, %].", dims, h, w);
|
|
}
|
|
|
|
regions := ocr_find_regions(probabilities, w, valid_w, valid_h, model.det_threshold);
|
|
defer {
|
|
for regions array_free(it.outline);
|
|
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;
|
|
|
|
rect := ocr_min_area_rect(it.outline);
|
|
rect_w := rect.u_max - rect.u_min;
|
|
rect_h := rect.v_max - rect.v_min;
|
|
if min(rect_w, rect_h) < 3 continue; // specks
|
|
|
|
// Grown as DB's unclip does: by area * ratio / perimeter on every side.
|
|
distance := (rect_w * rect_h * model.unclip_ratio) / (2.0 * (rect_w + rect_h));
|
|
rect.u_min -= distance; rect.u_max += distance;
|
|
rect.v_min -= distance; rect.v_max += distance;
|
|
corners := ocr_rect_corners(rect);
|
|
|
|
centre_y := tile_top + (corners[0].y + corners[2].y) / 2;
|
|
if centre_y < cast(float) band_top || centre_y >= cast(float) band_bottom continue; // the neighbouring strip's
|
|
|
|
box: Detected_Box;
|
|
box.score = score;
|
|
for corners {
|
|
box.corners[it_index].x = clamp(it.x * inverse, 0, cast(float) image.width);
|
|
box.corners[it_index].y = clamp((it.y + tile_top) * inverse, 0, cast(float) image.height);
|
|
}
|
|
if rect.u.x == 1 && rect.u.y == 0 {
|
|
// A level box takes in every pixel it touches: the recognizers read these crops best.
|
|
x0, y0, x1, y1 := ocr_corner_bounds(box.corners, image.width, image.height);
|
|
box.corners = .[.{xx x0, xx y0}, .{xx x1, xx y0}, .{xx x1, xx y1}, .{xx x0, xx y1}];
|
|
}
|
|
box.corners, box.turned = ocr_reading_order(box.corners);
|
|
x0, y0, x1, y1 := ocr_corner_bounds(box.corners, image.width, image.height);
|
|
if x1 - x0 >= 2 && y1 - y0 >= 2 array_add(boxes, box);
|
|
}
|
|
return true, "";
|
|
}
|
|
|
|
// The pixels the corners cover, clipped to the image: x0 and y0 inclusive, x1 and y1 exclusive.
|
|
ocr_corner_bounds :: (corners: [4] Vector2, image_width: s32, image_height: s32) -> x0: s32, y0: s32, x1: s32, y1: s32 {
|
|
low := corners[0];
|
|
high := corners[0];
|
|
for corners {
|
|
low.x = min(low.x, it.x); low.y = min(low.y, it.y);
|
|
high.x = max(high.x, it.x); high.y = max(high.y, it.y);
|
|
}
|
|
return clamp(cast(s32) floor(low.x), 0, image_width), clamp(cast(s32) floor(low.y), 0, image_height),
|
|
clamp(cast(s32) ceil(high.x), 0, image_width), clamp(cast(s32) ceil(high.y), 0, image_height);
|
|
}
|
|
|
|
// The same rectangle's corners started from another one, keeping them clockwise: by 1, turned a
|
|
// quarter (what was the right side becomes the top), by 2, upside down.
|
|
ocr_corners_turned :: (corners: [4] Vector2, quarters: s64) -> [4] Vector2 {
|
|
result: [4] Vector2;
|
|
for 0..3 result[it] = corners[(it + quarters) % 4];
|
|
return result;
|
|
}
|
|
|
|
// A rectangle as ranges along two perpendicular unit axes: u, and v = u turned a quarter clockwise
|
|
// on the screen (y down).
|
|
Rotated_Rect :: struct {
|
|
u: Vector2;
|
|
u_min, u_max, v_min, v_max: float;
|
|
}
|
|
|
|
ocr_rect_corners :: (rect: Rotated_Rect) -> [4] Vector2 {
|
|
u := rect.u;
|
|
v := Vector2.{-u.y, u.x};
|
|
return .[u * rect.u_min + v * rect.v_min, u * rect.u_max + v * rect.v_min,
|
|
u * rect.u_max + v * rect.v_max, u * rect.u_min + v * rect.v_max];
|
|
}
|
|
|
|
// The corners clockwise from the text's top left, as PaddleOCR orders them: of the two leftmost, the
|
|
// higher is the top left; of the two rightmost, the higher is the top right. A box at least
|
|
// TURN_RATIO times as tall as it is wide is then started from its top right, so a line read
|
|
// downwards comes out level.
|
|
ocr_reading_order :: (corners: [4] Vector2) -> [4] Vector2, turned: bool {
|
|
sorted := corners;
|
|
quick_sort(sorted, (a: Vector2, b: Vector2) -> s64 { return ifx a.x < b.x then -1 else ifx a.x > b.x then 1 else 0; });
|
|
top_left, bottom_left := sorted[0], sorted[1];
|
|
if bottom_left.y < top_left.y swap(*top_left, *bottom_left);
|
|
top_right, bottom_right := sorted[2], sorted[3];
|
|
if bottom_right.y < top_right.y swap(*top_right, *bottom_right);
|
|
result := Vector2.[top_left, top_right, bottom_right, bottom_left];
|
|
|
|
width := length(result[1] - result[0]);
|
|
height := length(result[3] - result[0]);
|
|
if height >= width * TURN_RATIO return ocr_corners_turned(result, 1), true;
|
|
return result, false;
|
|
}
|
|
|
|
// The smallest rectangle holding the points, at any angle, found by rotating calipers: one side of
|
|
// it lies along an edge of the points' convex hull. A level rectangle is kept unless a turned one is
|
|
// smaller by more than LEVEL_PREFERENCE.
|
|
ocr_min_area_rect :: (points: [] Vector2) -> Rotated_Rect {
|
|
level: Rotated_Rect;
|
|
level.u = .{1, 0};
|
|
level.u_min, level.v_min = FLOAT32_MAX, FLOAT32_MAX;
|
|
level.u_max, level.v_max = -FLOAT32_MAX, -FLOAT32_MAX;
|
|
for points {
|
|
level.u_min = min(level.u_min, it.x); level.u_max = max(level.u_max, it.x);
|
|
level.v_min = min(level.v_min, it.y); level.v_max = max(level.v_max, it.y);
|
|
}
|
|
if points.count < 3 return level;
|
|
level_area := (level.u_max - level.u_min) * (level.v_max - level.v_min);
|
|
|
|
hull := ocr_convex_hull(points);
|
|
best := level;
|
|
best_area := level_area;
|
|
for i: 0..hull.count - 1 {
|
|
edge := hull[(i + 1) % hull.count] - hull[i];
|
|
edge_length := length(edge);
|
|
if edge_length < 0.0001 continue;
|
|
rect: Rotated_Rect;
|
|
rect.u = edge / edge_length;
|
|
v := Vector2.{-rect.u.y, rect.u.x};
|
|
rect.u_min, rect.v_min = FLOAT32_MAX, FLOAT32_MAX;
|
|
rect.u_max, rect.v_max = -FLOAT32_MAX, -FLOAT32_MAX;
|
|
for hull {
|
|
a := dot(it, rect.u);
|
|
b := dot(it, v);
|
|
rect.u_min = min(rect.u_min, a); rect.u_max = max(rect.u_max, a);
|
|
rect.v_min = min(rect.v_min, b); rect.v_max = max(rect.v_max, b);
|
|
}
|
|
area := (rect.u_max - rect.u_min) * (rect.v_max - rect.v_min);
|
|
if area < best_area { best = rect; best_area = area; }
|
|
}
|
|
if level_area <= best_area * LEVEL_PREFERENCE return level;
|
|
return best;
|
|
}
|
|
|
|
// The convex hull of the points, counter-clockwise (Andrew's monotone chain). Temp.
|
|
ocr_convex_hull :: (points: [] Vector2) -> [] Vector2 {
|
|
sorted := array_copy(points,, temp);
|
|
quick_sort(sorted, (a: Vector2, b: Vector2) -> s64 {
|
|
if a.x != b.x return ifx a.x < b.x then -1 else 1;
|
|
return ifx a.y < b.y then -1 else ifx a.y > b.y then 1 else 0;
|
|
});
|
|
cross :: (o: Vector2, a: Vector2, b: Vector2) -> float { return (a.x - o.x) * (b.y - o.y) - (a.y - o.y) * (b.x - o.x); }
|
|
|
|
hull := NewArray(sorted.count * 2, Vector2,, temp);
|
|
k := 0;
|
|
for sorted { // lower hull
|
|
while k >= 2 && cross(hull[k - 2], hull[k - 1], it) <= 0 k -= 1;
|
|
hull[k] = it; k += 1;
|
|
}
|
|
lower := k + 1;
|
|
for < i: 0..sorted.count - 2 { // upper hull
|
|
p := sorted[i];
|
|
while k >= lower && cross(hull[k - 2], hull[k - 1], p) <= 0 k -= 1;
|
|
hull[k] = p; k += 1;
|
|
}
|
|
hull.count = max(k - 1, 1); // the last point is the first again
|
|
return hull;
|
|
}
|
|
|
|
align_up_32 :: (n: s32) -> s32 {
|
|
return (n + 31) & cast(s32) ~31;
|
|
}
|
|
|
|
// n rounded up to a multiple of step (at least one step).
|
|
align_up :: (n: s32, step: s32) -> s32 {
|
|
step = max(step, 1);
|
|
return max((n + step - 1) / step, 1) * step;
|
|
}
|
|
|
|
Region :: struct {
|
|
pixels: s64;
|
|
probability_sum: float64;
|
|
// The corners of the first and last pixel in each of its rows: what its convex hull needs.
|
|
outline: [..] Vector2; // default allocator
|
|
row: s32 = -1; // the row being scanned, and its first and last pixel
|
|
row_first, row_last: s32;
|
|
}
|
|
|
|
// 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 probability and outline.
|
|
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;
|
|
}
|
|
|
|
flush_row :: (r: *Region) {
|
|
if r.row < 0 return;
|
|
top, bottom := cast(float) r.row, cast(float) r.row + 1;
|
|
array_add(*r.outline, .{cast(float) r.row_first, top}, .{cast(float) r.row_first, bottom},
|
|
.{cast(float) r.row_last + 1, top}, .{cast(float) r.row_last + 1, bottom});
|
|
r.row = -1;
|
|
}
|
|
|
|
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.* = .{};
|
|
region.outline.allocator = context.default_allocator;
|
|
index = cast(s32) regions.count - 1;
|
|
region_of[root] = index + 1;
|
|
}
|
|
r := *regions[index];
|
|
if r.row != y {
|
|
flush_row(r);
|
|
r.row, r.row_first = y, x;
|
|
}
|
|
r.row_last = x; // rows are scanned left to right
|
|
r.pixels += 1;
|
|
r.probability_sum += probability[cast(s64) y * stride + x];
|
|
}
|
|
}
|
|
for * regions flush_row(it);
|
|
return regions;
|
|
}
|