oar_ocr_core/processors/

uvdoc_postprocess.rs

//! Document transformation post-processing functionality.

use crate::core::OcrResult;
use crate::core::errors::OCRError;
use std::str::FromStr;

/// Post-processor for document transformation results.
///
/// The `UVDocPostProcess` struct handles the post-processing of document
/// transformation model outputs, converting normalized coordinates back
/// to pixel coordinates and applying various transformations.
#[derive(Debug)]
pub struct UVDocPostProcess {
    /// Scale factor to convert normalized values back to pixel values.
    pub scale: f32,
}

impl UVDocPostProcess {
    /// Creates a new UVDocPostProcess instance.
    ///
    /// # Arguments
    ///
    /// * `scale` - Scale factor for converting normalized coordinates to pixels.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use oar_ocr_core::processors::UVDocPostProcess;
    ///
    /// let postprocessor = UVDocPostProcess::new(1.0);
    /// ```
    pub fn new(scale: f32) -> Self {
        Self { scale }
    }

    /// Gets the current scale factor.
    ///
    /// # Returns
    ///
    /// The scale factor used for coordinate conversion.
    pub fn scale(&self) -> f32 {
        self.scale
    }

    /// Sets a new scale factor.
    ///
    /// # Arguments
    ///
    /// * `scale` - New scale factor.
    pub fn set_scale(&mut self, scale: f32) {
        self.scale = scale;
    }

    /// Converts normalized coordinates to pixel coordinates.
    ///
    /// # Arguments
    ///
    /// * `normalized_coords` - Vector of normalized coordinates (0.0 to 1.0).
    ///
    /// # Returns
    ///
    /// * `Vec<f32>` - Vector of pixel coordinates.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use oar_ocr_core::processors::UVDocPostProcess;
    ///
    /// let postprocessor = UVDocPostProcess::new(100.0);
    /// let normalized = vec![0.1, 0.2, 0.8, 0.9];
    /// let pixels = postprocessor.denormalize_coordinates(&normalized);
    /// assert_eq!(pixels, vec![10.0, 20.0, 80.0, 90.0]);
    /// ```
    pub fn denormalize_coordinates(&self, normalized_coords: &[f32]) -> Vec<f32> {
        normalized_coords
            .iter()
            .map(|&coord| coord * self.scale)
            .collect()
    }

    /// Converts pixel coordinates to normalized coordinates.
    ///
    /// # Arguments
    ///
    /// * `pixel_coords` - Vector of pixel coordinates.
    ///
    /// # Returns
    ///
    /// * `Vec<f32>` - Vector of normalized coordinates (0.0 to 1.0).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use oar_ocr_core::processors::UVDocPostProcess;
    ///
    /// let postprocessor = UVDocPostProcess::new(100.0);
    /// let pixels = vec![10.0, 20.0, 80.0, 90.0];
    /// let normalized = postprocessor.normalize_coordinates(&pixels);
    /// assert_eq!(normalized, vec![0.1, 0.2, 0.8, 0.9]);
    /// ```
    pub fn normalize_coordinates(&self, pixel_coords: &[f32]) -> Vec<f32> {
        if self.scale == 0.0 {
            return vec![0.0; pixel_coords.len()];
        }
        pixel_coords
            .iter()
            .map(|&coord| coord / self.scale)
            .collect()
    }

    /// Processes a bounding box from normalized to pixel coordinates.
    ///
    /// # Arguments
    ///
    /// * `bbox` - Bounding box as [x1, y1, x2, y2] in normalized coordinates.
    ///
    /// # Returns
    ///
    /// * `[f32; 4]` - Bounding box in pixel coordinates.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use oar_ocr_core::processors::UVDocPostProcess;
    ///
    /// let postprocessor = UVDocPostProcess::new(100.0);
    /// let normalized_bbox = [0.1, 0.2, 0.8, 0.9];
    /// let pixel_bbox = postprocessor.process_bbox(&normalized_bbox);
    /// assert_eq!(pixel_bbox, [10.0, 20.0, 80.0, 90.0]);
    /// ```
    pub fn process_bbox(&self, bbox: &[f32; 4]) -> [f32; 4] {
        [
            bbox[0] * self.scale,
            bbox[1] * self.scale,
            bbox[2] * self.scale,
            bbox[3] * self.scale,
        ]
    }

    /// Processes multiple bounding boxes.
    ///
    /// # Arguments
    ///
    /// * `bboxes` - Vector of bounding boxes in normalized coordinates.
    ///
    /// # Returns
    ///
    /// * `Vec<[f32; 4]>` - Vector of bounding boxes in pixel coordinates.
    pub fn process_bboxes(&self, bboxes: &[[f32; 4]]) -> Vec<[f32; 4]> {
        bboxes.iter().map(|bbox| self.process_bbox(bbox)).collect()
    }

