106 lines
4.4 KiB
Plaintext
106 lines
4.4 KiB
Plaintext
// Recognition: the text of each detected line. Each box is cropped from the image at full size,
|
|
// scaled to rec_height 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.
|
|
|
|
Line_Job :: struct {
|
|
image: *Ocr_Image;
|
|
box: Detected_Box;
|
|
width: s32; // after scaling to rec_height
|
|
// 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
|
|
|
|
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 {
|
|
base := input.data + plane * 3 * n;
|
|
box := job.box;
|
|
bw := box.x1 - box.x0;
|
|
bh := box.y1 - box.y0;
|
|
image := job.image;
|
|
// Bilinear: the crops are mostly enlarged (small UI text up to 48 px), where nearest would be blocky.
|
|
x_scale := cast(float) bw / job.width;
|
|
y_scale := cast(float) bh / height;
|
|
for y: 0..height - 1 {
|
|
sy := (y + 0.5) * y_scale - 0.5 + box.y0;
|
|
y0 := clamp(cast(s32) floor(sy), 0, image.height - 1);
|
|
y1 := min(y0 + 1, image.height - 1);
|
|
fy := clamp(sy - y0, 0, 1);
|
|
row0 := image.pixels + cast(s64) y0 * image.stride;
|
|
row1 := image.pixels + cast(s64) y1 * image.stride;
|
|
for x: 0..job.width - 1 {
|
|
sx := (x + 0.5) * x_scale - 0.5 + box.x0;
|
|
x0 := clamp(cast(s32) floor(sx), 0, image.width - 1);
|
|
x1 := min(x0 + 1, image.width - 1);
|
|
fx := clamp(sx - x0, 0, 1);
|
|
i := cast(s64) y * width + 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;
|
|
base[i + plane * c] = (value / 255.0 - 0.5) / 0.5;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
shape := s64.[jobs.count, 3, height, width];
|
|
out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
|
|
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 a box is recognized at: rec_height high, keeping its shape.
|
|
ocr_line_width :: (box: Detected_Box, rec_height: s32) -> s32 {
|
|
bw := box.x1 - box.x0;
|
|
bh := max(box.y1 - box.y0, 1);
|
|
return clamp(cast(s32) (cast(float) bw * rec_height / bh + 0.5), 1, REC_MAX_WIDTH);
|
|
}
|
|
|
|
REC_MAX_WIDTH :: 4096; // a line longer than this, at rec_height, is squeezed to fit
|