oar_ocr_core/utils/

transform.rs

//! Image transformation utilities for OCR processing.
//!
//! This module provides functions for perspective transformation and image warping,
//! which are essential for correcting skewed text regions in images.

use crate::core::OCRError;
use crate::processors::Point;
use image::{Rgb, RgbImage, imageops};
use nalgebra::{Matrix3, Vector3};
use rayon::prelude::*;
use tracing::debug;

/// Calculates the Euclidean distance between two points.
///
/// # Arguments
///
/// * `p1` - First point
/// * `p2` - Second point
///
/// # Returns
///
/// The distance between the two points.
fn distance(p1: &Point, p2: &Point) -> f32 {
    (p1.x - p2.x).hypot(p1.y - p2.y)
}

/// Extracts a rotated and cropped image from a source image based on bounding box points.
///
/// This function takes a source image and a set of four points that define a quadrilateral
/// region in the image. It crops the image to the bounding box of these points, then applies
/// a perspective transformation to produce a rectified image of the region. If the resulting
/// image has an aspect ratio that suggests it's rotated, it will be automatically rotated.
///
/// # Arguments
///
/// * `src_image` - The source image to crop from
/// * `box_points` - Array of exactly 4 points defining the quadrilateral region
///
/// # Returns
///
/// A Result containing the cropped and transformed image, or an OCRError if the operation fails.
///
/// # Errors
///
/// Returns an OCRError if:
/// * The box_points array doesn't contain exactly 4 points
/// * The calculated crop region is invalid
/// * The calculated crop dimensions are zero
/// * The perspective transformation fails
pub fn get_rotate_crop_image(
    src_image: &RgbImage,
    box_points: &[Point],
) -> Result<RgbImage, OCRError> {
    // Validate input
    if box_points.len() != 4 {
        return Err(OCRError::InvalidInput {
            message: "Box must contain exactly 4 points".to_string(),
        });
    }

    // Find bounding box of the points
    let mut min_x = f32::INFINITY;
    let mut max_x = f32::NEG_INFINITY;
    let mut min_y = f32::INFINITY;
    let mut max_y = f32::NEG_INFINITY;

    for p in box_points {
        min_x = min_x.min(p.x);
        max_x = max_x.max(p.x);
        min_y = min_y.min(p.y);
        max_y = max_y.max(p.y);
    }

    // Calculate crop boundaries, clamping to image dimensions
    let left = min_x.max(0.0) as u32;
    let top = min_y.max(0.0) as u32;
    let right = max_x.min(src_image.width() as f32) as u32;
    let bottom = max_y.min(src_image.height() as f32) as u32;

    // Validate crop region
    if right <= left || bottom <= top {
        return Err(OCRError::InvalidInput {
            message: "Invalid crop region".to_string(),
        });
    }

    // Perform initial crop
    let crop_width = right - left;
    let crop_height = bottom - top;
    let img_crop = imageops::crop_imm(src_image, left, top, crop_width, crop_height).to_image();

    // Adjust points relative to the cropped image
    let points: Vec<Point> = box_points
        .iter()
        .map(|p| Point::new(p.x - left as f32, p.y - top as f32))
        .collect();

    // Reorder points to (top-left, top-right, bottom-right, bottom-left)
    // to keep width/height estimation stable when point order varies.
    let mut sorted = points.clone();
    sorted.sort_by(|a, b| a.x.partial_cmp(&b.x).unwrap_or(std::cmp::Ordering::Equal));
    let (mut index_a, mut index_d) = (0usize, 1usize);
    if sorted[1].y < sorted[0].y {
        index_a = 1;
        index_d = 0;
    }
    let (mut index_b, mut index_c) = (2usize, 3usize);
    if sorted[3].y < sorted[2].y {
        index_b = 3;
        index_c = 2;
    }
    let ordered = [
        sorted[index_a],
        sorted[index_b],
        sorted[index_c],
        sorted[index_d],
    ];

    // Calculate target image dimensions based on the max opposite-edge lengths
    let width1 = distance(&ordered[0], &ordered[1]);
    let width2 = distance(&ordered[2], &ordered[3]);
    let img_crop_width = width1.max(width2).round() as u32;

    let height1 = distance(&ordered[0], &ordered[3]);
    let height2 = distance(&ordered[1], &ordered[2]);
    let img_crop_height = height1.max(height2).round() as u32;

    // Validate target dimensions
    if img_crop_width == 0 || img_crop_height == 0 {
        return Err(OCRError::InvalidInput {
            message: "Invalid crop dimensions".to_string(),
        });
    }

    // Define standard points for the target rectangle
    let pts_std = [
        Point::new(0.0, 0.0),
        Point::new(img_crop_width as f32, 0.0),
        Point::new(img_crop_width as f32, img_crop_height as f32),
        Point::new(0.0, img_crop_height as f32),
    ];

    // Calculate perspective transformation matrix
    let transform_matrix = get_perspective_transform(&ordered, &pts_std)?;

    // Apply perspective transformation
    let dst_img = warp_perspective(
        &img_crop,
        &transform_matrix,
        img_crop_width,
        img_crop_height,
    )?;

    // Automatically rotate if the aspect ratio suggests the text is vertical
    if dst_img.height() as f32 >= dst_img.width() as f32 * 1.5 {
        debug!(
            "Rotating image due to aspect ratio: {}x{}",
            dst_img.width(),
            dst_img.height()
        );

        Ok(imageops::rotate270(&dst_img))
    } else {
        Ok(dst_img)
    }
}