    /// Processes a polygon from normalized to pixel coordinates.
    ///
    /// # Arguments
    ///
    /// * `polygon` - Vector of points as [x, y] pairs in normalized coordinates.
    ///
    /// # Returns
    ///
    /// * `Vec<[f32; 2]>` - Vector of points in pixel coordinates.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use oar_ocr_core::processors::UVDocPostProcess;
    ///
    /// let postprocessor = UVDocPostProcess::new(100.0);
    /// let normalized_polygon = vec![[0.1, 0.2], [0.8, 0.2], [0.8, 0.9], [0.1, 0.9]];
    /// let pixel_polygon = postprocessor.process_polygon(&normalized_polygon);
    /// assert_eq!(pixel_polygon[0], [10.0, 20.0]);
    /// ```
    pub fn process_polygon(&self, polygon: &[[f32; 2]]) -> Vec<[f32; 2]> {
        polygon
            .iter()
            .map(|&[x, y]| [x * self.scale, y * self.scale])
            .collect()
    }

    /// Clamps coordinates to valid ranges.
    ///
    /// # Arguments
    ///
    /// * `coords` - Vector of coordinates to clamp.
    /// * `min_val` - Minimum allowed value.
    /// * `max_val` - Maximum allowed value.
    ///
    /// # Returns
    ///
    /// * `Vec<f32>` - Vector of clamped coordinates.
    pub fn clamp_coordinates(&self, coords: &[f32], min_val: f32, max_val: f32) -> Vec<f32> {
        coords
            .iter()
            .map(|&coord| coord.clamp(min_val, max_val))
            .collect()
    }

    /// Validates that coordinates are within expected ranges.
    ///
    /// # Arguments
    ///
    /// * `coords` - Vector of coordinates to validate.
    /// * `min_val` - Minimum expected value.
    /// * `max_val` - Maximum expected value.
    ///
    /// # Returns
    ///
    /// * `true` - If all coordinates are within range.
    /// * `false` - If any coordinate is out of range.
    pub fn validate_coordinates(&self, coords: &[f32], min_val: f32, max_val: f32) -> bool {
        coords
            .iter()
            .all(|&coord| coord >= min_val && coord <= max_val)
    }

    /// Rounds coordinates to integer values.
    ///
    /// # Arguments
    ///
    /// * `coords` - Vector of coordinates to round.
    ///
    /// # Returns
    ///
    /// * `Vec<i32>` - Vector of rounded integer coordinates.
    pub fn round_coordinates(&self, coords: &[f32]) -> Vec<i32> {
        coords.iter().map(|&coord| coord.round() as i32).collect()
    }

    /// Applies batch processing to tensor output to produce rectified images.
    ///
    /// # Arguments
    ///
    /// * `output` - 4D tensor output from the model [batch, channels, height, width].
    ///
    /// # Returns
    ///
    /// * `OcrResult<Vec<image::RgbImage>>` - Vector of rectified images or error.
    pub fn apply_batch(&self, output: &ndarray::Array4<f32>) -> OcrResult<Vec<image::RgbImage>> {
        use image::{Rgb, RgbImage};

        let shape = output.shape();
        if shape.len() != 4 {
            return Err(OCRError::InvalidInput {
                message: "Expected 4D tensor [batch, channels, height, width]".to_string(),
            });
        }

        let batch_size = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];

        if channels != 3 {
            return Err(OCRError::InvalidInput {
                message: "Expected 3 channels (RGB)".to_string(),
            });
        }

        let mut images = Vec::with_capacity(batch_size);

        let scale = self.scale;
        let plane = height * width;

        for b in 0..batch_size {
            let mut img = RgbImage::new(width as u32, height as u32);

            // Model outputs are in BGR order; convert back to RGB. For the
            // common standard-layout tensor the three channel planes are
            // contiguous, so the SIMD kernel scales + clamps them straight into
            // the image buffer (no per-pixel `put_pixel`/strided indexing).
            // Fall back to indexed access for non-standard layouts.
            match output.as_slice() {
                Some(buf) => {
                    let base = b * 3 * plane;
                    crate::processors::simd::scale_clamp_bgr_planes_to_rgb(
                        &buf[base..base + plane],
                        &buf[base + plane..base + 2 * plane],
                        &buf[base + 2 * plane..base + 3 * plane],
                        scale,
                        img.as_mut(),
                    );
                }
                None => {
                    for y in 0..height {
                        for x in 0..width {
                            let b_val = (output[[b, 0, y, x]] * scale).clamp(0.0, 255.0) as u8;
                            let g_val = (output[[b, 1, y, x]] * scale).clamp(0.0, 255.0) as u8;
                            let r_val = (output[[b, 2, y, x]] * scale).clamp(0.0, 255.0) as u8;
                            img.put_pixel(x as u32, y as u32, Rgb([r_val, g_val, b_val]));
                        }
                    }
                }
            }

            images.push(img);
        }

