visual-cortex-ocr-onnx 0.7.0

use image::RgbImage;
use ndarray::Array4;
use visual_cortex_capture::{FrameView, PxRect};

/// ImageNet normalization used by the PP-OCR detection models.
const DET_MEAN: [f32; 3] = [0.485, 0.456, 0.406];
const DET_STD: [f32; 3] = [0.229, 0.224, 0.225];
/// Longest side the det model sees; larger views are scaled down.
const DET_MAX_SIDE: u32 = 960;
/// Minimum blob size (w, h) kept as a text box, in probability-map pixels.
const MIN_BOX_W: u32 = 3;
const MIN_BOX_H: u32 = 2;

/// Compute the det-model input size for a view: aspect-preserving scale to at
/// most `DET_MAX_SIDE`, then each dimension rounded to a multiple of 32
/// (minimum 32).
pub(crate) fn det_target_size(w: u32, h: u32) -> (u32, u32) {
    let scale = (DET_MAX_SIDE as f32 / w.max(h) as f32).min(1.0);
    let round32 = |v: f32| (((v / 32.0).round() as u32).max(1)) * 32;
    (round32(w as f32 * scale), round32(h as f32 * scale))
}

/// Convert a BGRA view to an RGB image (owned, contiguous).
pub(crate) fn view_to_rgb(view: &FrameView<'_>) -> RgbImage {
    let (w, h) = (view.width(), view.height());
    let mut img = RgbImage::new(w, h);
    for (y, row) in view.rows().enumerate() {
        for (x, px) in row.chunks_exact(4).enumerate() {
            img.put_pixel(x as u32, y as u32, image::Rgb([px[2], px[1], px[0]]));
        }
    }
    img
}

/// Resize + normalize an RGB image into the det model's NCHW input tensor.
pub(crate) fn preprocess_det(img: &RgbImage, tw: u32, th: u32) -> Array4<f32> {
    let resized = image::imageops::resize(img, tw, th, image::imageops::FilterType::Triangle);
    let mut input = Array4::<f32>::zeros((1, 3, th as usize, tw 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 - DET_MEAN[c]) / DET_STD[c];
        }
    }
    input
}

/// Extract axis-aligned text boxes from a binarized probability map via
/// 4-connected flood fill. A deliberate simplification of DB polygon
/// unclipping — adequate for horizontal UI/tooltip text. `pad` inflates each
/// box (clamped to the map bounds) as the unclip substitute.
pub(crate) fn extract_boxes(
    prob: &[f32],
    pw: usize,
    ph: usize,
    threshold: f32,
    pad: u32,
) -> Vec<PxRect> {
    let mut visited = vec![false; pw * ph];
    let mut boxes = Vec::new();
    for start in 0..pw * ph {
        if visited[start] || prob[start] < threshold {
            continue;
        }
        // Flood fill this component, tracking its bounding rect.
        let (mut min_x, mut min_y, mut max_x, mut max_y) =
            (start % pw, start / pw, start % pw, start / pw);
        let mut stack = vec![start];
        visited[start] = true;
        while let Some(idx) = stack.pop() {
            let (x, y) = (idx % pw, idx / pw);
            min_x = min_x.min(x);
            max_x = max_x.max(x);
            min_y = min_y.min(y);
            max_y = max_y.max(y);
            let neighbors = [
                (x > 0).then(|| idx - 1),
                (x + 1 < pw).then(|| idx + 1),
                (y > 0).then(|| idx - pw),
                (y + 1 < ph).then(|| idx + pw),
            ];
            for n in neighbors.into_iter().flatten() {
                if !visited[n] && prob[n] >= threshold {
                    visited[n] = true;
                    stack.push(n);
                }
            }
        }
        let (w, h) = ((max_x - min_x + 1) as u32, (max_y - min_y + 1) as u32);
        if w < MIN_BOX_W || h < MIN_BOX_H {
            continue;
        }
        // Pad each edge independently, clamped to the map bounds (padding
        // absorbed by a clamp is lost, not shifted to the other side).
        let x0 = (min_x as u32).saturating_sub(pad);
        let y0 = (min_y as u32).saturating_sub(pad);
        let x1 = (max_x as u32 + 1 + pad).min(pw as u32);
        let y1 = (max_y as u32 + 1 + pad).min(ph as u32);
        boxes.push(PxRect {
            x: x0,
            y: y0,
            w: x1 - x0,
            h: y1 - y0,
        });
    }
    boxes
}

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

    #[test]
    fn det_target_size_scales_and_rounds_to_32() {
        // Small input: rounded to multiples of 32, minimum 32.
        assert_eq!(det_target_size(100, 40), (96, 32));
        // Longest side capped at 960, aspect preserved before rounding.
        let (w, h) = det_target_size(3840, 1920);
        assert_eq!(w, 960);
        assert_eq!(h, 480);
        assert_eq!(w % 32, 0);
        assert_eq!(h % 32, 0);
    }

    #[test]
    fn extract_boxes_finds_two_blobs() {
        // 32x16 probability map with two bright rectangles.
        let (pw, ph) = (32usize, 16usize);
        let mut prob = vec![0.0f32; pw * ph];
        for y in 2..6 {
            for x in 2..12 {
                prob[y * pw + x] = 0.9;
            }
        }
        for y in 9..13 {
            for x in 18..30 {
                prob[y * pw + x] = 0.9;
            }
        }
        let mut boxes = extract_boxes(&prob, pw, ph, 0.3, 0);
        boxes.sort_by_key(|r| r.y);
        assert_eq!(boxes.len(), 2);
        assert_eq!(
            (boxes[0].x, boxes[0].y, boxes[0].w, boxes[0].h),
            (2, 2, 10, 4)
        );
        assert_eq!(
            (boxes[1].x, boxes[1].y, boxes[1].w, boxes[1].h),
            (18, 9, 12, 4)
        );
    }

    #[test]
    fn extract_boxes_pads_and_clamps_at_edges() {
        let (pw, ph) = (16usize, 8usize);
        let mut prob = vec![0.0f32; pw * ph];
        // Blob touching the top-left corner.
        for y in 0..3 {
            for x in 0..4 {
                prob[y * pw + x] = 1.0;
            }
        }
        let boxes = extract_boxes(&prob, pw, ph, 0.3, 2);
        assert_eq!(boxes.len(), 1);
        // Padding clamps at 0 and at the map edge.
        assert_eq!((boxes[0].x, boxes[0].y), (0, 0));
        assert_eq!((boxes[0].w, boxes[0].h), (6, 5));
    }

    #[test]
    fn extract_boxes_ignores_tiny_specks_and_empty_maps() {
        let (pw, ph) = (16usize, 8usize);
        let mut prob = vec![0.0f32; pw * ph];
        prob[3 * pw + 3] = 1.0; // single pixel: below the 3x2 minimum
        assert!(extract_boxes(&prob, pw, ph, 0.3, 0).is_empty());
        let empty = vec![0.0f32; pw * ph];
        assert!(extract_boxes(&empty, pw, ph, 0.3, 0).is_empty());
    }
}