/// Calculates the perspective transformation matrix that maps source points to destination points.
///
/// This function solves the linear system of equations to find the perspective transformation
/// matrix that maps four source points to four destination points.
///
/// # Arguments
///
/// * `src_points` - Array of exactly 4 source points
/// * `dst_points` - Array of exactly 4 destination points
///
/// # Returns
///
/// A Result containing the 3x3 transformation matrix, or an OCRError if the operation fails.
///
/// # Errors
///
/// Returns an OCRError if:
/// * Either array doesn't contain exactly 4 points
/// * The linear system cannot be solved
fn get_perspective_transform(
    src_points: &[Point],
    dst_points: &[Point],
) -> Result<Matrix3<f32>, OCRError> {
    // Validate input
    if src_points.len() != 4 || dst_points.len() != 4 {
        return Err(OCRError::InvalidInput {
            message: "Need exactly 4 points for perspective transformation".to_string(),
        });
    }

    // Set up the linear system of equations
    let mut a = nalgebra::DMatrix::<f32>::zeros(8, 8);
    let mut b = nalgebra::DVector::<f32>::zeros(8);

    // Fill the matrix A and vector b with the equations for perspective transformation
    for i in 0..4 {
        let src = &src_points[i];
        let dst = &dst_points[i];

        // First equation for x coordinate transformation
        a.set_row(
            i * 2,
            &nalgebra::RowDVector::from_row_slice(&[
                src.x,
                src.y,
                1.0,
                0.0,
                0.0,
                0.0,
                -src.x * dst.x,
                -src.y * dst.x,
            ]),
        );
        b[i * 2] = dst.x;

        // Second equation for y coordinate transformation
        a.set_row(
            i * 2 + 1,
            &nalgebra::RowDVector::from_row_slice(&[
                0.0,
                0.0,
                0.0,
                src.x,
                src.y,
                1.0,
                -src.x * dst.y,
                -src.y * dst.y,
            ]),
        );
        b[i * 2 + 1] = dst.y;
    }

    // Solve the linear system to find the transformation parameters
    let decomp = a.lu();
    let solution = decomp.solve(&b).ok_or_else(|| OCRError::InvalidInput {
        message: "Cannot solve perspective transformation".to_string(),
    })?;

    // Construct the 3x3 transformation matrix
    Ok(Matrix3::new(
        solution[0],
        solution[1],
        solution[2],
        solution[3],
        solution[4],
        solution[5],
        solution[6],
        solution[7],
        1.0,
    ))
}

