captcha-engine 0.4.10

//! Image preprocessing for the captcha model.

use image::{DynamicImage, imageops::FilterType};
use rten_tensor::Tensor;
use std::path::Path;

/// Image dimensions for the custom finetuned model.
/// Model expects: `[1, 3, 80, 215]` (NCHW format)
pub const IMG_HEIGHT: u32 = 80;
pub const IMG_WIDTH: u32 = 215;

/// `ImageNet` normalization parameters (used during training)
const MEAN: [f32; 3] = [0.485, 0.456, 0.406];
const STD: [f32; 3] = [0.229, 0.224, 0.225];

/// Preprocess an image for the captcha model.
///
/// Model expects: `[1, 3, 80, 215]` (NCHW format)
#[must_use]
pub fn preprocess(img: &DynamicImage) -> Tensor<f32> {
    // Resize to model input dimensions
    let img = img.resize_exact(IMG_WIDTH, IMG_HEIGHT, FilterType::Triangle);

    // Convert to RGB
    let img = img.to_rgb8();

    // Create tensor data in NCHW format: [1, 3, height, width]
    let mut data = Vec::with_capacity((3 * IMG_HEIGHT * IMG_WIDTH) as usize);

    // Iterating by channel first to be cache friendly?
    // No, standard image iteration is y,x then c.
    // NCHW means we need all Red, then all Green, then all Blue.

    // We can pre-calculate normalized values or better yet, construct plane by plane.
    for c in 0..3 {
        for y in 0..IMG_HEIGHT {
            for x in 0..IMG_WIDTH {
                let pixel = img.get_pixel(x, y);
                let pixel_value = f32::from(pixel[c]) / 255.0;
                // ImageNet normalization: (value - mean) / std
                let normalized = (pixel_value - MEAN[c]) / STD[c];
                data.push(normalized);
            }
        }
    }

    Tensor::from_data(&[1, 3, IMG_HEIGHT as usize, IMG_WIDTH as usize], data)
}

/// Load and preprocess an image from a file path.
///
/// # Errors
///
/// Returns an error if the image cannot be loaded.
pub fn preprocess_file<P: AsRef<Path>>(path: P) -> crate::Result<Tensor<f32>> {
    let img = image::open(path)?;
    Ok(preprocess(&img))
}

#[cfg(test)]
mod tests {
    use super::*;
    use image::RgbImage;
    use rten_tensor::prelude::*;

    #[test]
    fn test_preprocess_output_shape() {
        // Create a dummy image - 100x50
        let img = DynamicImage::ImageRgb8(RgbImage::new(100, 50));
        let tensor = preprocess(&img);

        // Check shape: [1, 3, 80, 215]
        assert_eq!(tensor.shape(), &[1, 3, 80, 215]);
    }

    #[test]
    fn test_preprocess_normalization() {
        // Create an all-white image (255, 255, 255)
        let mut img = RgbImage::new(10, 10);
        for pixel in img.pixels_mut() {
            *pixel = image::Rgb([255, 255, 255]);
        }
        let img = DynamicImage::ImageRgb8(img);
        let tensor = preprocess(&img);

        // Expected value for 1.0 (255/255)
        // Red: (1.0 - 0.485) / 0.229 ≈ 2.2489
        // Green: (1.0 - 0.456) / 0.224 ≈ 2.4286
        // Blue: (1.0 - 0.406) / 0.225 ≈ 2.6400

        // Index manually into the flattened data or use NCHW logic
        // Tensor is [1, 3, 80, 215]
        // data layout is [C0... C1... C2...]
        // so first pixel of C0 is at index 0
        // first pixel of C1 is at index 80*215
        // first pixel of C2 is at index 2*80*215

        // rten Tensor allows indexing? It implements Index but maybe by slice.
        // Let's verify by iterating or just checking specific elements if possible.
        // We can just get `data()` slice.

        let data = tensor.data().expect("should be contiguous");
        let stride = (IMG_HEIGHT * IMG_WIDTH) as usize;

        let first_pixel_r = data[0];
        let first_pixel_g = data[stride];
        let first_pixel_b = data[stride * 2];

        assert!((first_pixel_r - 2.2489).abs() < 0.001);
        assert!((first_pixel_g - 2.4286).abs() < 0.001);
        assert!((first_pixel_b - 2.6400).abs() < 0.001);
    }
}