oar_ocr_core/utils/
image.rs

1//! Utility functions for image processing.
2//!
3//! This module provides functions for loading, converting, and manipulating images
4//! in the OCR pipeline. It includes functions for converting between different
5//! image formats, loading single or batch images from files, creating images
6//! from raw data, and resize-and-pad operations.
7
8use crate::core::OCRError;
9use crate::core::errors::ImageProcessError;
10use image::{DynamicImage, GrayImage, ImageBuffer, ImageError, ImageReader, RgbImage};
11use std::fs::File;
12use std::io::BufReader;
13use std::path::Path;
14
15/// Converts a DynamicImage to an RgbImage.
16///
17/// This function takes a DynamicImage (which can be in any format) and converts
18/// it to an RgbImage (8-bit RGB format).
19///
20/// # Arguments
21///
22/// * `img` - The DynamicImage to convert
23///
24/// # Returns
25///
26/// * `RgbImage` - The converted RGB image
27pub fn dynamic_to_rgb(img: DynamicImage) -> RgbImage {
28    img.to_rgb8()
29}
30
31/// Converts a DynamicImage to a GrayImage.
32///
33/// This function takes a DynamicImage (which can be in any format) and converts
34/// it to a GrayImage (8-bit grayscale format).
35///
36/// # Arguments
37///
38/// * `img` - The DynamicImage to convert
39///
40/// # Returns
41///
42/// * `GrayImage` - The converted grayscale image
43pub fn dynamic_to_gray(img: DynamicImage) -> GrayImage {
44    img.to_luma8()
45}
46
47/// Loads an image from the given bytes and converts it to RgbImage.
48///
49/// This function decodes an image from a byte slice and converts it
50/// to an RgbImage. It handles any image format supported by the image crate.
51///
52/// # Arguments
53///
54/// * `bytes` - A byte slice containing the encoded image data
55///
56/// # Returns
57///
58/// * `Ok(RgbImage)` - The decoded and converted RGB image
59/// * `Err(OCRError)` - An error if the image could not be decoded or converted
60///
61/// # Errors
62///
63/// This function will return an `OCRError::ImageLoad` error if the image cannot
64/// be decoded from the provided bytes, or if there is an error during conversion.
65pub fn load_image_from_memory(bytes: &[u8]) -> Result<RgbImage, OCRError> {
66    let img = image::load_from_memory(bytes).map_err(OCRError::ImageLoad)?;
67    Ok(dynamic_to_rgb(img))
68}
69
70/// Loads an image from a file path and converts it to RgbImage.
71///
72/// This function opens an image from the specified file path and converts it
73/// to an RgbImage. It handles any image format supported by the image crate.
74///
75/// # Arguments
76///
77/// * `path` - A reference to the path of the image file to load
78///
79/// # Returns
80///
81/// * `Ok(RgbImage)` - The loaded and converted RGB image
82/// * `Err(OCRError)` - An error if the image could not be loaded or converted
83///
84/// # Errors
85///
86/// This function will return an `OCRError::ImageLoad` error if the image cannot
87/// be loaded from the specified path, or if there is an error during conversion.
88pub fn load_image<P: AsRef<Path>>(path: P) -> Result<RgbImage, OCRError> {
89    let img = open_image_any_format(path.as_ref()).map_err(OCRError::ImageLoad)?;
90    Ok(dynamic_to_rgb(img))
91}
92
93fn open_image_any_format(path: &Path) -> Result<DynamicImage, ImageError> {
94    let file = File::open(path)?;
95    let reader = BufReader::new(file);
96    let reader = ImageReader::new(reader).with_guessed_format()?;
97    reader.decode()
98}
99
100/// Creates an RgbImage from raw pixel data.
101///
102/// This function creates an RgbImage from raw pixel data. The data must be
103/// in RGB format (3 bytes per pixel) and the length must match the specified
104/// width and height.
105///
106/// # Arguments
107///
108/// * `width` - The width of the image in pixels
109/// * `height` - The height of the image in pixels
110/// * `data` - A vector containing the raw pixel data (RGB format)
111///
112/// # Returns
113///
114/// * `Some(RgbImage)` - The created RGB image if the data is valid
115/// * `None` - If the data length doesn't match the specified dimensions
116pub fn create_rgb_image(width: u32, height: u32, data: Vec<u8>) -> Option<RgbImage> {
117    if data.len() != (width * height * 3) as usize {
118        return None;
119    }
120
121    ImageBuffer::from_raw(width, height, data)
122}
123
124/// Checks if the given image size is valid (non-zero dimensions).
125pub fn check_image_size(size: &[u32; 2]) -> Result<(), ImageProcessError> {
126    if size[0] == 0 || size[1] == 0 {
127        return Err(ImageProcessError::InvalidCropSize);
128    }
129    Ok(())
130}
131
132/// Extracts a rectangular region from an RGB image.
133pub fn slice_image(
134    img: &RgbImage,
135    coords: (u32, u32, u32, u32),
136) -> Result<RgbImage, ImageProcessError> {
137    let (x1, y1, x2, y2) = coords;
138    let (img_width, img_height) = img.dimensions();
139
140    if x1 >= x2 || y1 >= y2 {
141        return Err(ImageProcessError::InvalidCropCoordinates);
142    }
143
144    if x2 > img_width || y2 > img_height {
145        return Err(ImageProcessError::CropOutOfBounds);
146    }
147
148    let crop_width = x2 - x1;
149    let crop_height = y2 - y1;
150
151    Ok(image::imageops::crop_imm(img, x1, y1, crop_width, crop_height).to_image())
152}
153
154/// Extracts a rectangular region from a grayscale image.
155pub fn slice_gray_image(
156    img: &GrayImage,
157    coords: (u32, u32, u32, u32),
158) -> Result<GrayImage, ImageProcessError> {
159    let (x1, y1, x2, y2) = coords;
160    let (img_width, img_height) = img.dimensions();
161
162    if x1 >= x2 || y1 >= y2 {
163        return Err(ImageProcessError::InvalidCropCoordinates);
164    }
165
166    if x2 > img_width || y2 > img_height {
167        return Err(ImageProcessError::CropOutOfBounds);
168    }
169
170    let crop_width = x2 - x1;
171    let crop_height = y2 - y1;
172
173    Ok(image::imageops::crop_imm(img, x1, y1, crop_width, crop_height).to_image())
174}
175
176/// Calculates centered crop coordinates for a target size.
177pub fn calculate_center_crop_coords(
178    img_width: u32,
179    img_height: u32,
180    crop_width: u32,
181    crop_height: u32,
182) -> Result<(u32, u32), ImageProcessError> {
183    if crop_width > img_width || crop_height > img_height {
184        return Err(ImageProcessError::CropSizeTooLarge);
185    }
186
187    let x = (img_width - crop_width) / 2;
188    let y = (img_height - crop_height) / 2;
189
190    Ok((x, y))
191}
192
193/// Validates that crop coordinates stay within image bounds.
194pub fn validate_crop_bounds(
195    img_width: u32,
196    img_height: u32,
197    x: u32,
198    y: u32,
199    crop_width: u32,
200    crop_height: u32,
201) -> Result<(), ImageProcessError> {
202    if x + crop_width > img_width || y + crop_height > img_height {
203        return Err(ImageProcessError::CropOutOfBounds);
204    }
205    Ok(())
206}
207
208/// Resizes an RGB image to the target dimensions using Lanczos3 filtering.
209///
210/// # Errors
211///
212/// Returns `ImageProcessError::InvalidCropSize` if width or height is 0.
213pub fn resize_image(
214    img: &RgbImage,
215    width: u32,
216    height: u32,
217) -> Result<RgbImage, ImageProcessError> {
218    if width == 0 || height == 0 {
219        return Err(ImageProcessError::InvalidCropSize);
220    }
221    Ok(image::imageops::resize(
222        img,
223        width,
224        height,
225        image::imageops::FilterType::Lanczos3,
226    ))
227}
228
229/// Resizes a grayscale image to the target dimensions using Lanczos3 filtering.
230///
231/// # Errors
232///
233/// Returns `ImageProcessError::InvalidCropSize` if width or height is 0.
234pub fn resize_gray_image(
235    img: &GrayImage,
236    width: u32,
237    height: u32,
238) -> Result<GrayImage, ImageProcessError> {
239    if width == 0 || height == 0 {
240        return Err(ImageProcessError::InvalidCropSize);
241    }
242    Ok(image::imageops::resize(
243        img,
244        width,
245        height,
246        image::imageops::FilterType::Lanczos3,
247    ))
248}
249
250/// Converts an RGB image to grayscale.
251pub fn rgb_to_grayscale(img: &RgbImage) -> GrayImage {
252    image::imageops::grayscale(img)
253}
254
255/// Pads an image to the specified dimensions with a fill color.
256pub fn pad_image(
257    img: &RgbImage,
258    target_width: u32,
259    target_height: u32,
260    fill_color: [u8; 3],
261) -> Result<RgbImage, ImageProcessError> {
262    let (src_width, src_height) = img.dimensions();
263
264    if target_width < src_width || target_height < src_height {
265        return Err(ImageProcessError::InvalidCropSize);
266    }
267
268    if target_width == src_width && target_height == src_height {
269        return Ok(img.clone());
270    }
271
272    let mut padded = RgbImage::from_pixel(target_width, target_height, image::Rgb(fill_color));
273    let x_offset = (target_width - src_width) / 2;
274    let y_offset = (target_height - src_height) / 2;
275    image::imageops::overlay(&mut padded, img, x_offset as i64, y_offset as i64);
276
277    Ok(padded)
278}
279
280/// Loads a batch of images from file paths.
281///
282/// This function loads multiple images from the specified file paths and
283/// converts them to RgbImages. It uses parallel processing when the number
284/// of images exceeds the default parallel threshold.
285///
286/// # Arguments
287///
288/// * `paths` - A slice of paths to the image files to load
289///
290/// # Returns
291///
292/// * `Ok(Vec<RgbImage>)` - A vector of loaded RGB images
293/// * `Err(OCRError)` - An error if any image could not be loaded
294///
295/// # Errors
296///
297/// This function will return an `OCRError` if any image cannot be loaded
298/// from its specified path.
299pub fn load_images<P: AsRef<std::path::Path> + Send + Sync>(
300    paths: &[P],
301) -> Result<Vec<RgbImage>, OCRError> {
302    load_images_batch_with_threshold(paths, None)
303}
304
305/// Loads a batch of images from file paths with a custom parallel threshold.
306///
307/// This function loads multiple images from the specified file paths and
308/// converts them to RgbImages. It uses parallel processing when the number
309/// of images exceeds the specified threshold, or the default threshold if
310/// none is provided.
311///
312/// # Arguments
313///
314/// * `paths` - A slice of paths to the image files to load
315/// * `parallel_threshold` - An optional threshold for parallel processing.
316///   If `None`, the default threshold from `DEFAULT_PARALLEL_THRESHOLD` is used.
317///
318/// # Returns
319///
320/// * `Ok(Vec<RgbImage>)` - A vector of loaded RGB images
321/// * `Err(OCRError)` - An error if any image could not be loaded
322///
323/// # Errors
324///
325/// This function will return an `OCRError` if any image cannot be loaded
326/// from its specified path.
327pub fn load_images_batch_with_threshold<P: AsRef<std::path::Path> + Send + Sync>(
328    paths: &[P],
329    parallel_threshold: Option<usize>,
330) -> Result<Vec<RgbImage>, OCRError> {
331    use crate::core::constants::DEFAULT_PARALLEL_THRESHOLD;
332
333    let threshold = parallel_threshold.unwrap_or(DEFAULT_PARALLEL_THRESHOLD);
334
335    if paths.len() > threshold {
336        use rayon::prelude::*;
337        paths.par_iter().map(|p| load_image(p.as_ref())).collect()
338    } else {
339        paths.iter().map(|p| load_image(p.as_ref())).collect()
340    }
341}
342
343/// Load multiple images from file paths using centralized parallel policy.
344///
345/// This function loads images from the provided file paths using the utility threshold
346/// from the centralized ParallelPolicy. If the number of paths exceeds the threshold,
347/// the loading is performed in parallel using rayon. Otherwise, images are loaded
348/// sequentially.
349///
350/// # Arguments
351///
352/// * `paths` - A slice of paths to image files
353/// * `policy` - The parallel policy containing the utility threshold
354///
355/// # Returns
356///
357/// A Result containing a vector of loaded RgbImages, or an OCRError if any image fails to load.
358///
359/// # Errors
360///
361/// This function will return an `OCRError` if any image cannot be loaded
362/// from its specified path.
363pub fn load_images_batch_with_policy<P: AsRef<std::path::Path> + Send + Sync>(
364    paths: &[P],
365    policy: &crate::core::config::ParallelPolicy,
366) -> Result<Vec<RgbImage>, OCRError> {
367    if paths.len() > policy.utility_threshold {
368        use rayon::prelude::*;
369        paths.par_iter().map(|p| load_image(p.as_ref())).collect()
370    } else {
371        paths.iter().map(|p| load_image(p.as_ref())).collect()
372    }
373}
374
375/// Padding strategy for resize-and-pad operations.
376#[derive(Debug, Clone, Copy, PartialEq, Default)]
377pub enum PaddingStrategy {
378    /// Pad with a solid color
379    SolidColor([u8; 3]),
380    /// Pad with black (equivalent to SolidColor([0, 0, 0]))
381    #[default]
382    Black,
383    /// Left-align the resized image (no centering)
384    LeftAlign([u8; 3]),
385}
386
387/// Configuration for resize-and-pad operations.
388#[derive(Debug, Clone)]
389pub struct ResizePadConfig {
390    /// Target dimensions (width, height)
391    pub target_dims: (u32, u32),
392    /// Padding strategy to use
393    pub padding_strategy: PaddingStrategy,
394    /// Filter type for resizing
395    pub filter_type: image::imageops::FilterType,
396}
397
398impl ResizePadConfig {
399    /// Create a new resize-pad configuration.
400    pub fn new(target_dims: (u32, u32)) -> Self {
401        Self {
402            target_dims,
403            padding_strategy: PaddingStrategy::default(),
404            filter_type: image::imageops::FilterType::Triangle,
405        }
406    }
407
408    /// Set the padding strategy.
409    pub fn with_padding_strategy(mut self, strategy: PaddingStrategy) -> Self {
410        self.padding_strategy = strategy;
411        self
412    }
413
414    /// Set the filter type for resizing.
415    pub fn with_filter_type(mut self, filter_type: image::imageops::FilterType) -> Self {
416        self.filter_type = filter_type;
417        self
418    }
419}
420
421/// Resize an image to fit within target dimensions while maintaining aspect ratio,
422/// then pad to exact target dimensions.
423///
424/// This function provides a unified approach to resize-and-pad operations that
425/// can replace the duplicated logic found in various processors.
426///
427/// # Arguments
428///
429/// * `image` - The input RGB image to resize and pad
430/// * `config` - Configuration for the resize-and-pad operation
431///
432/// # Returns
433///
434/// A resized and padded RGB image with exact target dimensions.
435///
436/// # Errors
437///
438/// Returns `ImageProcessError::InvalidCropSize` if target dimensions are 0.
439pub fn resize_and_pad(
440    image: &RgbImage,
441    config: &ResizePadConfig,
442) -> Result<RgbImage, ImageProcessError> {
443    let (target_width, target_height) = config.target_dims;
444
445    if target_width == 0 || target_height == 0 {
446        return Err(ImageProcessError::InvalidCropSize);
447    }
448
449    let (orig_width, orig_height) = image.dimensions();
450
451    // Calculate scaling factor to fit within target dimensions while maintaining aspect ratio
452    let scale_w = target_width as f32 / orig_width as f32;
453    let scale_h = target_height as f32 / orig_height as f32;
454    let scale = scale_w.min(scale_h);
455
456    // Calculate new dimensions
457    let new_width = (orig_width as f32 * scale) as u32;
458    let new_height = (orig_height as f32 * scale) as u32;
459
460    // Resize the image
461    let resized = image::imageops::resize(image, new_width, new_height, config.filter_type);
462
463    // Create padded image with target dimensions
464    let padding_color = match config.padding_strategy {
465        PaddingStrategy::SolidColor(color) => color,
466        PaddingStrategy::Black => [0, 0, 0],
467        PaddingStrategy::LeftAlign(color) => color,
468    };
469    let padding_rgb = image::Rgb(padding_color);
470    let mut padded = ImageBuffer::from_pixel(target_width, target_height, padding_rgb);
471
472    // Calculate padding offsets
473    let (pad_x, pad_y) = match config.padding_strategy {
474        PaddingStrategy::LeftAlign(_) => (0, 0),
475        _ => {
476            // Center the image
477            let pad_x = (target_width - new_width) / 2;
478            let pad_y = (target_height - new_height) / 2;
479            (pad_x, pad_y)
480        }
481    };
482
483    // Copy resized image to padded image using efficient overlay
484    image::imageops::overlay(&mut padded, &resized, pad_x as i64, pad_y as i64);
485
486    Ok(padded)
487}
488
489/// Configuration for OCR-style resize-and-pad operations with width constraints.
490#[derive(Debug, Clone)]
491pub struct OCRResizePadConfig {
492    /// Target height
493    pub target_height: u32,
494    /// Maximum allowed width
495    pub max_width: u32,
496    /// Padding strategy to use
497    pub padding_strategy: PaddingStrategy,
498    /// Filter type for resizing
499    pub filter_type: image::imageops::FilterType,
500}
501
502impl OCRResizePadConfig {
503    /// Create a new OCR resize-pad configuration.
504    ///
505    /// Uses Triangle (bilinear) interpolation to match OpenCV's cv2.resize default behavior.
506    pub fn new(target_height: u32, max_width: u32) -> Self {
507        Self {
508            target_height,
509            max_width,
510            padding_strategy: PaddingStrategy::default(),
511            // Use Triangle (bilinear) to match cv2.resize INTER_LINEAR
512            filter_type: image::imageops::FilterType::Triangle,
513        }
514    }
515
516    /// Set the padding strategy.
517    pub fn with_padding_strategy(mut self, strategy: PaddingStrategy) -> Self {
518        self.padding_strategy = strategy;
519        self
520    }
521
522    /// Set the filter type for resizing.
523    pub fn with_filter_type(mut self, filter_type: image::imageops::FilterType) -> Self {
524        self.filter_type = filter_type;
525        self
526    }
527}
528
529/// Resize an image for OCR processing with width constraints and padding.
530///
531/// This function handles the specific resize-and-pad logic used in OCR processing,
532/// where images are resized to a fixed height while maintaining aspect ratio,
533/// with a maximum width constraint, and then padded to a target width.
534///
535/// # Arguments
536///
537/// * `image` - The input RGB image to resize and pad
538/// * `config` - Configuration for the OCR resize-and-pad operation
539/// * `target_width_ratio` - Optional ratio to calculate target width from height.
540///   If None, uses the image's original aspect ratio.
541///
542/// # Returns
543///
544/// A tuple containing:
545/// - The resized and padded RGB image
546/// - The actual width used for the padded image
547///
548/// # Errors
549///
550/// Returns `ImageProcessError::InvalidCropSize` if target height is 0.
551pub fn ocr_resize_and_pad(
552    image: &RgbImage,
553    config: &OCRResizePadConfig,
554    target_width_ratio: Option<f32>,
555) -> Result<(RgbImage, u32), ImageProcessError> {
556    if config.target_height == 0 {
557        return Err(ImageProcessError::InvalidCropSize);
558    }
559
560    let (original_w, original_h) = image.dimensions();
561    let original_ratio = original_w as f32 / original_h as f32;
562
563    // Calculate target width based on ratio or original aspect ratio
564    let mut target_w = if let Some(ratio) = target_width_ratio {
565        (config.target_height as f32 * ratio) as u32
566    } else {
567        (config.target_height as f32 * original_ratio).ceil() as u32
568    };
569
570    // Apply maximum width constraint
571    let resized_w = if target_w > config.max_width {
572        target_w = config.max_width;
573        config.max_width
574    } else {
575        // Calculate actual resized width based on aspect ratio
576        let ratio = original_w as f32 / original_h as f32;
577        if (config.target_height as f32 * ratio).ceil() as u32 > target_w {
578            target_w
579        } else {
580            (config.target_height as f32 * ratio).ceil() as u32
581        }
582    };
583
584    // Resize the image
585    let resized_image =
586        image::imageops::resize(image, resized_w, config.target_height, config.filter_type);
587
588    // Create padded image with target dimensions
589    let padding_color = match config.padding_strategy {
590        PaddingStrategy::SolidColor(color) => color,
591        PaddingStrategy::Black => [0, 0, 0],
592        PaddingStrategy::LeftAlign(color) => color,
593    };
594    let padding_rgb = image::Rgb(padding_color);
595    let mut padded_image = ImageBuffer::from_pixel(target_w, config.target_height, padding_rgb);
596
597    // Copy resized image to padded image (left-aligned for OCR)
598    image::imageops::overlay(&mut padded_image, &resized_image, 0, 0);
599
600    Ok((padded_image, target_w))
601}
602
603/// Resizes a batch of images to the specified dimensions.
604///
605/// This function provides a unified approach to batch image resizing that can replace
606/// duplicated resize loops found in various predictors. It supports both functional
607/// and imperative styles and can optionally apply post-processing operations.
608///
609/// # Arguments
610///
611/// * `images` - A slice of RGB images to resize
612/// * `target_width` - Target width for all images
613/// * `target_height` - Target height for all images
614/// * `filter_type` - The filter type to use for resizing (defaults to Lanczos3 if None)
615///
616/// # Returns
617///
618/// A vector of resized RGB images.
619///
620/// # Example
621///
622/// ```rust,no_run
623/// use oar_ocr_core::utils::resize_images_batch;
624/// use image::RgbImage;
625///
626/// let images = vec![RgbImage::new(100, 100), RgbImage::new(200, 150)];
627/// let resized = resize_images_batch(&images, 224, 224, None);
628/// assert_eq!(resized.len(), 2);
629/// assert_eq!(resized[0].dimensions(), (224, 224));
630/// ```
631pub fn resize_images_batch(
632    images: &[RgbImage],
633    target_width: u32,
634    target_height: u32,
635    filter_type: Option<image::imageops::FilterType>,
636) -> Vec<RgbImage> {
637    let filter = filter_type.unwrap_or(image::imageops::FilterType::Lanczos3);
638
639    images
640        .iter()
641        .map(|img| image::imageops::resize(img, target_width, target_height, filter))
642        .collect()
643}
644
645/// Resizes a batch of images and converts them to DynamicImage format.
646///
647/// This function combines batch resizing with conversion to DynamicImage format,
648/// which is commonly needed in OCR preprocessing pipelines.
649///
650/// # Arguments
651///
652/// * `images` - A slice of RGB images to resize
653/// * `target_width` - Target width for all images
654/// * `target_height` - Target height for all images
655/// * `filter_type` - The filter type to use for resizing (defaults to Lanczos3 if None)
656///
657/// # Returns
658///
659/// A vector of resized images as DynamicImage instances.
660pub fn resize_images_batch_to_dynamic(
661    images: &[RgbImage],
662    target_width: u32,
663    target_height: u32,
664    filter_type: Option<image::imageops::FilterType>,
665) -> Vec<DynamicImage> {
666    let filter = filter_type.unwrap_or(image::imageops::FilterType::Lanczos3);
667
668    images
669        .iter()
670        .map(|img| {
671            let resized = image::imageops::resize(img, target_width, target_height, filter);
672            DynamicImage::ImageRgb8(resized)
673        })
674        .collect()
675}
676
677/// Masks a rectangular region in an RGB image with a solid color.
678///
679/// This function fills the specified rectangular region with a solid color,
680/// which is useful for masking formula regions before text detection to prevent
681/// formulas from being incorrectly detected as text (as done in PP-StructureV3).
682///
683/// # Arguments
684///
685/// * `image` - A mutable reference to the RGB image to mask
686/// * `x1` - Left coordinate of the region
687/// * `y1` - Top coordinate of the region
688/// * `x2` - Right coordinate of the region
689/// * `y2` - Bottom coordinate of the region
690/// * `fill_color` - The color to fill the masked region with (default: white [255, 255, 255])
691///
692/// # Returns
693///
694/// Returns `Ok(())` if masking succeeds, or an error if coordinates are invalid.
695///
696/// # Example
697///
698/// ```rust,no_run
699/// use oar_ocr_core::utils::mask_region;
700/// use image::RgbImage;
701///
702/// # fn main() -> Result<(), oar_ocr_core::core::errors::ImageProcessError> {
703/// let mut image = RgbImage::new(100, 100);
704/// // Mask a formula region from (10, 10) to (50, 30) with white
705/// mask_region(&mut image, 10, 10, 50, 30, [255, 255, 255])?;
706/// # Ok(())
707/// # }
708/// ```
709pub fn mask_region(
710    image: &mut RgbImage,
711    x1: u32,
712    y1: u32,
713    x2: u32,
714    y2: u32,
715    fill_color: [u8; 3],
716) -> Result<(), ImageProcessError> {
717    let (img_width, img_height) = image.dimensions();
718
719    // Clamp coordinates to image bounds
720    let x1 = x1.min(img_width);
721    let y1 = y1.min(img_height);
722    let x2 = x2.min(img_width);
723    let y2 = y2.min(img_height);
724
725    if x1 >= x2 || y1 >= y2 {
726        return Err(ImageProcessError::InvalidCropCoordinates);
727    }
728
729    let rgb = image::Rgb(fill_color);
730    for y in y1..y2 {
731        for x in x1..x2 {
732            image.put_pixel(x, y, rgb);
733        }
734    }
735
736    Ok(())
737}
738
739/// Masks multiple bounding box regions in an RGB image.
740///
741/// This function masks multiple regions by filling them with a solid color.
742/// It is useful for batch masking multiple formula or other regions before
743/// text detection (as done in PP-StructureV3).
744///
745/// # Arguments
746///
747/// * `image` - A mutable reference to the RGB image to mask
748/// * `bboxes` - A slice of bounding boxes to mask. Each bbox should provide
749///   `x_min()`, `y_min()`, `x_max()`, `y_max()` methods.
750/// * `fill_color` - The color to fill the masked regions with
751///
752/// # Example
753///
754/// ```rust,no_run
755/// use oar_ocr_core::utils::mask_regions;
756/// use oar_ocr_core::processors::BoundingBox;
757/// use image::RgbImage;
758///
759/// let mut image = RgbImage::new(100, 100);
760/// let bboxes = vec![
761///     BoundingBox::from_coords(10.0, 10.0, 30.0, 30.0),
762///     BoundingBox::from_coords(50.0, 50.0, 70.0, 70.0),
763/// ];
764/// mask_regions(&mut image, &bboxes, [255, 255, 255]);
765/// ```
766pub fn mask_regions(
767    image: &mut RgbImage,
768    bboxes: &[crate::processors::BoundingBox],
769    fill_color: [u8; 3],
770) {
771    for bbox in bboxes {
772        let x1 = bbox.x_min() as u32;
773        let y1 = bbox.y_min() as u32;
774        let x2 = bbox.x_max() as u32;
775        let y2 = bbox.y_max() as u32;
776
777        // Ignore errors for individual regions (they might be out of bounds)
778        let _ = mask_region(image, x1, y1, x2, y2, fill_color);
779    }
780}
781
782#[cfg(test)]
783mod tests {
784    use super::*;
785    use ::image::{GenericImageView, GrayImage, ImageBuffer, Rgb, RgbImage};
786
787    fn create_test_image(width: u32, height: u32, color: [u8; 3]) -> RgbImage {
788        ImageBuffer::from_pixel(width, height, Rgb(color))
789    }
790
791    #[test]
792    fn basic_size_checks() {
793        assert!(check_image_size(&[100, 100]).is_ok());
794        assert!(check_image_size(&[0, 50]).is_err());
795    }
796
797    #[test]
798    fn slice_rgb_image_region() -> Result<(), ImageProcessError> {
799        let img = RgbImage::from_pixel(10, 10, Rgb([255, 0, 0]));
800        let cropped = slice_image(&img, (2, 2, 6, 6))?;
801        assert_eq!(cropped.dimensions(), (4, 4));
802        assert!(slice_image(&img, (6, 6, 2, 2)).is_err());
803        Ok(())
804    }
805
806    #[test]
807    fn slice_gray_image_region() -> Result<(), ImageProcessError> {
808        let img = GrayImage::from_pixel(10, 10, image::Luma([128]));
809        let cropped = slice_gray_image(&img, (1, 1, 5, 5))?;
810        assert_eq!(cropped.dimensions(), (4, 4));
811        Ok(())
812    }
813
814    #[test]
815    fn center_crop_coordinates() -> Result<(), ImageProcessError> {
816        let coords = calculate_center_crop_coords(100, 60, 40, 20)?;
817        assert_eq!(coords, (30, 20));
818        assert!(calculate_center_crop_coords(20, 20, 40, 10).is_err());
819        Ok(())
820    }
821
822    #[test]
823    fn crop_bounds_validation() {
824        assert!(validate_crop_bounds(100, 80, 10, 10, 40, 40).is_ok());
825        assert!(validate_crop_bounds(100, 80, 70, 10, 40, 40).is_err());
826    }
827
828    #[test]
829    fn pad_image_to_target() -> Result<(), ImageProcessError> {
830        let img = RgbImage::from_pixel(20, 20, Rgb([10, 20, 30]));
831        let padded = pad_image(&img, 40, 40, [0, 0, 0])?;
832        assert_eq!(padded.dimensions(), (40, 40));
833        assert!(pad_image(&img, 10, 10, [0, 0, 0]).is_err());
834        Ok(())
835    }
836
837    #[test]
838    fn test_resize_and_pad_with_custom_padding() -> Result<(), ImageProcessError> {
839        let image = create_test_image(50, 100, [255, 0, 0]); // 1:2 aspect ratio (tall)
840        let config = ResizePadConfig::new((80, 80))
841            .with_padding_strategy(PaddingStrategy::SolidColor([0, 255, 0])); // Green padding
842
843        let result = resize_and_pad(&image, &config)?;
844
845        assert_eq!(result.dimensions(), (80, 80));
846
847        // The resized image should be 40x80 (maintaining 1:2 ratio), centered in 80x80
848        // So there should be 20 pixels of padding on left and right
849        let center_pixel = result.get_pixel(40, 40); // Center of image
850        assert_eq!(*center_pixel, Rgb([255, 0, 0])); // Should be red (original image)
851
852        let left_padding = result.get_pixel(10, 40); // Left padding area
853        assert_eq!(*left_padding, Rgb([0, 255, 0])); // Should be green (custom padding)
854        Ok(())
855    }
856
857    #[test]
858    fn test_resize_and_pad_left_align() -> Result<(), ImageProcessError> {
859        let image = create_test_image(50, 100, [0, 0, 255]); // 1:2 aspect ratio (tall)
860        let config = ResizePadConfig::new((80, 80))
861            .with_padding_strategy(PaddingStrategy::LeftAlign([255, 255, 0])); // Yellow padding, left-aligned
862
863        let result = resize_and_pad(&image, &config)?;
864
865        assert_eq!(result.dimensions(), (80, 80));
866
867        // The resized image should be 40x80, left-aligned in 80x80
868        let left_edge_pixel = result.get_pixel(20, 40); // Should be in the resized image
869        assert_eq!(*left_edge_pixel, Rgb([0, 0, 255])); // Should be blue (original image)
870
871        let right_padding = result.get_pixel(60, 40); // Right padding area
872        assert_eq!(*right_padding, Rgb([255, 255, 0])); // Should be yellow (padding)
873        Ok(())
874    }
875
876    #[test]
877    fn test_resize_images_batch() {
878        // Create test images with different sizes
879        let img1 = create_test_image(100, 50, [255, 0, 0]); // Red
880        let img2 = create_test_image(200, 100, [0, 255, 0]); // Green
881        let images = vec![img1, img2];
882
883        // Resize batch to 64x64
884        let resized = resize_images_batch(&images, 64, 64, None);
885
886        assert_eq!(resized.len(), 2);
887        assert_eq!(resized[0].dimensions(), (64, 64));
888        assert_eq!(resized[1].dimensions(), (64, 64));
889
890        // Check that the colors are preserved (approximately)
891        let pixel1 = resized[0].get_pixel(32, 32);
892        let pixel2 = resized[1].get_pixel(32, 32);
893
894        // Red image should still be predominantly red
895        assert!(pixel1[0] > pixel1[1] && pixel1[0] > pixel1[2]);
896        // Green image should still be predominantly green
897        assert!(pixel2[1] > pixel2[0] && pixel2[1] > pixel2[2]);
898    }
899
900    #[test]
901    fn test_resize_images_batch_to_dynamic() {
902        // Create test images
903        let img1 = create_test_image(100, 50, [255, 0, 0]);
904        let img2 = create_test_image(200, 100, [0, 255, 0]);
905        let images = vec![img1, img2];
906
907        // Resize batch to 32x32 and convert to DynamicImage
908        let resized = resize_images_batch_to_dynamic(&images, 32, 32, None);
909
910        assert_eq!(resized.len(), 2);
911
912        // Check that they are DynamicImage::ImageRgb8 variants
913        for dynamic_img in &resized {
914            assert_eq!(dynamic_img.dimensions(), (32, 32));
915            assert!(
916                matches!(dynamic_img, DynamicImage::ImageRgb8(_)),
917                "Expected ImageRgb8 variant"
918            );
919        }
920    }
921
922    #[test]
923    fn test_resize_images_batch_empty() {
924        let images: Vec<RgbImage> = vec![];
925        let resized = resize_images_batch(&images, 64, 64, None);
926        assert!(resized.is_empty());
927    }
928
929    #[test]
930    fn test_resize_images_batch_custom_filter() {
931        let img = create_test_image(100, 100, [128, 128, 128]);
932        let images = vec![img];
933
934        // Test with different filter types
935        let resized_lanczos =
936            resize_images_batch(&images, 50, 50, Some(image::imageops::FilterType::Lanczos3));
937        let resized_nearest =
938            resize_images_batch(&images, 50, 50, Some(image::imageops::FilterType::Nearest));
939
940        assert_eq!(resized_lanczos.len(), 1);
941        assert_eq!(resized_nearest.len(), 1);
942        assert_eq!(resized_lanczos[0].dimensions(), (50, 50));
943        assert_eq!(resized_nearest[0].dimensions(), (50, 50));
944    }
945
946    #[test]
947    fn test_ocr_resize_and_pad_with_max_width_constraint() -> Result<(), ImageProcessError> {
948        let image = create_test_image(400, 100, [200, 100, 50]); // 4:1 aspect ratio
949        let config = OCRResizePadConfig::new(32, 100); // Height 32, max width 100
950
951        let (result, actual_width) = ocr_resize_and_pad(&image, &config, None)?;
952
953        assert_eq!(result.height(), 32);
954        assert_eq!(actual_width, 100); // Should be constrained to max width
955        assert_eq!(result.width(), 100);
956
957        // Check that the image is left-aligned
958        let left_pixel = result.get_pixel(0, 16); // Left edge, middle height
959        assert_eq!(*left_pixel, Rgb([200, 100, 50])); // Should be original color
960        Ok(())
961    }
962
963    #[test]
964    fn test_ocr_resize_and_pad_with_target_ratio() -> Result<(), ImageProcessError> {
965        let image = create_test_image(100, 50, [255, 128, 64]); // 2:1 aspect ratio
966        let config = OCRResizePadConfig::new(32, 200); // Height 32, max width 200
967        let target_ratio = 3.0; // Force 3:1 ratio
968
969        let (result, actual_width) = ocr_resize_and_pad(&image, &config, Some(target_ratio))?;
970
971        assert_eq!(result.height(), 32);
972        assert_eq!(actual_width, 96); // 32 * 3.0 = 96
973        assert_eq!(result.width(), 96);
974        Ok(())
975    }
976
977    #[test]
978    fn test_resize_pad_config_builder() {
979        let config = ResizePadConfig::new((100, 50))
980            .with_padding_strategy(PaddingStrategy::SolidColor([255, 0, 0]))
981            .with_filter_type(image::imageops::FilterType::Lanczos3);
982
983        assert_eq!(config.target_dims, (100, 50));
984        assert_eq!(
985            config.padding_strategy,
986            PaddingStrategy::SolidColor([255, 0, 0])
987        );
988        assert_eq!(config.filter_type, image::imageops::FilterType::Lanczos3);
989    }
990
991    #[test]
992    fn test_ocr_resize_pad_config_builder() {
993        let config = OCRResizePadConfig::new(64, 320)
994            .with_padding_strategy(PaddingStrategy::SolidColor([100, 100, 100]))
995            .with_filter_type(image::imageops::FilterType::Nearest);
996
997        assert_eq!(config.target_height, 64);
998        assert_eq!(config.max_width, 320);
999        assert_eq!(
1000            config.padding_strategy,
1001            PaddingStrategy::SolidColor([100, 100, 100])
1002        );
1003        assert_eq!(config.filter_type, image::imageops::FilterType::Nearest);
1004    }
1005}
oar_ocr_core/utils/image.rs

oar_ocr_core/utils/
image.rs