/// Applies a perspective transformation to an image.
///
/// This function transforms an image using a given perspective transformation matrix.
/// It uses inverse mapping with bicubic interpolation to produce the output image.
///
/// # Arguments
///
/// * `src_image` - The source image to transform
/// * `transform_matrix` - The 3x3 perspective transformation matrix
/// * `dst_width` - Width of the output image
/// * `dst_height` - Height of the output image
///
/// # Returns
///
/// A Result containing the transformed image, or an OCRError if the operation fails.
///
/// # Errors
///
/// Returns an OCRError if:
/// * The transformation matrix cannot be inverted
fn warp_perspective(
    src_image: &RgbImage,
    transform_matrix: &Matrix3<f32>,
    dst_width: u32,
    dst_height: u32,
) -> Result<RgbImage, OCRError> {
    // Calculate the inverse transformation matrix for inverse mapping
    let inv_matrix = transform_matrix
        .try_inverse()
        .ok_or_else(|| OCRError::InvalidInput {
            message: "Cannot invert transformation matrix".to_string(),
        })?;

    // Create the destination image
    let mut dst_image = RgbImage::new(dst_width, dst_height);
    let buffer: &mut [u8] = dst_image.as_mut();

    // Inverse-map each destination row via `warp_row`, which walks the source
    // coordinate incrementally along the row instead of recomputing the 3x3
    // homography mat-vec per pixel. Small-image fast path avoids rayon overhead.
    // Bicubic with border replication (matches cv2.warpPerspective, INTER_CUBIC,
    // BORDER_REPLICATE).
    if dst_height <= 1 {
        let row_buffer = &mut buffer[0..(dst_width * 3) as usize];
        warp_row(&inv_matrix, src_image, 0, dst_width, row_buffer);
    } else {
        buffer
            .par_chunks_mut((dst_width * 3) as usize)
            .enumerate()
            .for_each(|(dst_y, row_buffer)| {
                warp_row(&inv_matrix, src_image, dst_y as u32, dst_width, row_buffer);
            });
    }

    Ok(dst_image)
}

/// Inverse-maps a single destination row through the perspective `inv_matrix`,
/// writing bicubic-sampled RGB pixels into `row_buffer`.
///
/// The source coordinate is recomputed per pixel with a full `inv_matrix *
/// [dst_x, dst_y, 1]` mat-vec, bit-identical to `cv2.warpPerspective`. A
/// per-row incremental variant (carry `src` and add the column-0 step each
/// pixel) was benchmarked and gave no measurable speedup — the cost is
/// dominated by the bicubic sampling and the two perspective divisions, not the
/// mat-vec — so the exact form is kept. This helper exists to share one row
/// loop between the sequential and rayon paths.
#[inline]
fn warp_row(
    inv_matrix: &Matrix3<f32>,
    src_image: &RgbImage,
    dst_y: u32,
    dst_width: u32,
    row_buffer: &mut [u8],
) {
    for dst_x in 0..dst_width {
        let src_point = inv_matrix * Vector3::new(dst_x as f32, dst_y as f32, 1.0);
        let final_pixel = if src_point.z.abs() > f32::EPSILON {
            let src_x = src_point.x / src_point.z;
            let src_y = src_point.y / src_point.z;
            // bicubic_interpolate handles out-of-bounds via border replication
            bicubic_interpolate(src_image, src_x, src_y)
        } else {
            // Degenerate case: replicate top-left corner pixel
            *src_image.get_pixel(0, 0)
        };
        let index = (dst_x * 3) as usize;
        row_buffer[index..index + 3].copy_from_slice(&final_pixel.0);
    }
}

/// Gets a pixel value with border replication for out-of-bounds coordinates.
///
/// This function implements OpenCV's BORDER_REPLICATE behavior:
/// when coordinates are outside the image, the nearest edge pixel is used.
///
/// Now used only by the test-only bilinear reference; `bicubic_interpolate`
/// inlines the equivalent clamping against the raw buffer.
///
/// # Arguments
///
/// * `image` - The source image
/// * `x` - X coordinate (can be negative or >= width)
/// * `y` - Y coordinate (can be negative or >= height)
///
/// # Returns
///
/// The pixel value at the clamped coordinates.
#[cfg(test)]
#[inline]
fn get_pixel_replicate(image: &RgbImage, x: i32, y: i32) -> Rgb<u8> {
    let clamped_x = x.clamp(0, image.width() as i32 - 1) as u32;
    let clamped_y = y.clamp(0, image.height() as i32 - 1) as u32;
    *image.get_pixel(clamped_x, clamped_y)
}

