visual-cortex-ocr-onnx 0.1.0

use image::RgbImage;
use ndarray::Array4;

/// Rec model input height (PP-OCRv3) and maximum width.
pub(crate) const REC_HEIGHT: u32 = 48;
/// Crops wider than 320px at height 48 are squashed (aspect not preserved
/// past the cap), not sliding-windowed — very long lines may misrecognize.
pub(crate) const REC_MAX_WIDTH: u32 = 320;

/// Build the CTC charset from the dictionary file contents:
/// index 0 = blank, then one entry per line, then space (use_space_char).
pub(crate) fn build_charset(dict: &str) -> Vec<String> {
    let mut charset = vec![String::new()]; // blank
    charset.extend(dict.lines().map(|l| l.to_string()));
    charset.push(" ".to_string());
    charset
}

/// Resize a crop to height 48 (aspect preserved, width capped at 320) and
/// normalize into the rec model's zero-padded NCHW input tensor.
/// Degenerate (zero-height) crops produce a blank tensor rather than panicking.
pub(crate) fn preprocess_rec(crop: &RgbImage) -> Array4<f32> {
    let scale = REC_HEIGHT as f32 / crop.height() as f32;
    let w = ((crop.width() as f32 * scale).round() as u32).clamp(1, REC_MAX_WIDTH);
    let resized =
        image::imageops::resize(crop, w, REC_HEIGHT, image::imageops::FilterType::Triangle);
    let mut input = Array4::<f32>::zeros((1, 3, REC_HEIGHT as usize, REC_MAX_WIDTH as usize));
    for (x, y, pixel) in resized.enumerate_pixels() {
        for c in 0..3 {
            input[[0, c, y as usize, x as usize]] = (pixel.0[c] as f32 / 255.0 - 0.5) / 0.5;
        }
    }
    input
}

/// Greedy CTC decode over row-major `[steps, classes]` probabilities
/// (post-softmax). Collapses repeats, drops blanks (class 0). Returns the
/// text and the mean probability of kept characters (0.0 when none kept —
/// never NaN).
pub(crate) fn ctc_decode(
    probs: &[f32],
    steps: usize,
    classes: usize,
    charset: &[String],
) -> (String, f32) {
    let mut text = String::new();
    let mut conf_sum = 0.0f32;
    let mut kept = 0u32;
    let mut prev_class = 0usize;
    for t in 0..steps {
        let row = &probs[t * classes..(t + 1) * classes];
        let (best, &best_p) = row
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .expect("classes > 0");
        if best != 0 && best != prev_class {
            if let Some(ch) = charset.get(best) {
                text.push_str(ch);
                conf_sum += best_p;
                kept += 1;
            }
        }
        prev_class = best;
    }
    let confidence = if kept == 0 {
        0.0
    } else {
        conf_sum / kept as f32
    };
    (text, confidence)
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A 4-class charset: blank + "a" + "b" + space.
    fn tiny_charset() -> Vec<String> {
        build_charset("a\nb")
    }

    /// One-hot probability rows for a timestep sequence.
    fn probs(steps: &[usize], classes: usize) -> Vec<f32> {
        let mut out = vec![0.0f32; steps.len() * classes];
        for (t, &c) in steps.iter().enumerate() {
            out[t * classes + c] = 0.8;
        }
        out
    }

    #[test]
    fn build_charset_adds_blank_and_space() {
        let cs = tiny_charset();
        assert_eq!(cs.len(), 4); // blank + a + b + space
        assert_eq!(cs[1], "a");
        assert_eq!(cs[3], " ");
    }

    #[test]
    fn ctc_collapses_repeats_and_drops_blanks() {
        let cs = tiny_charset();
        // a a blank a b b -> "aab"
        let p = probs(&[1, 1, 0, 1, 2, 2], cs.len());
        let (text, conf) = ctc_decode(&p, 6, cs.len(), &cs);
        assert_eq!(text, "aab");
        assert!((conf - 0.8).abs() < 1e-6);
    }

    #[test]
    fn ctc_decodes_space_class() {
        let cs = tiny_charset();
        // a space b -> "a b"
        let p = probs(&[1, 3, 2], cs.len());
        let (text, _) = ctc_decode(&p, 3, cs.len(), &cs);
        assert_eq!(text, "a b");
    }

    #[test]
    fn ctc_all_blank_is_empty_with_zero_confidence() {
        let cs = tiny_charset();
        let p = probs(&[0, 0, 0], cs.len());
        let (text, conf) = ctc_decode(&p, 3, cs.len(), &cs);
        assert_eq!(text, "");
        assert_eq!(conf, 0.0);
    }

    #[test]
    fn preprocess_rec_shape_and_padding() {
        // A 24x24 gray crop scales to 48x48; columns past 48 stay zero-padded.
        let crop = image::RgbImage::from_pixel(24, 24, image::Rgb([128, 128, 128]));
        let t = preprocess_rec(&crop);
        assert_eq!(t.shape(), &[1, 3, 48, 320]);
        // In-image region is normalized ~0 ((128/255 - 0.5) / 0.5 ≈ 0.0039)
        assert!((t[[0, 0, 0, 0]] - 0.003_921_57).abs() < 1e-4);
        // Padding region stays exactly zero.
        assert_eq!(t[[0, 0, 0, 319]], 0.0);
        assert_eq!(t[[0, 2, 47, 100]], 0.0);
    }
}