use std::hint::black_box;
use std::path::{Path, PathBuf};
use image::GrayImage;
use ndarray::Array2;
use crate::config::{DetectionConfig, Language};
use crate::recognize::Charset;
use crate::types::Image;
const FALLBACK_IMAGE: &str = "tests/data/images/english.png";
const BLOB_SCORE: f32 = 0.9;
const LINK_BASELINE: f32 = 0.1;
const BLOB_BLOCK: usize = 6;
const BOX_STRIDE: f32 = 45.0;
const BOX_WIDTH: f32 = 40.0;
const BOX_TOP: f32 = 10.0;
const BOX_HEIGHT: f32 = 20.0;
const LOGIT_SLOPE: f32 = 0.5;
const BENCH_INV_RATIO: f32 = 2.0;
pub fn load_corpus_image(relative: &str) -> Image {
let manifest = Path::new(env!("CARGO_MANIFEST_DIR"));
if let Some(root) = repository_root_with_corpus(manifest) {
let candidate = root.join("test_documents").join(relative);
if let Ok(image) = Image::from_path(&candidate) {
return image;
}
}
let fallback = manifest.join(FALLBACK_IMAGE);
Image::from_path(&fallback)
.unwrap_or_else(|error| panic!("failed to decode fallback image {}: {error}", fallback.display()))
}
fn repository_root_with_corpus(start: &Path) -> Option<PathBuf> {
start
.ancestors()
.find(|dir| dir.join("Cargo.toml").is_file() && dir.join("test_documents").is_dir())
.map(Path::to_path_buf)
}
pub fn synthetic_heatmaps(height: usize, width: usize) -> (Array2<f32>, Array2<f32>) {
let mut region = Array2::<f32>::zeros((height, width));
let mut link = Array2::<f32>::zeros((height, width));
let step = BLOB_BLOCK * 2;
let mut y = BLOB_BLOCK;
while y + BLOB_BLOCK < height {
let mut x = BLOB_BLOCK;
while x + BLOB_BLOCK < width {
for row in y..(y + BLOB_BLOCK) {
for col in x..(x + BLOB_BLOCK) {
region[[row, col]] = BLOB_SCORE;
}
}
x += step;
}
y += step;
}
link.fill(LINK_BASELINE);
(region, link)
}
pub fn synthetic_boxes(count: usize) -> Vec<[[f32; 2]; 4]> {
(0..count)
.map(|index| {
let x0 = index as f32 * BOX_STRIDE;
let x1 = x0 + BOX_WIDTH;
let (y0, y1) = (BOX_TOP, BOX_TOP + BOX_HEIGHT);
[[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
})
.collect()
}
pub fn synthetic_logits(timesteps: usize, classes: usize) -> Array2<f32> {
let mut logits = Array2::<f32>::zeros((timesteps, classes));
let span = classes.max(1);
for t in 0..timesteps {
let peak = (t % span) as f32;
for c in 0..classes {
logits[[t, c]] = -(c as f32 - peak).abs() * LOGIT_SLOPE;
}
}
logits
}
pub fn english_class_count() -> usize {
english_charset().num_classes()
}
pub(crate) fn english_charset() -> Charset {
Charset::for_language(Language::English)
}
pub fn detect_preprocess_for_benchmark(image: &Image, canvas_size: u32, mag_ratio: f32) {
let tensor = crate::detect::bench_preprocess(image, canvas_size, mag_ratio).expect("corpus image preprocesses");
black_box(&tensor);
}
pub fn detect_preprocess_reference_for_benchmark(image: &Image, canvas_size: u32, mag_ratio: f32) {
let tensor =
crate::detect::bench_preprocess_reference(image, canvas_size, mag_ratio).expect("corpus image preprocesses");
black_box(&tensor);
}
pub fn detect_postprocess_for_benchmark(
region: &Array2<f32>,
link: &Array2<f32>,
text_threshold: f32,
link_threshold: f32,
low_text: f32,
) {
let boxes =
crate::detect::bench_postprocess(region, link, text_threshold, link_threshold, low_text, BENCH_INV_RATIO)
.expect("synthetic heat-maps decode to boxes");
black_box(&boxes);
}
pub fn detect_group_for_benchmark(boxes: &[[[f32; 2]; 4]], config: &DetectionConfig) {
let grouped = crate::detect::bench_group(boxes, config);
black_box(&grouped);
}
pub fn recognize_crop_preprocess_for_benchmark(gray: &GrayImage, corners: &[[f32; 2]; 4]) {
let tensor = crate::recognize::bench_crop_preprocess(gray, corners).expect("crop preprocesses to a batch");
black_box(&tensor);
}
pub fn recognize_crop_preprocess_reference_for_benchmark(gray: &GrayImage, corners: &[[f32; 2]; 4]) {
let tensor =
crate::recognize::bench_crop_preprocess_reference(gray, corners).expect("crop preprocesses to a batch");
black_box(&tensor);
}
pub fn ctc_decode_for_benchmark(logits: &Array2<f32>) {
let decoded = crate::recognize::bench_ctc_decode(logits.view());
black_box(&decoded);
}
pub fn ctc_decode_reference_for_benchmark(logits: &Array2<f32>) {
let decoded = crate::recognize::bench_ctc_decode_reference(logits.view());
black_box(&decoded);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn should_load_a_valid_fallback_image_when_corpus_is_absent() {
let image = load_corpus_image("images/does_not_exist_in_any_corpus.png");
assert!(image.width() > 0, "fallback image has non-zero width");
assert!(image.height() > 0, "fallback image has non-zero height");
}
#[test]
fn should_build_heatmaps_with_requested_shape_and_a_high_score_blob() {
let (region, link) = synthetic_heatmaps(40, 60);
assert_eq!(region.dim(), (40, 60));
assert_eq!(link.dim(), (40, 60));
let region_max = region.iter().copied().fold(f32::NEG_INFINITY, f32::max);
assert!(
(region_max - BLOB_SCORE).abs() < 1e-6,
"region contains a high-score blob"
);
}
#[test]
fn should_build_the_requested_number_of_boxes() {
let boxes = synthetic_boxes(7);
assert_eq!(boxes.len(), 7);
}
#[test]
fn should_build_logits_with_requested_shape() {
let logits = synthetic_logits(12, 97);
assert_eq!(logits.dim(), (12, 97));
}
#[test]
fn english_class_count_includes_the_blank() {
assert!(english_class_count() > 1, "charset has classes beyond the blank");
}
}