/// Cubic interpolation kernel function.
///
/// This implements the standard cubic convolution kernel used in bicubic interpolation.
/// The kernel is defined as:
/// - For |t| <= 1: (a+2)|t|³ - (a+3)|t|² + 1
/// - For 1 < |t| < 2: a|t|³ - 5a|t|² + 8a|t| - 4a
/// - Otherwise: 0
///
/// Where a = -0.5 (Catmull-Rom spline, same as OpenCV's default)
#[inline]
fn cubic_kernel(t: f32) -> f32 {
    const A: f32 = -0.5; // Catmull-Rom spline coefficient (OpenCV default)
    let t_abs = t.abs();

    if t_abs <= 1.0 {
        (A + 2.0) * t_abs * t_abs * t_abs - (A + 3.0) * t_abs * t_abs + 1.0
    } else if t_abs < 2.0 {
        A * t_abs * t_abs * t_abs - 5.0 * A * t_abs * t_abs + 8.0 * A * t_abs - 4.0 * A
    } else {
        0.0
    }
}

/// Performs bicubic interpolation to get a pixel value at non-integer coordinates.
///
/// This function calculates the pixel value at a fractional (x, y) coordinate
/// by interpolating using a 4x4 neighborhood of pixels with cubic convolution.
/// Uses border replication for edge handling (same as OpenCV's BORDER_REPLICATE).
///
/// # Arguments
///
/// * `image` - The source image
/// * `x` - X coordinate (can be fractional)
/// * `y` - Y coordinate (can be fractional)
///
/// # Returns
///
/// The interpolated pixel value.
fn bicubic_interpolate(image: &RgbImage, x: f32, y: f32) -> Rgb<u8> {
    let x_int = x.floor() as i32;
    let y_int = y.floor() as i32;
    let dx = x - x_int as f32;
    let dy = y - y_int as f32;

    // Calculate x-direction weights
    let wx = [
        cubic_kernel(dx + 1.0),
        cubic_kernel(dx),
        cubic_kernel(dx - 1.0),
        cubic_kernel(dx - 2.0),
    ];

    // Calculate y-direction weights
    let wy = [
        cubic_kernel(dy + 1.0),
        cubic_kernel(dy),
        cubic_kernel(dy - 1.0),
        cubic_kernel(dy - 2.0),
    ];

    // Precompute the four border-replicated sample columns and rows once (8
    // clamps total) instead of re-clamping inside the 4x4 loop (which did 16
    // clamped `get_pixel` lookups with their bounds checks). Then index the raw
    // interleaved buffer directly. This is bit-identical to the previous
    // `get_pixel_replicate`-based version: same clamps, same accumulation order
    // (row-major j,i with channel innermost), same round/clamp.
    let w = image.width() as i32;
    let h = image.height() as i32;
    let raw = image.as_raw();
    let stride = (w as usize) * 3;
    let cx = [
        (x_int - 1).clamp(0, w - 1) as usize * 3,
        x_int.clamp(0, w - 1) as usize * 3,
        (x_int + 1).clamp(0, w - 1) as usize * 3,
        (x_int + 2).clamp(0, w - 1) as usize * 3,
    ];
    let cy = [
        (y_int - 1).clamp(0, h - 1) as usize * stride,
        y_int.clamp(0, h - 1) as usize * stride,
        (y_int + 1).clamp(0, h - 1) as usize * stride,
        (y_int + 2).clamp(0, h - 1) as usize * stride,
    ];

    let mut result = [0.0f32; 3];
    for (j, &weight_y) in wy.iter().enumerate() {
        let row = cy[j];
        for (i, &weight_x) in wx.iter().enumerate() {
            let weight = weight_x * weight_y;
            let idx = row + cx[i];
            result[0] += weight * raw[idx] as f32;
            result[1] += weight * raw[idx + 1] as f32;
            result[2] += weight * raw[idx + 2] as f32;
        }
    }

    // Clamp and convert to u8
    Rgb([
        result[0].round().clamp(0.0, 255.0) as u8,
        result[1].round().clamp(0.0, 255.0) as u8,
        result[2].round().clamp(0.0, 255.0) as u8,
    ])
}

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

    /// Performs bilinear interpolation to get a pixel value at non-integer coordinates.
    ///
    /// This function calculates the pixel value at a fractional (x, y) coordinate
    /// by interpolating between the four nearest pixels.
    /// Uses border replication for edge handling (same as OpenCV's BORDER_REPLICATE).
    fn bilinear_interpolate(image: &RgbImage, x: f32, y: f32) -> Rgb<u8> {
        let x_int = x.floor() as i32;
        let y_int = y.floor() as i32;

        // Calculate the fractional parts
        let dx = x - x_int as f32;
        let dy = y - y_int as f32;

        // Get the four neighboring pixels with border replication
        let p11 = get_pixel_replicate(image, x_int, y_int);
        let p12 = get_pixel_replicate(image, x_int, y_int + 1);
        let p21 = get_pixel_replicate(image, x_int + 1, y_int);
        let p22 = get_pixel_replicate(image, x_int + 1, y_int + 1);

        // Interpolate each color channel
        let mut result = [0u8; 3];
        for (i, result_channel) in result.iter_mut().enumerate() {
            let val = (1.0 - dx) * (1.0 - dy) * p11.0[i] as f32
                + dx * (1.0 - dy) * p21.0[i] as f32
                + (1.0 - dx) * dy * p12.0[i] as f32
                + dx * dy * p22.0[i] as f32;
            *result_channel = val.round().clamp(0.0, 255.0) as u8;
        }

        Rgb(result)
    }

    /// Independent reference for `bicubic_interpolate`, written exactly as the
    /// previous `get_pixel_replicate`-based implementation. Used to prove the
    /// raw-buffer rewrite is bit-identical.
    fn bicubic_reference(image: &RgbImage, x: f32, y: f32) -> Rgb<u8> {
        let x_int = x.floor() as i32;
        let y_int = y.floor() as i32;
        let dx = x - x_int as f32;
        let dy = y - y_int as f32;
        let wx = [
            cubic_kernel(dx + 1.0),
            cubic_kernel(dx),
            cubic_kernel(dx - 1.0),
            cubic_kernel(dx - 2.0),
        ];
        let wy = [
            cubic_kernel(dy + 1.0),
            cubic_kernel(dy),
            cubic_kernel(dy - 1.0),
            cubic_kernel(dy - 2.0),
        ];
        let mut result = [0.0f32; 3];
        for (j, &weight_y) in wy.iter().enumerate() {
            let sample_y = y_int - 1 + j as i32;
            for (i, &weight_x) in wx.iter().enumerate() {
                let sample_x = x_int - 1 + i as i32;
                let weight = weight_x * weight_y;
                let pixel = get_pixel_replicate(image, sample_x, sample_y);
                for (c, result_c) in result.iter_mut().enumerate().take(3) {
                    *result_c += weight * pixel.0[c] as f32;
                }
            }
        }
        Rgb([
            result[0].round().clamp(0.0, 255.0) as u8,
            result[1].round().clamp(0.0, 255.0) as u8,
            result[2].round().clamp(0.0, 255.0) as u8,
        ])
    }

    #[test]
    fn bicubic_raw_buffer_matches_reference_bit_exact() {
        // Deterministic pseudo-random image.
        let (w, h) = (17u32, 11u32);
        let img = RgbImage::from_fn(w, h, |x, y| {
            let i = y * w + x;
            Rgb([
                (i.wrapping_mul(37).wrapping_add(11) % 256) as u8,
                (i.wrapping_mul(59).wrapping_add(7) % 256) as u8,
                (i.wrapping_mul(101).wrapping_add(3) % 256) as u8,
            ])
        });
        // Sample fractional coords across the interior, edges, and out-of-bounds
        // (negative and beyond width/height) to exercise border replication.
        for yi in -3..(h as i32 + 3) {
            for xi in -3..(w as i32 + 3) {
                for &fx in &[0.0f32, 0.25, 0.5, 0.75] {
                    for &fy in &[0.0f32, 0.33, 0.66] {
                        let x = xi as f32 + fx;
                        let y = yi as f32 + fy;
                        assert_eq!(
                            bicubic_interpolate(&img, x, y),
                            bicubic_reference(&img, x, y),
                            "mismatch at ({x}, {y})"
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_distance() {
        let p1 = Point::new(0.0, 0.0);
        let p2 = Point::new(3.0, 4.0);
        let dist = distance(&p1, &p2);
        assert_eq!(dist, 5.0);
    }

    #[test]
    fn test_get_perspective_transform() -> Result<(), OCRError> {
        // Define a simple square in source and destination
        let src_points = [
            Point::new(0.0, 0.0),
            Point::new(1.0, 0.0),
            Point::new(1.0, 1.0),
            Point::new(0.0, 1.0),
        ];

        let dst_points = [
            Point::new(0.0, 0.0),
            Point::new(2.0, 0.0),
            Point::new(2.0, 2.0),
            Point::new(0.0, 2.0),
        ];

        let transform = get_perspective_transform(&src_points, &dst_points)?;

        // Check that the transformation matrix is valid (all elements are finite)
        assert!(transform.iter().all(|&x| x.is_finite()));
        Ok(())
    }

    #[test]
    fn test_get_perspective_transform_invalid_input() {
        // Test with wrong number of points
        let src_points = [Point::new(0.0, 0.0), Point::new(1.0, 0.0)];

        let dst_points = [
            Point::new(0.0, 0.0),
            Point::new(2.0, 0.0),
            Point::new(2.0, 2.0),
            Point::new(0.0, 2.0),
        ];

        let result = get_perspective_transform(&src_points, &dst_points);
        assert!(result.is_err());
    }

    #[test]
    fn test_get_rotate_crop_image_invalid_points() {
        // Create a simple 4x4 image
        let image = RgbImage::new(4, 4);

        // Test with wrong number of points
        let points = vec![Point::new(0.0, 0.0), Point::new(1.0, 0.0)];

        let result = get_rotate_crop_image(&image, &points);
        assert!(result.is_err());
    }

    #[test]
    fn test_get_rotate_crop_image_success() -> Result<(), OCRError> {
        // Create a simple 4x4 image with distinct colors
        let mut image = RgbImage::new(4, 4);
        for y in 0..4 {
            for x in 0..4 {
                // Create a gradient
                let r = (x * 64) as u8;
                let g = (y * 64) as u8;
                let b = ((x + y) * 32) as u8;
                image.put_pixel(x, y, Rgb([r, g, b]));
            }
        }

        // Define a simple square region
        let points = vec![
            Point::new(1.0, 1.0),
            Point::new(3.0, 1.0),
            Point::new(3.0, 3.0),
            Point::new(1.0, 3.0),
        ];

        let cropped_image = get_rotate_crop_image(&image, &points)?;
        // Check that we got an image back
        assert!(cropped_image.width() > 0);
        assert!(cropped_image.height() > 0);
        Ok(())
    }

    #[test]
    fn test_warp_perspective_invalid_matrix() {
        // Create a simple 2x2 image
        let image = RgbImage::new(2, 2);

        // Create a singular matrix (non-invertible)
        let matrix = Matrix3::new(1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0);

        let result = warp_perspective(&image, &matrix, 2, 2);
        assert!(result.is_err());
    }

    #[test]
    fn test_bilinear_interpolate() {
        // Create a simple 2x2 image with distinct colors
        let mut image = RgbImage::new(2, 2);
        image.put_pixel(0, 0, Rgb([255, 0, 0])); // Red
        image.put_pixel(1, 0, Rgb([0, 255, 0])); // Green
        image.put_pixel(0, 1, Rgb([0, 0, 255])); // Blue
        image.put_pixel(1, 1, Rgb([255, 255, 0])); // Yellow

        // Test interpolation at the center
        let pixel = bilinear_interpolate(&image, 0.5, 0.5);
        // Expected: average of all four colors
        // Red + Green + Blue + Yellow = (255, 0, 0) + (0, 255, 0) + (0, 0, 255) + (255, 255, 0)
        // = (510, 510, 255) / 4 = (127.5, 127.5, 63.75) ≈ (128, 128, 64)
        assert_eq!(pixel.0[0], 128);
        assert_eq!(pixel.0[1], 128);
        assert_eq!(pixel.0[2], 64);
    }
}