visual-cortex-ocr-onnx 0.4.0

use std::path::Path;

use ndarray::Array4;
use ort::session::Session;
use ort::value::TensorRef;
use visual_cortex_capture::{FrameView, PxRect};
use visual_cortex_vision::{DetectorError, OcrEngine, TextSpan};

use crate::det::{det_target_size, extract_boxes, preprocess_det, view_to_rgb};
use crate::error::OcrError;
use crate::models::ensure_all_cached;
use crate::rec::{build_charset, ctc_decode, preprocess_rec};

/// Binarization threshold for the det probability map. Tuned against the
/// hp_1234 e2e fixture: lowered from PaddleOCR's stock 0.3 while diagnosing a
/// clipped leading digit (no box change on its own, kept as the more
/// permissive tested value; false positives are filtered by MIN_CONFIDENCE).
const DET_THRESHOLD: f32 = 0.2;
/// Box padding in det-map pixels (the DB unclip substitute). Tuned 3 -> 5
/// against the hp_1234 e2e fixture: at 3 the rec crop clipped the leading
/// "1" and read "HP 234".
const BOX_PAD: u32 = 5;
/// Spans below this confidence are discarded.
const MIN_CONFIDENCE: f32 = 0.3;

/// PaddleOCR det+rec pipeline behind visual-cortex-vision's `OcrEngine`.
pub struct PaddleOcr {
    det: Session,
    rec: Session,
    charset: Vec<String>,
}

impl PaddleOcr {
    /// Download (if absent), verify, and load the pinned models from the
    /// platform cache dir (`~/.cache/visual_cortex` or the macOS/Windows equivalent).
    pub async fn new() -> Result<Self, OcrError> {
        let (det, rec, dict) = ensure_all_cached().await?;
        Self::from_paths(&det, &rec, &dict)
    }

    /// Load models from explicit local paths. Never touches the network.
    pub fn from_paths(det: &Path, rec: &Path, dict: &Path) -> Result<Self, OcrError> {
        let det = Session::builder()
            .and_then(|mut b| b.commit_from_file(det))
            .map_err(|e| OcrError::ModelLoad(format!("det: {e}")))?;
        let rec = Session::builder()
            .and_then(|mut b| b.commit_from_file(rec))
            .map_err(|e| OcrError::ModelLoad(format!("rec: {e}")))?;
        let dict =
            std::fs::read_to_string(dict).map_err(|e| OcrError::ModelLoad(format!("dict: {e}")))?;
        Ok(Self {
            det,
            rec,
            charset: build_charset(&dict),
        })
    }

    fn run_det(&mut self, input: Array4<f32>, tw: u32, th: u32) -> Result<Vec<f32>, OcrError> {
        let outputs = self
            .det
            .run(ort::inputs!["x" => TensorRef::from_array_view(input.view())
                .map_err(|e| OcrError::Inference(format!("det input: {e}")))?])
            .map_err(|e| OcrError::Inference(format!("det run: {e}")))?;
        let map = outputs["sigmoid_0.tmp_0"]
            .try_extract_array::<f32>()
            .map_err(|e| OcrError::Inference(format!("det output: {e}")))?;
        let flat: Vec<f32> = map.iter().copied().collect();
        let expected = tw as usize * th as usize;
        if flat.len() != expected {
            return Err(OcrError::Inference(format!(
                "det output has {} elements, expected {tw}x{th} = {expected}",
                flat.len()
            )));
        }
        Ok(flat)
    }

    fn run_rec(&mut self, input: Array4<f32>) -> Result<(String, f32), OcrError> {
        let outputs = self
            .rec
            .run(ort::inputs!["x" => TensorRef::from_array_view(input.view())
                .map_err(|e| OcrError::Inference(format!("rec input: {e}")))?])
            .map_err(|e| OcrError::Inference(format!("rec run: {e}")))?;
        let probs = outputs["softmax_2.tmp_0"]
            .try_extract_array::<f32>()
            .map_err(|e| OcrError::Inference(format!("rec output: {e}")))?;
        let shape = probs.shape().to_vec(); // [1, T, C]
        if shape.len() != 3 || shape[2] != self.charset.len() {
            return Err(OcrError::Inference(format!(
                "rec output shape {shape:?} does not match charset len {}",
                self.charset.len()
            )));
        }
        let flat: Vec<f32> = probs.iter().copied().collect();
        Ok(ctc_decode(&flat, shape[1], shape[2], &self.charset))
    }
}

impl OcrEngine for PaddleOcr {
    fn recognize(&mut self, view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
        let rgb = view_to_rgb(view);
        let (vw, vh) = (rgb.width(), rgb.height());
        let (tw, th) = det_target_size(vw, vh);
        let det_input = preprocess_det(&rgb, tw, th);
        let prob = self
            .run_det(det_input, tw, th)
            .map_err(|e| DetectorError::Ocr(e.to_string()))?;
        let boxes = extract_boxes(&prob, tw as usize, th as usize, DET_THRESHOLD, BOX_PAD);

        // Map det-map boxes back to view coordinates.
        let (sx, sy) = (vw as f32 / tw as f32, vh as f32 / th as f32);
        let mut spans = Vec::new();
        for b in boxes {
            let x = ((b.x as f32 * sx) as u32).min(vw - 1);
            let y = ((b.y as f32 * sy) as u32).min(vh - 1);
            let w = ((b.w as f32 * sx).ceil() as u32).clamp(1, vw - x);
            let h = ((b.h as f32 * sy).ceil() as u32).clamp(1, vh - y);
            let crop = image::imageops::crop_imm(&rgb, x, y, w, h).to_image();
            let (text, confidence) = self
                .run_rec(preprocess_rec(&crop))
                .map_err(|e| DetectorError::Ocr(e.to_string()))?;
            if !text.is_empty() && confidence >= MIN_CONFIDENCE {
                spans.push(TextSpan {
                    text,
                    confidence,
                    bbox: PxRect { x, y, w, h },
                });
            }
        }
        Ok(spans)
    }
}