captchaforge 0.2.38

//! CRNN text recognition for distorted text CAPTCHAs.
//!
//! Uses an ONNX-exported CRNN model (CNN feature extraction +
//! BiLSTM sequence modelling + CTC decoding) trained on synthetic
//! CAPTCHA data.
//!
//! # Model
//!
//! The solver expects a model at `~/.captchaforge/models/crnn_text/crnn_text.onnx`.
//! If the model is missing, the solver is inert and the chain falls
//! back to Tesseract OCR or VLM.
//!
//! To train a model, run the Python script shipped in `scripts/train_crnn.py`.

use anyhow::{Context, Result};
use image::{imageops, DynamicImage};

pub const CRNN_INPUT_WIDTH: u32 = 160;
pub const CRNN_INPUT_HEIGHT: u32 = 60;

/// CRNN text recognizer backed by ONNX Runtime.
pub struct CrnnRecognizer {
    session: ort::session::Session,
    charset: Vec<char>,
}

impl CrnnRecognizer {
    /// Load a CRNN ONNX model from disk.
    pub fn load(model_path: &std::path::Path, charset: Vec<char>) -> Result<Self> {
        let session = ort::session::Session::builder()
            .map_err(|e| anyhow::anyhow!("create ONNX session builder: {e}"))?
            .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)
            .map_err(|e| anyhow::anyhow!("set graph optimization: {e}"))?
            .with_execution_providers([
                ort::execution_providers::CUDAExecutionProvider::default().build(),
                ort::execution_providers::TensorRTExecutionProvider::default().build(),
                ort::execution_providers::CPUExecutionProvider::default().build(),
            ])
            .map_err(|e| anyhow::anyhow!("set execution providers: {e}"))?
            .commit_from_file(model_path)
            .map_err(|e| anyhow::anyhow!("load ONNX model from {}: {e}", model_path.display()))?;

        tracing::info!(model = %model_path.display(), charset_len = charset.len(), "CRNN recognizer loaded");
        Ok(Self { session, charset })
    }

    /// Recognize text in a CAPTCHA image.
    pub fn recognize(&mut self, image: &DynamicImage) -> Result<(String, f32)> {
        // Resize to CRNN input size.
        let resized = image.resize_exact(CRNN_INPUT_WIDTH, CRNN_INPUT_HEIGHT, imageops::FilterType::Triangle);
        let rgb = resized.to_rgb8();

        // NCHW float tensor, normalized.
        let (w, h) = (CRNN_INPUT_WIDTH as usize, CRNN_INPUT_HEIGHT as usize);
        let mut pixel_data = vec![0.0f32; 3 * h * w];
        for y in 0..h {
            for x in 0..w {
                let pixel = rgb.get_pixel(x as u32, y as u32);
                pixel_data[y * w + x] = pixel[0] as f32 / 255.0;
                pixel_data[h * w + y * w + x] = pixel[1] as f32 / 255.0;
                pixel_data[2 * h * w + y * w + x] = pixel[2] as f32 / 255.0;
            }
        }

        let shape = vec![1_i64, 3, h as i64, w as i64];
        let input = ort::value::Tensor::from_array((shape, pixel_data))
            .context("build input tensor")?;

        let outputs = self
            .session
            .run(ort::inputs! { "input" => input })
            .context("run CRNN inference")?;

        let output_view = outputs[0]
            .try_extract_array::<f32>()
            .context("extract output tensor")?;

        let shape = output_view.shape().to_vec();
        let flat: Vec<f32> = output_view.iter().copied().collect();

        // CTC decode: output shape is [T, 1, C] or [1, T, C].
        let (time_steps, num_classes) = if shape.len() == 3 {
            if shape[0] == 1 {
                (shape[1], shape[2])
            } else {
                (shape[0], shape[2])
            }
        } else {
            anyhow::bail!("unexpected CRNN output shape: {:?}", shape);
        };

        let mut decoded = String::new();
        let mut prev_idx = usize::MAX;
        let mut total_conf = 0.0f32;
        let mut count = 0usize;

        for t in 0..time_steps {
            let offset = t * num_classes;
            let frame = &flat[offset..offset + num_classes];
            let (best_idx, best_val) = frame.iter().enumerate()
                .fold((0, frame[0]), |(bi, bv), (i, &v)| {
                    if v > bv { (i, v) } else { (bi, bv) }
                });

            // CTC blank index is typically the last class.
            let blank_idx = num_classes - 1;
            if best_idx != blank_idx && best_idx != prev_idx {
                if let Some(&ch) = self.charset.get(best_idx) {
                    decoded.push(ch);
                }
                total_conf += best_val;
                count += 1;
            }
            prev_idx = best_idx;
        }

        let confidence = if count > 0 { total_conf / count as f32 } else { 0.0 };
        Ok((decoded, confidence.min(1.0)))
    }
}

/// Default charset for alphanumeric CAPTCHAs (digits + uppercase letters).
/// Index 0–9 = '0'–'9', 10–35 = 'A'–'Z', 36 = CTC blank.
pub fn default_charset() -> Vec<char> {
    let mut chars: Vec<char> = ('0'..='9').chain('A'..='Z').collect();
    // Ensure blank token exists at the end.
    chars.push('-');
    chars
}

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

    #[test]
    fn default_charset_length() {
        let cs = default_charset();
        assert_eq!(cs.len(), 37); // 10 digits + 26 letters + 1 blank
    }
}