use anyhow::{Context, Result};
use image::{imageops::FilterType, GrayImage};
use ndarray::Array4;
use ort::session::{builder::GraphOptimizationLevel, Session};
use std::borrow::Cow;
use std::fmt;
use std::path::{Path, PathBuf};
use crate::model_manager::ModelManager;
use crate::segmenter::{tile_line, LineSegment, LineSegmenter, DEFAULT_DENSITY_THRESHOLD_RATIO};
use crate::utils::calculate_accuracy;
use crate::OcrResult;
const DEFAULT_CHARSET: &str = include_str!("charset.txt");
pub const EXPECTED_INPUT_HEIGHT: u32 = 160;
pub const DEFAULT_INPUT_WIDTH: u32 = 1024;
const POLARITY_CORNER_FRACTION: u32 = 10;
const POLARITY_CORNER_FLOOR: u32 = 3;
const DARK_BACKGROUND_MEDIAN: u8 = 128;
fn background_is_dark(image: &GrayImage) -> bool {
let (width, height) = image.dimensions();
if width == 0 || height == 0 {
return false;
}
let ch = (height / POLARITY_CORNER_FRACTION)
.max(POLARITY_CORNER_FLOOR)
.min(height);
let cw = (width / POLARITY_CORNER_FRACTION)
.max(POLARITY_CORNER_FLOOR)
.min(width);
let mut samples = Vec::with_capacity((4 * ch * cw) as usize);
for (ox, oy) in [
(0, 0),
(width - cw, 0),
(0, height - ch),
(width - cw, height - ch),
] {
for y in 0..ch {
for x in 0..cw {
samples.push(image.get_pixel(ox + x, oy + y)[0]);
}
}
}
samples.sort_unstable();
let n = samples.len();
let median = (samples[n / 2 - 1] as f64 + samples[n / 2] as f64) / 2.0;
median < DARK_BACKGROUND_MEDIAN as f64
}
pub fn normalize_polarity(image: &GrayImage) -> Cow<'_, GrayImage> {
if !background_is_dark(image) {
return Cow::Borrowed(image);
}
let mut inverted = image.clone();
for pixel in inverted.pixels_mut() {
pixel[0] = 255 - pixel[0];
}
Cow::Owned(inverted)
}
fn page_for_segmentation(image_path: &Path) -> Result<GrayImage> {
let page = image::open(image_path)
.with_context(|| format!("cannot open {}", image_path.display()))?
.to_luma8();
Ok(normalize_polarity(&page).into_owned())
}
fn segment_page(segmenter: &LineSegmenter, image_path: &Path) -> Result<Vec<LineSegment>> {
let page = page_for_segmentation(image_path)?;
segmenter.segment_image(&page)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelContractError(pub String);
impl fmt::Display for ModelContractError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "model contract violation: {}", self.0)
}
}
impl std::error::Error for ModelContractError {}
pub fn normalize_charset(charset: &str) -> &str {
charset
.trim_start_matches(['\n', '\r'])
.trim_end_matches(['\n', '\r'])
}
fn static_dim(shape: &[i64], axis: usize) -> Option<usize> {
match shape.get(axis) {
Some(&d) if d > 0 => Some(d as usize),
_ => None,
}
}
fn check_contract(
charset_len: usize,
model_classes: Option<usize>,
model_height: Option<usize>,
source: &str,
) -> Result<(), ModelContractError> {
if charset_len == 0 {
return Err(ModelContractError(
"no charset available; cannot decode model output".to_string(),
));
}
if let Some(classes) = model_classes {
let expected = charset_len + 1;
if classes != expected {
return Err(ModelContractError(format!(
"charset/model mismatch.\n \
charset: {charset_len} characters -> expects {expected} classes \
({charset_len} + CTC blank)\n \
model ({source}): {classes} classes\n\
Every index above the first divergence would decode to the wrong character."
)));
}
}
if let Some(height) = model_height {
if height != EXPECTED_INPUT_HEIGHT as usize {
return Err(ModelContractError(format!(
"input height mismatch: this binding preprocesses to height {EXPECTED_INPUT_HEIGHT} \
but {source} expects {height}"
)));
}
}
Ok(())
}
fn check_density_ratio(ratio: f32) -> Result<f32> {
if !ratio.is_finite() || ratio <= 0.0 {
anyhow::bail!(
"density_threshold_ratio must be finite and greater than 0, got {ratio}; \
at or below 0 every row clears the gap threshold and the page comes back \
as a single band"
);
}
Ok(ratio)
}
pub struct MonOcrBuilder {
model_path: Option<PathBuf>,
charset: Option<String>,
min_line_height: u32,
smooth_window: u32,
density_threshold_ratio: f32,
tile_wide_lines: bool,
}
impl Default for MonOcrBuilder {
fn default() -> Self {
Self {
model_path: None,
charset: None,
min_line_height: 10,
smooth_window: 3,
density_threshold_ratio: DEFAULT_DENSITY_THRESHOLD_RATIO,
tile_wide_lines: true,
}
}
}
impl MonOcrBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn model_path(mut self, path: impl AsRef<Path>) -> Self {
self.model_path = Some(path.as_ref().to_path_buf());
self
}
pub fn charset(mut self, charset: impl Into<String>) -> Self {
self.charset = Some(charset.into());
self
}
pub fn min_line_height(mut self, height: u32) -> Self {
self.min_line_height = height;
self
}
pub fn smooth_window(mut self, window: u32) -> Self {
self.smooth_window = window;
self
}
pub fn tile_wide_lines(mut self, tile: bool) -> Self {
self.tile_wide_lines = tile;
self
}
pub fn density_threshold_ratio(mut self, ratio: f32) -> Self {
self.density_threshold_ratio = ratio;
self
}
pub async fn build(self) -> Result<MonOcr> {
MonOcr::new(
self.model_path,
self.charset,
self.min_line_height,
self.smooth_window,
check_density_ratio(self.density_threshold_ratio)?,
self.tile_wide_lines,
)
.await
}
}
pub struct MonOcr {
session: Session,
charset: Vec<char>,
segmenter: LineSegmenter,
target_height: u32,
target_width: u32,
tile_wide_lines: bool,
}
#[derive(Debug, Clone)]
pub struct LineResult {
pub text: String,
pub bbox: BBox,
}
#[derive(Debug, Clone, Copy)]
pub struct BBox {
pub x: u32,
pub y: u32,
pub w: u32,
pub h: u32,
}
fn union_bbox(a: BBox, b: BBox) -> BBox {
let x = a.x.min(b.x);
let y = a.y.min(b.y);
let right = (a.x + a.w).max(b.x + b.w);
let bottom = (a.y + a.h).max(b.y + b.h);
BBox {
x,
y,
w: right - x,
h: bottom - y,
}
}
pub fn page_text(lines: &[LineResult]) -> String {
lines
.iter()
.map(|l| l.text.as_str())
.collect::<Vec<_>>()
.join("\n")
}
impl MonOcr {
pub fn builder() -> MonOcrBuilder {
MonOcrBuilder::new()
}
async fn new(
model_path: Option<PathBuf>,
charset: Option<String>,
min_line_height: u32,
smooth_window: u32,
density_threshold_ratio: f32,
tile_wide_lines: bool,
) -> Result<Self> {
let (model_path, published_charset) = match model_path {
Some(path) => (path, None),
None => {
tokio::task::spawn_blocking(|| {
let manager = ModelManager::new();
let path = manager.get_model_path()?;
let published = manager.get_charset().ok();
Ok::<_, anyhow::Error>((path, published))
})
.await
.context("the model download task did not finish")??
}
};
let charset_str = charset
.or(published_charset)
.unwrap_or_else(|| DEFAULT_CHARSET.to_string());
let charset: Vec<char> = normalize_charset(&charset_str).chars().collect();
let session = Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.commit_from_file(&model_path)?;
let source = model_path.display().to_string();
let in_shape = session
.inputs()
.first()
.and_then(|i| i.dtype().tensor_shape())
.ok_or_else(|| ModelContractError(format!("{source} has no tensor input")))?
.to_vec();
let out_shape = session
.outputs()
.first()
.and_then(|o| o.dtype().tensor_shape())
.ok_or_else(|| ModelContractError(format!("{source} has no tensor output")))?
.to_vec();
if in_shape.len() != 4 {
return Err(ModelContractError(format!(
"expected a 4-D [batch, channel, height, width] input, {source} declares {in_shape:?}"
))
.into());
}
if out_shape.len() != 3 {
return Err(ModelContractError(format!(
"expected a 3-D [batch, sequence, classes] output, {source} declares {out_shape:?}"
))
.into());
}
let model_height = static_dim(&in_shape, 2);
let model_classes = static_dim(&out_shape, 2);
check_contract(charset.len(), model_classes, model_height, &source)?;
let segmenter = LineSegmenter::with_density_ratio(
min_line_height,
smooth_window,
density_threshold_ratio,
);
Ok(Self {
session,
charset,
segmenter,
target_height: model_height
.map(|h| h as u32)
.unwrap_or(EXPECTED_INPUT_HEIGHT),
target_width: DEFAULT_INPUT_WIDTH,
tile_wide_lines,
})
}
pub async fn read_image(&mut self, image_path: impl AsRef<Path>) -> Result<String> {
let results = self.predict_page(image_path).await?;
Ok(page_text(&results))
}
pub async fn read_images(&mut self, image_paths: &[impl AsRef<Path>]) -> Result<Vec<String>> {
let mut results = Vec::new();
for path in image_paths {
let text = self.read_image(path).await?;
results.push(text);
}
Ok(results)
}
pub async fn read_pdf(&mut self, pdf_path: impl AsRef<Path>) -> Result<Vec<String>> {
let pages = self.predict_pdf(pdf_path).await?;
Ok(pages.iter().map(|lines| page_text(lines)).collect())
}
pub async fn predict_pdf(
&mut self,
pdf_path: impl AsRef<Path>,
) -> Result<Vec<Vec<LineResult>>> {
use std::process::Stdio;
use tokio::process::Command;
let pdf_path = pdf_path.as_ref();
let check = Command::new("which").arg("pdftoppm").output().await;
if check.is_err() || !check.as_ref().map(|o| o.status.success()).unwrap_or(false) {
anyhow::bail!("pdftoppm not found: please install poppler-utils");
}
if check.as_ref().map(|o| o.stdout.is_empty()).unwrap_or(true) {
anyhow::bail!("pdftoppm not found: please install poppler-utils");
}
let temp_dir = tempfile::tempdir()?;
let output_prefix = temp_dir.path().join("page");
let output = Command::new("pdftoppm")
.args(["-png", "-r", "300"])
.arg(pdf_path)
.arg(&output_prefix)
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await?;
if !output.success() {
anyhow::bail!("Failed to convert PDF to images");
}
let mut entries: Vec<_> = std::fs::read_dir(temp_dir.path())?
.filter_map(|e| e.ok())
.filter(|e| {
e.path()
.extension()
.map(|ext| ext == "png")
.unwrap_or(false)
})
.collect();
entries.sort_by(|a, b| {
let name_a = a.file_name();
let name_b = b.file_name();
let num_a: u32 = name_a
.to_string_lossy()
.split('-')
.next_back()
.and_then(|s| s.trim_end_matches(".png").parse().ok())
.unwrap_or(0);
let num_b: u32 = name_b
.to_string_lossy()
.split('-')
.next_back()
.and_then(|s| s.trim_end_matches(".png").parse().ok())
.unwrap_or(0);
num_a.cmp(&num_b)
});
if entries.is_empty() {
anyhow::bail!("No images generated from PDF");
}
let mut pages = Vec::new();
for entry in entries {
let lines = self.predict_page(entry.path()).await?;
pages.push(lines);
}
Ok(pages)
}
pub async fn read_image_with_accuracy(
&mut self,
image_path: impl AsRef<Path>,
ground_truth: &str,
) -> Result<OcrResult> {
let text = self.read_image(image_path).await?;
let accuracy = calculate_accuracy(&text, ground_truth);
Ok(OcrResult { text, accuracy })
}
async fn predict_line(&mut self, image: &GrayImage) -> Result<String> {
let input_tensor = self.preprocess(image)?;
let input = ort::value::Tensor::from_array(input_tensor)?;
let outputs = self.session.run(ort::inputs![input])?;
let output = outputs[0].downcast_ref::<ort::value::DynTensorValueType>()?;
let (shape, data) = output.try_extract_tensor::<f32>()?;
let output_shape: Vec<usize> = shape.iter().cloned().map(|x| x as usize).collect();
let output_data: Vec<f32> = data.to_vec();
drop(outputs);
self.decode_owned(&output_data, &output_shape)
}
pub async fn predict_page(&mut self, image_path: impl AsRef<Path>) -> Result<Vec<LineResult>> {
let image_path = image_path.as_ref();
let lines = segment_page(&self.segmenter, image_path)?;
let mut results = Vec::new();
for line in lines {
let origin = BBox {
x: line.bbox.x,
y: line.bbox.y,
w: line.bbox.w,
h: line.bbox.h,
};
results.push(self.read_line_crop(&line.img, origin).await?);
}
Ok(results)
}
pub async fn predict_single_line(
&mut self,
image_path: impl AsRef<Path>,
) -> Result<LineResult> {
let image_path = image_path.as_ref();
let crop = image::open(image_path)
.with_context(|| format!("cannot open {}", image_path.display()))?
.to_luma8();
let (w, h) = crop.dimensions();
if w == 0 || h == 0 {
anyhow::bail!(
"{} is {w}x{h}: there is nothing to read",
image_path.display()
);
}
self.read_line_crop(&crop, BBox { x: 0, y: 0, w, h }).await
}
async fn read_line_crop(&mut self, crop: &GrayImage, origin: BBox) -> Result<LineResult> {
let tiles = if self.tile_wide_lines {
tile_line(crop, self.target_height, self.target_width)
} else {
vec![crop.clone()]
};
let mut text = String::new();
let mut bbox: Option<BBox> = None;
let mut x_offset = 0u32;
for tile in &tiles {
let (tile_w, tile_h) = tile.dimensions();
text.push_str(&self.predict_line(tile).await?);
let tile_bbox = BBox {
x: origin.x + x_offset,
y: origin.y,
w: tile_w,
h: tile_h,
};
bbox = Some(match bbox {
Some(current) => union_bbox(current, tile_bbox),
None => tile_bbox,
});
x_offset += tile_w;
}
Ok(LineResult {
text,
bbox: bbox.unwrap_or(origin),
})
}
fn preprocess(&self, image: &GrayImage) -> Result<Array4<f32>> {
Ok(preprocess_line(
image,
self.target_height,
self.target_width,
))
}
fn decode_owned(&self, data: &[f32], shape: &[usize]) -> Result<String> {
decode_ctc(&self.charset, data, shape)
}
}
fn preprocess_line(image: &GrayImage, target_height: u32, target_width: u32) -> Array4<f32> {
let normalized = normalize_polarity(image);
let image = normalized.as_ref();
let (width, height) = image.dimensions();
let scale = target_height as f32 / height as f32;
let new_width = (width as f32 * scale).round() as u32;
let new_width = new_width.min(target_width);
let resized = image::imageops::resize(image, new_width, target_height, FilterType::Triangle);
let mut tensor = Array4::<f32>::zeros((1, 1, target_height as usize, target_width as usize));
for y in 0..target_height {
for x in 0..target_width {
let value = if x < new_width {
let pixel = resized.get_pixel(x, y);
(pixel[0] as f32 / 127.5) - 1.0 } else {
1.0 };
tensor[[0, 0, y as usize, x as usize]] = value;
}
}
tensor
}
fn decode_ctc(charset: &[char], data: &[f32], shape: &[usize]) -> Result<String> {
if shape.len() != 3 {
return Err(ModelContractError(format!(
"expected a 3-D [batch, sequence, classes] output tensor, got shape {shape:?}"
))
.into());
}
let sequence_length = shape[1];
let num_classes = shape[2];
if sequence_length == 0 || num_classes == 0 {
return Err(ModelContractError(format!(
"output tensor has an empty axis: shape {shape:?}"
))
.into());
}
let expected = charset.len() + 1;
if num_classes != expected {
return Err(ModelContractError(format!(
"charset/model mismatch at decode time: charset has {} characters -> \
expects {expected} classes, tensor has {num_classes}",
charset.len()
))
.into());
}
if data.len() < sequence_length * num_classes {
return Err(ModelContractError(format!(
"output tensor holds {} values, shape {shape:?} needs {}",
data.len(),
sequence_length * num_classes
))
.into());
}
let mut decoded = String::new();
let mut prev_idx: i32 = -1;
for t in 0..sequence_length {
let mut max_val = f32::NEG_INFINITY;
let mut max_idx = 0;
let base = t * num_classes;
for c in 0..num_classes {
let val = data[base + c];
if val > max_val {
max_val = val;
max_idx = c;
}
}
if max_idx != 0 && max_idx as i32 != prev_idx {
decoded.push(charset[max_idx - 1]);
}
prev_idx = max_idx as i32;
}
Ok(decoded)
}
#[cfg(test)]
mod tests {
use super::*;
use image::Luma;
fn drawn_page(width: u32, height: u32, band_h: u32, glyph_w: u32, pitch: u32) -> GrayImage {
let mut img = GrayImage::from_pixel(width, height, Luma([255u8]));
let margin = height / 10 + 4;
let mut y = margin;
while y + band_h < height - margin {
for yy in y..y + band_h {
let mut x = margin;
while x + glyph_w < width - margin {
for i in 0..glyph_w {
img.put_pixel(x + i, yy, Luma([0u8]));
}
x += pitch;
}
}
y += band_h * 2;
}
img
}
fn inverted(img: &GrayImage) -> GrayImage {
let mut out = img.clone();
for pixel in out.pixels_mut() {
pixel[0] = 255 - pixel[0];
}
out
}
fn ink_fraction(img: &GrayImage) -> f64 {
let dark = img.pixels().filter(|p| p[0] < 128).count();
dark as f64 / (img.width() * img.height()) as f64
}
#[test]
fn a_light_page_is_not_inverted() {
let page = drawn_page(400, 300, 20, 10, 18);
let out = normalize_polarity(&page);
assert!(
matches!(out, Cow::Borrowed(_)),
"a light page must be handed back borrowed, not copied"
);
assert_eq!(out.as_raw(), page.as_raw(), "a light page was modified");
}
#[test]
fn a_dark_page_is_inverted() {
let page = drawn_page(400, 300, 20, 10, 18);
let dark = inverted(&page);
let out = normalize_polarity(&dark);
assert!(
matches!(out, Cow::Owned(_)),
"a light-on-dark page must be inverted"
);
assert_eq!(
out.as_raw(),
page.as_raw(),
"inverting a dark page must reproduce the light original exactly"
);
}
#[test]
fn a_dense_page_is_not_mistaken_for_dark_mode() {
let mut page = GrayImage::from_pixel(400, 300, Luma([255u8]));
let (margin_x, margin_y) = (34, 34);
for y in margin_y..300 - margin_y {
for x in margin_x..400 - margin_x {
page.put_pixel(x, y, Luma([0u8]));
}
}
let covered = ink_fraction(&page);
assert!(
covered > 0.5,
"the fixture must be more than half ink for this to test anything, \
got {covered:.3}"
);
assert!(
matches!(normalize_polarity(&page), Cow::Borrowed(_)),
"a dense but correctly-polarised page was inverted"
);
}
#[test]
fn polarity_is_idempotent() {
let dark = inverted(&drawn_page(400, 300, 20, 10, 18));
let once = normalize_polarity(&dark).into_owned();
let twice = normalize_polarity(&once).into_owned();
assert_eq!(
once.as_raw(),
twice.as_raw(),
"a second pass changed the image; the two call sites would fight"
);
}
#[test]
fn a_tiny_crop_does_not_panic() {
for (w, h) in [(1u32, 1u32), (1, 40), (40, 1), (5, 5), (2, 7)] {
let dark = GrayImage::from_pixel(w, h, Luma([10u8]));
assert!(
matches!(normalize_polarity(&dark), Cow::Owned(_)),
"{w}x{h}: a solid dark crop must be inverted"
);
let light = GrayImage::from_pixel(w, h, Luma([240u8]));
assert!(
matches!(normalize_polarity(&light), Cow::Borrowed(_)),
"{w}x{h}: a solid light crop must be left alone"
);
}
}
#[test]
fn a_dark_mode_page_segments_into_the_same_lines_as_its_upright_twin() {
let page = drawn_page(400, 300, 20, 10, 18);
let dir = tempfile::tempdir().unwrap();
let light_path = dir.path().join("light.png");
let dark_path = dir.path().join("dark.png");
page.save(&light_path).unwrap();
inverted(&page).save(&dark_path).unwrap();
let seg = LineSegmenter::new(10, 3);
let light = seg.segment(&light_path).unwrap();
assert!(
light.len() > 1,
"the upright control found {} line(s); the comparison below needs a \
page that actually segments",
light.len()
);
let uncorrected = seg.segment(&dark_path).unwrap();
assert_ne!(
uncorrected.len(),
light.len(),
"the dark page segmented correctly without the probe, so this test \
cannot show the probe is needed — pick a harder fixture"
);
let corrected = segment_page(&seg, &dark_path).unwrap();
assert_eq!(
corrected.len(),
light.len(),
"a dark-mode page gave {} line(s) against {} for the same page \
upright; polarity is not being corrected before segmentation",
corrected.len(),
light.len()
);
for (c, l) in corrected.iter().zip(light.iter()) {
assert_eq!(
(c.bbox.y, c.bbox.h),
(l.bbox.y, l.bbox.h),
"corrected bands must land on the same rows as the upright page"
);
}
}
#[test]
fn a_dark_crop_preprocesses_to_the_same_tensor_as_its_upright_twin() {
let crop = drawn_page(300, 40, 20, 10, 18);
let upright = preprocess_line(&crop, EXPECTED_INPUT_HEIGHT, DEFAULT_INPUT_WIDTH);
let dark = preprocess_line(&inverted(&crop), EXPECTED_INPUT_HEIGHT, DEFAULT_INPUT_WIDTH);
assert!(
upright.iter().any(|v| *v < 0.0),
"the upright control has no ink in it, so equality proves nothing"
);
assert_eq!(
upright, dark,
"a light-on-dark crop reached the model as a different tensor from \
the same crop upright; the per-crop polarity probe is missing"
);
}
const PINNED_CLASSES: usize = 277;
const PINNED_CHAR_LEN: usize = 276;
fn charset_of_len(n: usize) -> Vec<char> {
std::iter::once(' ')
.chain(std::iter::repeat_n('x', n - 1))
.collect()
}
#[test]
fn contract_accepts_the_pinned_model() {
check_contract(
PINNED_CHAR_LEN,
Some(PINNED_CLASSES),
Some(EXPECTED_INPUT_HEIGHT as usize),
"model.onnx",
)
.expect("the pinned pair should pass");
}
#[test]
fn contract_rejects_charset_mismatch() {
let err = check_contract(
225,
Some(PINNED_CLASSES),
Some(EXPECTED_INPUT_HEIGHT as usize),
"model.onnx",
)
.expect_err("225 characters vs 277 classes must be refused");
assert!(err.0.contains("226") && err.0.contains("277"), "{err}");
}
#[test]
fn contract_rejects_off_by_one_charset() {
check_contract(
PINNED_CHAR_LEN - 1,
Some(PINNED_CLASSES),
Some(EXPECTED_INPUT_HEIGHT as usize),
"model.onnx",
)
.expect_err("275 characters vs 277 classes must be refused");
}
#[test]
fn contract_rejects_height_mismatch() {
check_contract(
PINNED_CHAR_LEN,
Some(PINNED_CLASSES),
Some(64),
"stale.onnx",
)
.expect_err("a 64-pixel input must be refused");
}
#[test]
fn contract_rejects_empty_charset() {
check_contract(
0,
Some(PINNED_CLASSES),
Some(EXPECTED_INPUT_HEIGHT as usize),
"model.onnx",
)
.expect_err("an empty charset must be refused");
}
#[test]
fn contract_skips_dynamic_axes() {
check_contract(PINNED_CHAR_LEN, None, None, "dynamic.onnx")
.expect("dynamic axes should defer the check");
}
#[test]
fn static_dim_reads_only_fixed_axes() {
let shape = [1i64, 1, 128, -1];
assert_eq!(static_dim(&shape, 2), Some(128));
assert_eq!(static_dim(&shape, 3), None, "dynamic axis");
assert_eq!(static_dim(&shape, 9), None, "out of range");
}
fn synthetic_logits(seq_len: usize, num_classes: usize) -> Vec<f32> {
(0..seq_len * num_classes)
.map(|i| (i as f32 * 0.37).sin())
.collect()
}
#[test]
fn decode_stride_comes_from_the_tensor() {
let seq_len = 128;
let data = synthetic_logits(seq_len, PINNED_CLASSES);
let shape = [1, seq_len, PINNED_CLASSES];
let text = decode_ctc(&charset_of_len(PINNED_CHAR_LEN), &data, &shape)
.expect("the matching charset should decode");
assert!(!text.is_empty());
decode_ctc(&charset_of_len(PINNED_CHAR_LEN - 1), &data, &shape)
.expect_err("a 275-character charset against a 277-class tensor must be refused");
decode_ctc(&charset_of_len(225), &data, &shape)
.expect_err("a 225-character charset against a 277-class tensor must be refused");
}
#[test]
fn decode_rejects_unexpected_shapes() {
let charset = charset_of_len(PINNED_CHAR_LEN);
let data = synthetic_logits(8, PINNED_CLASSES);
decode_ctc(&charset, &data, &[8, PINNED_CLASSES]).expect_err("2-D output");
decode_ctc(&charset, &data, &[1, 0, PINNED_CLASSES]).expect_err("empty sequence axis");
decode_ctc(&charset, &data, &[1, 16, PINNED_CLASSES])
.expect_err("shape larger than the buffer");
}
#[test]
fn union_of_adjacent_tiles_is_the_line() {
let line = BBox {
x: 100,
y: 40,
w: 900,
h: 60,
};
let widths = [254u32, 255, 255, 136];
let mut x = line.x;
let mut acc: Option<BBox> = None;
for w in widths {
let tile = BBox {
x,
y: line.y,
w,
h: line.h,
};
acc = Some(match acc {
Some(current) => union_bbox(current, tile),
None => tile,
});
x += w;
}
let got = acc.expect("at least one tile");
assert_eq!(
(got.x, got.y, got.w, got.h),
(line.x, line.y, line.w, line.h)
);
}
#[test]
fn union_covers_boxes_in_any_order() {
let a = BBox {
x: 10,
y: 5,
w: 4,
h: 2,
};
let b = BBox {
x: 2,
y: 9,
w: 3,
h: 6,
};
let u = union_bbox(a, b);
assert_eq!((u.x, u.y, u.w, u.h), (2, 5, 12, 10));
let flipped = union_bbox(b, a);
assert_eq!(
(flipped.x, flipped.y, flipped.w, flipped.h),
(u.x, u.y, u.w, u.h)
);
}
#[test]
fn density_ratio_default_is_unchanged() {
assert_eq!(DEFAULT_DENSITY_THRESHOLD_RATIO, 0.05);
assert_eq!(
MonOcrBuilder::default().density_threshold_ratio,
DEFAULT_DENSITY_THRESHOLD_RATIO
);
assert_eq!(
MonOcr::builder()
.density_threshold_ratio(0.3)
.density_threshold_ratio,
0.3
);
}
#[test]
fn density_ratio_rejects_useless_values() {
for bad in [0.0, -0.05, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
check_density_ratio(bad).expect_err(&format!("{bad} must be refused"));
}
for good in [DEFAULT_DENSITY_THRESHOLD_RATIO, 0.12, 0.5, 1.0] {
assert_eq!(check_density_ratio(good).expect("valid ratio"), good);
}
}
#[test]
fn decode_ctc_semantics() {
let charset: Vec<char> = "abc".chars().collect();
let num_classes = charset.len() + 1;
let argmax = [1usize, 1, 0, 1, 2, 3];
let mut data = vec![0.0f32; argmax.len() * num_classes];
for (t, &want) in argmax.iter().enumerate() {
data[t * num_classes + want] = 1.0;
}
let got = decode_ctc(&charset, &data, &[1, argmax.len(), num_classes]).unwrap();
assert_eq!(got, "aabc");
}
}