oar_ocr_core/processors/
normalization.rs

1//! Image normalization utilities for OCR processing.
2//!
3//! This module provides functionality to normalize images for OCR processing,
4//! including standard normalization with mean and standard deviation, as well as
5//! specialized normalization for OCR recognition tasks.
6
7use crate::core::OCRError;
8use crate::processors::types::{ColorOrder, TensorLayout};
9use image::{DynamicImage, RgbImage};
10use rayon::prelude::*;
11
12/// Normalizes images for OCR processing.
13///
14/// This struct encapsulates the parameters needed to normalize images,
15/// including scaling factors, mean values, standard deviations, and channel ordering.
16/// It provides methods to apply normalization to single images or batches of images.
17#[derive(Debug)]
18pub struct NormalizeImage {
19    /// Scaling factors for each channel (alpha = scale / std)
20    pub alpha: Vec<f32>,
21    /// Offset values for each channel (beta = -mean / std)
22    pub beta: Vec<f32>,
23    /// Tensor data layout (CHW or HWC)
24    pub order: TensorLayout,
25    /// Color channel order (RGB or BGR)
26    pub color_order: ColorOrder,
27}
28
29impl NormalizeImage {
30    const PARALLEL_NORMALIZE_MIN_BYTES: usize = 1_048_576;
31
32    fn should_parallelize(batch_size: usize, total_output_bytes: usize) -> bool {
33        batch_size > 1 && total_output_bytes > Self::PARALLEL_NORMALIZE_MIN_BYTES
34    }
35
36    fn src_channels(&self) -> [usize; 3] {
37        match self.color_order {
38            ColorOrder::RGB => [0, 1, 2],
39            ColorOrder::BGR => [2, 1, 0],
40        }
41    }
42
43    fn image_len(width: u32, height: u32, channels: usize) -> usize {
44        width as usize * height as usize * channels
45    }
46
47    /// Creates a new NormalizeImage instance with the specified parameters.
48    ///
49    /// # Arguments
50    ///
51    /// * `scale` - Optional scaling factor (defaults to 1.0/255.0)
52    /// * `mean` - Optional mean values for each channel (defaults to [0.485, 0.456, 0.406])
53    /// * `std` - Optional standard deviation values for each channel (defaults to [0.229, 0.224, 0.225])
54    /// * `order` - Optional tensor data layout (defaults to CHW)
55    /// * `color_order` - Optional color channel order (defaults to RGB)
56    ///
57    /// # Returns
58    ///
59    /// A Result containing the new NormalizeImage instance or an OCRError if validation fails.
60    ///
61    /// # Errors
62    ///
63    /// Returns an error if:
64    /// * Scale is less than or equal to 0
65    /// * Mean or std vectors don't have exactly 3 elements
66    /// * Any standard deviation value is less than or equal to 0
67    pub fn new(
68        scale: Option<f32>,
69        mean: Option<Vec<f32>>,
70        std: Option<Vec<f32>>,
71        order: Option<TensorLayout>,
72        color_order: Option<ColorOrder>,
73    ) -> Result<Self, OCRError> {
74        Self::with_color_order(scale, mean, std, order, color_order)
75    }
76
77    /// Creates a new NormalizeImage instance with the specified parameters including color order.
78    ///
79    /// # Arguments
80    ///
81    /// * `scale` - Optional scaling factor (defaults to 1.0/255.0)
82    /// * `mean` - Optional mean values for each channel (defaults to [0.485, 0.456, 0.406])
83    /// * `std` - Optional standard deviation values for each channel (defaults to [0.229, 0.224, 0.225])
84    /// * `order` - Optional tensor data layout (defaults to CHW)
85    /// * `color_order` - Optional color channel order (defaults to RGB)
86    ///
87    /// # Mean/Std Semantics
88    ///
89    /// `mean` and `std` must be provided in the **output channel order** specified by `color_order`.
90    /// For example, if `color_order` is BGR, pass mean/std as `[B_mean, G_mean, R_mean]`.
91    ///
92    /// **Note:** This function does not validate that mean/std values match the specified
93    /// `color_order`. Ensuring consistency is the caller's responsibility. If you have stats
94    /// expressed in RGB order but need BGR output, prefer using
95    /// [`NormalizeImage::with_color_order_from_rgb_stats`] or
96    /// [`NormalizeImage::imagenet_bgr_from_rgb_stats`] which handle the reordering automatically.
97    ///
98    /// # Returns
99    ///
100    /// A Result containing the new NormalizeImage instance or an OCRError if validation fails.
101    pub fn with_color_order(
102        scale: Option<f32>,
103        mean: Option<Vec<f32>>,
104        std: Option<Vec<f32>>,
105        order: Option<TensorLayout>,
106        color_order: Option<ColorOrder>,
107    ) -> Result<Self, OCRError> {
108        let scale = scale.unwrap_or(1.0 / 255.0);
109        let mean = mean.unwrap_or_else(|| vec![0.485, 0.456, 0.406]);
110        let std = std.unwrap_or_else(|| vec![0.229, 0.224, 0.225]);
111        let order = order.unwrap_or(TensorLayout::CHW);
112        let color_order = color_order.unwrap_or_default();
113
114        if scale <= 0.0 {
115            return Err(OCRError::ConfigError {
116                message: "Scale must be greater than 0".to_string(),
117            });
118        }
119
120        if mean.len() != 3 {
121            return Err(OCRError::ConfigError {
122                message: "Mean must have exactly 3 elements (3-channel normalization)".to_string(),
123            });
124        }
125
126        if std.len() != 3 {
127            return Err(OCRError::ConfigError {
128                message: "Std must have exactly 3 elements (3-channel normalization)".to_string(),
129            });
130        }
131
132        for (i, &s) in std.iter().enumerate() {
133            if s <= 0.0 {
134                return Err(OCRError::ConfigError {
135                    message: format!(
136                        "Standard deviation at index {i} must be greater than 0, got {s}"
137                    ),
138                });
139            }
140        }
141
142        let alpha: Vec<f32> = std.iter().map(|s| scale / s).collect();
143        let beta: Vec<f32> = mean.iter().zip(&std).map(|(m, s)| -m / s).collect();
144
145        Ok(Self {
146            alpha,
147            beta,
148            order,
149            color_order,
150        })
151    }
152
153    /// Validates the configuration of the NormalizeImage instance.
154    ///
155    /// # Returns
156    ///
157    /// A Result indicating success or an OCRError if validation fails.
158    ///
159    /// # Errors
160    ///
161    /// Returns an error if:
162    /// * Alpha or beta vectors don't have exactly 3 elements
163    /// * Any alpha or beta value is not finite
164    pub fn validate_config(&self) -> Result<(), OCRError> {
165        if self.alpha.len() != 3 || self.beta.len() != 3 {
166            return Err(OCRError::ConfigError {
167                message: "Alpha and beta must have exactly 3 elements (3-channel normalization)"
168                    .to_string(),
169            });
170        }
171
172        for (i, &alpha) in self.alpha.iter().enumerate() {
173            if !alpha.is_finite() {
174                return Err(OCRError::ConfigError {
175                    message: format!("Alpha value at index {i} is not finite: {alpha}"),
176                });
177            }
178        }
179
180        for (i, &beta) in self.beta.iter().enumerate() {
181            if !beta.is_finite() {
182                return Err(OCRError::ConfigError {
183                    message: format!("Beta value at index {i} is not finite: {beta}"),
184                });
185            }
186        }
187
188        Ok(())
189    }
190
191    /// Creates a NormalizeImage instance with parameters suitable for OCR recognition.
192    ///
193    /// This creates a normalization configuration with:
194    /// * Scale: 2.0/255.0
195    /// * Mean: [1.0, 1.0, 1.0]
196    /// * Std: [1.0, 1.0, 1.0]
197    /// * Order: CHW
198    ///
199    /// # Returns
200    ///
201    /// A Result containing the new NormalizeImage instance or an OCRError.
202    pub fn for_ocr_recognition() -> Result<Self, OCRError> {
203        Self::new(
204            Some(2.0 / 255.0),
205            Some(vec![1.0, 1.0, 1.0]),
206            Some(vec![1.0, 1.0, 1.0]),
207            Some(TensorLayout::CHW),
208            Some(ColorOrder::BGR),
209        )
210    }
211
212    /// Creates an ImageNet-style RGB normalizer (mean/std in RGB order).
213    pub fn imagenet_rgb() -> Result<Self, OCRError> {
214        Self::with_color_order(
215            None,
216            Some(vec![0.485, 0.456, 0.406]),
217            Some(vec![0.229, 0.224, 0.225]),
218            Some(TensorLayout::CHW),
219            Some(ColorOrder::RGB),
220        )
221    }
222
223    /// Creates an ImageNet-style BGR normalizer from RGB stats.
224    ///
225    /// This is useful for PaddlePaddle-exported models that expect BGR input,
226    /// while configuration commonly provides ImageNet mean/std in RGB order.
227    pub fn imagenet_bgr_from_rgb_stats() -> Result<Self, OCRError> {
228        Self::with_color_order(
229            None,
230            Some(vec![0.406, 0.456, 0.485]),
231            Some(vec![0.225, 0.224, 0.229]),
232            Some(TensorLayout::CHW),
233            Some(ColorOrder::BGR),
234        )
235    }
236
237    /// Builds a normalizer for a given output `color_order` using RGB mean/std stats.
238    ///
239    /// Invariant: `mean`/`std` passed to `with_color_order` are interpreted in the output channel
240    /// order (`ColorOrder`). This helper makes the conversion explicit at call sites.
241    pub fn with_color_order_from_rgb_stats(
242        scale: Option<f32>,
243        mean_rgb: Vec<f32>,
244        std_rgb: Vec<f32>,
245        order: Option<TensorLayout>,
246        output_color_order: ColorOrder,
247    ) -> Result<Self, OCRError> {
248        if mean_rgb.len() != 3 || std_rgb.len() != 3 {
249            return Err(OCRError::ConfigError {
250                message: format!(
251                    "mean/std must have exactly 3 elements (got mean={}, std={})",
252                    mean_rgb.len(),
253                    std_rgb.len()
254                ),
255            });
256        }
257
258        let (mean, std) = match output_color_order {
259            ColorOrder::RGB => (mean_rgb, std_rgb),
260            ColorOrder::BGR => (
261                vec![mean_rgb[2], mean_rgb[1], mean_rgb[0]],
262                vec![std_rgb[2], std_rgb[1], std_rgb[0]],
263            ),
264        };
265
266        Self::with_color_order(
267            scale,
268            Some(mean),
269            Some(std),
270            order,
271            Some(output_color_order),
272        )
273    }
274
275    /// Applies normalization to a vector of images.
276    ///
277    /// # Arguments
278    ///
279    /// * `imgs` - A vector of DynamicImage instances to normalize
280    ///
281    /// # Returns
282    ///
283    /// A vector of normalized images represented as vectors of f32 values
284    pub fn apply(&self, imgs: Vec<DynamicImage>) -> Vec<Vec<f32>> {
285        imgs.into_iter().map(|img| self.normalize(img)).collect()
286    }
287
288    /// Normalizes a single image.
289    ///
290    /// # Arguments
291    ///
292    /// * `img` - The DynamicImage to normalize
293    ///
294    /// # Returns
295    ///
296    /// A vector of normalized pixel values as f32
297    fn normalize(&self, img: DynamicImage) -> Vec<f32> {
298        let rgb_img = into_rgb8_no_copy(img);
299        self.normalize_rgb(&rgb_img)
300    }
301
302    fn normalize_rgb(&self, rgb_img: &RgbImage) -> Vec<f32> {
303        let (width, height) = rgb_img.dimensions();
304        let channels = 3usize;
305        let mut result = vec![0.0f32; Self::image_len(width, height, channels)];
306        self.normalize_rgb_into(rgb_img, &mut result);
307        result
308    }
309
310    /// Normalizes `rgb_img` into a pre-allocated, tightly packed output buffer.
311    ///
312    /// `out` must have length `width * height * 3` and is laid out per
313    /// [`self.order`](TensorLayout). Writes through the SIMD kernels (see
314    /// [`crate::processors::simd`]) operating on the raw interleaved RGB bytes,
315    /// avoiding per-pixel `get_pixel`/`pixels()` overhead.
316    fn normalize_rgb_into(&self, rgb_img: &RgbImage, out: &mut [f32]) {
317        let (width, height) = rgb_img.dimensions();
318        let (width, height) = (width as usize, height as usize);
319        let src_channels = self.src_channels();
320        let alpha = [self.alpha[0], self.alpha[1], self.alpha[2]];
321        let beta = [self.beta[0], self.beta[1], self.beta[2]];
322        let rgb = rgb_img.as_raw();
323
324        match self.order {
325            TensorLayout::CHW => crate::processors::simd::normalize_chw_into(
326                rgb,
327                width,
328                height,
329                src_channels,
330                &alpha,
331                &beta,
332                out,
333            ),
334            TensorLayout::HWC => crate::processors::simd::normalize_hwc_into(
335                rgb,
336                width,
337                height,
338                src_channels,
339                &alpha,
340                &beta,
341                out,
342            ),
343        }
344    }
345
346    /// Normalizes a single image and returns it as a 4D tensor.
347    ///
348    /// # Arguments
349    ///
350    /// * `img` - The DynamicImage to normalize
351    ///
352    /// # Returns
353    ///
354    /// A Result containing the normalized image as a 4D tensor or an OCRError.
355    pub fn normalize_to(&self, img: DynamicImage) -> Result<ndarray::Array4<f32>, OCRError> {
356        let rgb_img = into_rgb8_no_copy(img);
357        let (width, height) = rgb_img.dimensions();
358        let channels = 3usize;
359        let image_len = Self::image_len(width, height, channels);
360
361        match self.order {
362            TensorLayout::CHW => {
363                let result = self.normalize_rgb(&rgb_img);
364
365                ndarray::Array4::from_shape_vec(
366                    (1, channels, height as usize, width as usize),
367                    result,
368                )
369                .map_err(|e| {
370                    OCRError::tensor_operation_error(
371                        "normalization_tensor_creation_chw",
372                        &[1, channels, height as usize, width as usize],
373                        &[image_len],
374                        &format!("Failed to create CHW normalization tensor for {}x{} image with {} channels",
375                            width, height, channels),
376                        e,
377                    )
378                })
379            }
380            TensorLayout::HWC => {
381                let result = self.normalize_rgb(&rgb_img);
382
383                ndarray::Array4::from_shape_vec(
384                    (1, height as usize, width as usize, channels),
385                    result,
386                )
387                .map_err(|e| {
388                    OCRError::tensor_operation_error(
389                        "normalization_tensor_creation_hwc",
390                        &[1, height as usize, width as usize, channels],
391                        &[image_len],
392                        &format!("Failed to create HWC normalization tensor for {}x{} image with {} channels",
393                            width, height, channels),
394                        e,
395                    )
396                })
397            }
398        }
399    }
400
401    /// Normalizes a batch of images and returns them as a 4D tensor.
402    ///
403    /// # Arguments
404    ///
405    /// * `imgs` - A vector of DynamicImage instances to normalize
406    ///
407    /// # Returns
408    ///
409    /// A Result containing the normalized batch as a 4D tensor or an OCRError.
410    ///
411    /// # Errors
412    ///
413    /// Returns an error if:
414    /// * Images in the batch don't all have the same dimensions
415    pub fn normalize_batch_to(
416        &self,
417        imgs: Vec<DynamicImage>,
418    ) -> Result<ndarray::Array4<f32>, OCRError> {
419        let rgb_imgs: Vec<_> = imgs.into_iter().map(into_rgb8_no_copy).collect();
420        let refs: Vec<_> = rgb_imgs.iter().collect();
421        self.normalize_batch_refs(&refs)
422    }
423
424    /// Normalizes a batch of borrowed RGB images into a 4D tensor.
425    ///
426    /// This is equivalent to [`Self::normalize_batch_to`] but lets model
427    /// preprocessors keep shared page images borrowed when resizing is not
428    /// required. Only the output tensor is allocated.
429    pub fn normalize_batch_refs(
430        &self,
431        imgs: &[&RgbImage],
432    ) -> Result<ndarray::Array4<f32>, OCRError> {
433        if imgs.is_empty() {
434            return Ok(ndarray::Array4::zeros((0, 0, 0, 0)));
435        }
436
437        let batch_size = imgs.len();
438        let dimensions: Vec<_> = imgs.iter().map(|img| img.dimensions()).collect();
439
440        let (first_width, first_height) = dimensions.first().copied().unwrap_or((0, 0));
441        for (i, &(width, height)) in dimensions.iter().enumerate() {
442            if width != first_width || height != first_height {
443                return Err(OCRError::InvalidInput {
444                    message: format!(
445                        "All images in batch must have the same dimensions. Image 0: {first_width}x{first_height}, Image {i}: {width}x{height}"
446                    ),
447                });
448            }
449        }
450
451        let (width, height) = (first_width, first_height);
452        let channels = 3usize;
453        let img_size = Self::image_len(width, height, channels);
454        let mut result = vec![0.0f32; batch_size * img_size];
455
456        // Parallelism is gated by total batch output size (not per-image size)
457        // so tiny OCR crops stay serial unless the batch is large enough to
458        // amortize rayon overhead. Each image is normalized into its own
459        // contiguous `img_size` slice via the shared SIMD-backed kernel.
460        let use_parallel =
461            Self::should_parallelize(batch_size, result.len() * std::mem::size_of::<f32>());
462        if !use_parallel {
463            for (rgb_img, batch_slice) in imgs.iter().zip(result.chunks_mut(img_size)) {
464                self.normalize_rgb_into(rgb_img, batch_slice);
465            }
466        } else {
467            result
468                .par_chunks_mut(img_size)
469                .zip(imgs.par_iter())
470                .for_each(|(batch_slice, rgb_img)| {
471                    self.normalize_rgb_into(rgb_img, batch_slice);
472                });
473        }
474
475        let shape = match self.order {
476            TensorLayout::CHW => (batch_size, channels, height as usize, width as usize),
477            TensorLayout::HWC => (batch_size, height as usize, width as usize, channels),
478        };
479        ndarray::Array4::from_shape_vec(shape, result).map_err(|e| {
480            OCRError::tensor_operation("Failed to create batch normalization tensor", e)
481        })
482    }
483}
484
485fn into_rgb8_no_copy(img: DynamicImage) -> RgbImage {
486    match img {
487        DynamicImage::ImageRgb8(img) => img,
488        img => img.to_rgb8(),
489    }
490}
491
492#[cfg(test)]
493mod tests {
494    use super::*;
495    use image::{Rgb, RgbImage};
496    use ndarray::Axis;
497
498    #[test]
499    fn test_normalize_image_color_order_rgb_vs_bgr_chw() -> Result<(), OCRError> {
500        let mut img = RgbImage::new(1, 1);
501        img.put_pixel(0, 0, Rgb([10, 20, 30])); // R, G, B
502
503        let rgb = NormalizeImage::with_color_order(
504            Some(1.0),
505            Some(vec![0.0, 0.0, 0.0]),
506            Some(vec![1.0, 1.0, 1.0]),
507            Some(TensorLayout::CHW),
508            Some(ColorOrder::RGB),
509        )?;
510        let bgr = NormalizeImage::with_color_order(
511            Some(1.0),
512            Some(vec![0.0, 0.0, 0.0]),
513            Some(vec![1.0, 1.0, 1.0]),
514            Some(TensorLayout::CHW),
515            Some(ColorOrder::BGR),
516        )?;
517
518        let rgb_out = rgb.apply(vec![DynamicImage::ImageRgb8(img.clone())]);
519        let bgr_out = bgr.apply(vec![DynamicImage::ImageRgb8(img)]);
520
521        assert_eq!(rgb_out.len(), 1);
522        assert_eq!(bgr_out.len(), 1);
523        assert_eq!(rgb_out[0], vec![10.0, 20.0, 30.0]);
524        assert_eq!(bgr_out[0], vec![30.0, 20.0, 10.0]);
525        Ok(())
526    }
527
528    #[test]
529    fn test_normalize_image_mean_std_applied_in_output_channel_order() -> Result<(), OCRError> {
530        let mut img = RgbImage::new(1, 1);
531        img.put_pixel(0, 0, Rgb([11, 22, 33])); // R, G, B
532
533        let rgb = NormalizeImage::with_color_order(
534            Some(1.0),
535            Some(vec![1.0, 2.0, 3.0]), // RGB means
536            Some(vec![2.0, 4.0, 5.0]), // RGB stds
537            Some(TensorLayout::CHW),
538            Some(ColorOrder::RGB),
539        )?;
540        let bgr = NormalizeImage::with_color_order(
541            Some(1.0),
542            Some(vec![3.0, 2.0, 1.0]), // BGR means
543            Some(vec![5.0, 4.0, 2.0]), // BGR stds
544            Some(TensorLayout::CHW),
545            Some(ColorOrder::BGR),
546        )?;
547
548        let rgb_out = rgb.apply(vec![DynamicImage::ImageRgb8(img.clone())]);
549        let bgr_out = bgr.apply(vec![DynamicImage::ImageRgb8(img)]);
550
551        assert_eq!(rgb_out[0], vec![5.0, 5.0, 6.0]); // (R-1)/2, (G-2)/4, (B-3)/5
552        assert_eq!(bgr_out[0], vec![6.0, 5.0, 5.0]); // (B-3)/5, (G-2)/4, (R-1)/2
553        Ok(())
554    }
555
556    #[test]
557    fn test_should_parallelize_threshold_behavior() {
558        assert!(!NormalizeImage::should_parallelize(
559            1,
560            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES * 4,
561        ));
562        assert!(!NormalizeImage::should_parallelize(
563            4,
564            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES,
565        ));
566        assert!(NormalizeImage::should_parallelize(
567            4,
568            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES + 1,
569        ));
570    }
571
572    #[test]
573    fn test_normalize_batch_to_matches_single_image_paths_for_serial_and_parallel()
574    -> Result<(), OCRError> {
575        let normalizer = NormalizeImage::with_color_order(
576            Some(1.0),
577            Some(vec![0.0, 0.0, 0.0]),
578            Some(vec![1.0, 1.0, 1.0]),
579            Some(TensorLayout::CHW),
580            Some(ColorOrder::RGB),
581        )?;
582
583        let mut small_a = RgbImage::new(1, 1);
584        small_a.put_pixel(0, 0, Rgb([1, 2, 3]));
585        let mut small_b = RgbImage::new(1, 1);
586        small_b.put_pixel(0, 0, Rgb([4, 5, 6]));
587        let small_batch = vec![
588            DynamicImage::ImageRgb8(small_a.clone()),
589            DynamicImage::ImageRgb8(small_b.clone()),
590        ];
591        let serial = normalizer.normalize_batch_to(small_batch)?;
592        let serial_expected = [
593            normalizer.normalize_to(DynamicImage::ImageRgb8(small_a))?,
594            normalizer.normalize_to(DynamicImage::ImageRgb8(small_b))?,
595        ];
596
597        assert_eq!(serial.len_of(Axis(0)), serial_expected.len());
598        for (idx, expected) in serial_expected.iter().enumerate() {
599            assert_eq!(
600                serial.index_axis(Axis(0), idx).to_owned(),
601                expected.index_axis(Axis(0), 0)
602            );
603        }
604
605        let large_a =
606            RgbImage::from_fn(512, 512, |x, y| Rgb([(x % 251) as u8, (y % 241) as u8, 7]));
607        let large_b =
608            RgbImage::from_fn(512, 512, |x, y| Rgb([11, (x % 239) as u8, (y % 233) as u8]));
609        let parallel_batch = vec![
610            DynamicImage::ImageRgb8(large_a.clone()),
611            DynamicImage::ImageRgb8(large_b.clone()),
612        ];
613        let parallel = normalizer.normalize_batch_to(parallel_batch)?;
614        let parallel_expected = [
615            normalizer.normalize_to(DynamicImage::ImageRgb8(large_a))?,
616            normalizer.normalize_to(DynamicImage::ImageRgb8(large_b))?,
617        ];
618
619        assert_eq!(parallel.len_of(Axis(0)), parallel_expected.len());
620        for (idx, expected) in parallel_expected.iter().enumerate() {
621            assert_eq!(
622                parallel.index_axis(Axis(0), idx).to_owned(),
623                expected.index_axis(Axis(0), 0)
624            );
625        }
626
627        Ok(())
628    }
629
630    #[test]
631    fn test_normalize_batch_to_preserves_batch_and_layout_semantics() -> Result<(), OCRError> {
632        let chw = NormalizeImage::with_color_order(
633            Some(1.0),
634            Some(vec![0.0, 0.0, 0.0]),
635            Some(vec![1.0, 1.0, 1.0]),
636            Some(TensorLayout::CHW),
637            Some(ColorOrder::RGB),
638        )?;
639        let hwc = NormalizeImage::with_color_order(
640            Some(1.0),
641            Some(vec![0.0, 0.0, 0.0]),
642            Some(vec![1.0, 1.0, 1.0]),
643            Some(TensorLayout::HWC),
644            Some(ColorOrder::RGB),
645        )?;
646
647        let img_a = RgbImage::from_fn(2, 2, |x, y| {
648            let base = (y * 2 + x) as u8 * 3 + 1;
649            Rgb([base, base + 1, base + 2])
650        });
651        let img_b = RgbImage::from_fn(2, 2, |x, y| {
652            let base = (y * 2 + x) as u8 * 3 + 21;
653            Rgb([base, base + 1, base + 2])
654        });
655
656        let chw_batch = chw.normalize_batch_to(vec![
657            DynamicImage::ImageRgb8(img_a.clone()),
658            DynamicImage::ImageRgb8(img_b.clone()),
659        ])?;
660        assert_eq!(chw_batch.shape(), &[2, 3, 2, 2]);
661        assert_eq!(
662            chw_batch.iter().copied().collect::<Vec<_>>(),
663            vec![
664                1.0, 4.0, 7.0, 10.0, 2.0, 5.0, 8.0, 11.0, 3.0, 6.0, 9.0, 12.0, 21.0, 24.0, 27.0,
665                30.0, 22.0, 25.0, 28.0, 31.0, 23.0, 26.0, 29.0, 32.0,
666            ]
667        );
668
669        let hwc_batch = hwc.normalize_batch_to(vec![
670            DynamicImage::ImageRgb8(img_a),
671            DynamicImage::ImageRgb8(img_b),
672        ])?;
673        assert_eq!(hwc_batch.shape(), &[2, 2, 2, 3]);
674        assert_eq!(
675            hwc_batch.iter().copied().collect::<Vec<_>>(),
676            vec![
677                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 21.0, 22.0, 23.0,
678                24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 32.0,
679            ]
680        );
681
682        Ok(())
683    }
684
685    #[test]
686    fn normalize_batch_refs_matches_owned_path_bit_exact() -> Result<(), OCRError> {
687        let normalizer = NormalizeImage::with_color_order(
688            Some(1.0 / 255.0),
689            Some(vec![0.485, 0.456, 0.406]),
690            Some(vec![0.229, 0.224, 0.225]),
691            Some(TensorLayout::CHW),
692            Some(ColorOrder::BGR),
693        )?;
694        let images = [
695            RgbImage::from_fn(96, 64, |x, y| {
696                Rgb([(x % 251) as u8, (y % 241) as u8, ((x + y) % 239) as u8])
697            }),
698            RgbImage::from_fn(96, 64, |x, y| {
699                Rgb([(y % 233) as u8, ((x * 3) % 229) as u8, 17])
700            }),
701        ];
702        let refs: Vec<_> = images.iter().collect();
703        let borrowed = normalizer.normalize_batch_refs(&refs)?;
704        let owned = normalizer
705            .normalize_batch_to(images.into_iter().map(DynamicImage::ImageRgb8).collect())?;
706
707        assert_eq!(borrowed, owned);
708        Ok(())
709    }
710}
oar_ocr_core/processors/normalization.rs

oar_ocr_core/processors/
normalization.rs