        Ok(images)
    }
}

impl Default for UVDocPostProcess {
    /// Creates a default UVDocPostProcess with scale factor 1.0.
    fn default() -> Self {
        Self::new(1.0)
    }
}

impl FromStr for UVDocPostProcess {
    type Err = std::num::ParseFloatError;

    /// Creates a UVDocPostProcess from a string representation of the scale factor.
    ///
    /// # Arguments
    ///
    /// * `s` - String representation of the scale factor.
    ///
    /// # Returns
    ///
    /// * `Ok(UVDocPostProcess)` - If the string can be parsed as a float.
    /// * `Err(ParseFloatError)` - If the string cannot be parsed.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let scale = s.parse::<f32>()?;
        Ok(Self::new(scale))
    }
}

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

    #[test]
    fn test_denormalize_coordinates() {
        let postprocessor = UVDocPostProcess::new(100.0);
        let normalized = vec![0.1, 0.2, 0.8, 0.9];
        let pixels = postprocessor.denormalize_coordinates(&normalized);
        assert_eq!(pixels, vec![10.0, 20.0, 80.0, 90.0]);
    }

    #[test]
    fn test_normalize_coordinates() {
        let postprocessor = UVDocPostProcess::new(100.0);
        let pixels = vec![10.0, 20.0, 80.0, 90.0];
        let normalized = postprocessor.normalize_coordinates(&pixels);
        assert_eq!(normalized, vec![0.1, 0.2, 0.8, 0.9]);
    }

    #[test]
    fn test_process_bbox() {
        let postprocessor = UVDocPostProcess::new(100.0);
        let normalized_bbox = [0.1, 0.2, 0.8, 0.9];
        let pixel_bbox = postprocessor.process_bbox(&normalized_bbox);
        assert_eq!(pixel_bbox, [10.0, 20.0, 80.0, 90.0]);
    }

    #[test]
    fn test_process_polygon() {
        let postprocessor = UVDocPostProcess::new(100.0);
        let normalized_polygon = vec![[0.1, 0.2], [0.8, 0.2], [0.8, 0.9], [0.1, 0.9]];
        let pixel_polygon = postprocessor.process_polygon(&normalized_polygon);
        assert_eq!(pixel_polygon[0], [10.0, 20.0]);
        assert_eq!(pixel_polygon[1], [80.0, 20.0]);
    }

    #[test]
    fn test_clamp_coordinates() {
        let postprocessor = UVDocPostProcess::new(1.0);
        let coords = vec![-10.0, 50.0, 150.0];
        let clamped = postprocessor.clamp_coordinates(&coords, 0.0, 100.0);
        assert_eq!(clamped, vec![0.0, 50.0, 100.0]);
    }

    #[test]
    fn test_validate_coordinates() {
        let postprocessor = UVDocPostProcess::new(1.0);
        let valid_coords = vec![10.0, 50.0, 90.0];
        let invalid_coords = vec![10.0, 150.0, 90.0];

        assert!(postprocessor.validate_coordinates(&valid_coords, 0.0, 100.0));
        assert!(!postprocessor.validate_coordinates(&invalid_coords, 0.0, 100.0));
    }

    #[test]
    fn test_round_coordinates() {
        let postprocessor = UVDocPostProcess::new(1.0);
        let coords = vec![10.3, 20.7, 30.5];
        let rounded = postprocessor.round_coordinates(&coords);
        assert_eq!(rounded, vec![10, 21, 31]);
    }

    #[test]
    fn test_from_str() -> Result<(), std::num::ParseFloatError> {
        let postprocessor: UVDocPostProcess = "2.5".parse()?;
        assert_eq!(postprocessor.scale(), 2.5);

        assert!("invalid".parse::<UVDocPostProcess>().is_err());
        Ok(())
    }

    #[test]
    fn test_zero_scale_normalize() {
        let postprocessor = UVDocPostProcess::new(0.0);
        let pixels = vec![10.0, 20.0];
        let normalized = postprocessor.normalize_coordinates(&pixels);
        assert_eq!(normalized, vec![0.0, 0.0]);
    }
}