Skip to main content

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        if imgs.is_empty() {
420            return Ok(ndarray::Array4::zeros((0, 0, 0, 0)));
421        }
422
423        let batch_size = imgs.len();
424
425        let rgb_imgs: Vec<_> = imgs.into_iter().map(into_rgb8_no_copy).collect();
426        let dimensions: Vec<_> = rgb_imgs.iter().map(|img| img.dimensions()).collect();
427
428        let (first_width, first_height) = dimensions.first().copied().unwrap_or((0, 0));
429        for (i, &(width, height)) in dimensions.iter().enumerate() {
430            if width != first_width || height != first_height {
431                return Err(OCRError::InvalidInput {
432                    message: format!(
433                        "All images in batch must have the same dimensions. Image 0: {first_width}x{first_height}, Image {i}: {width}x{height}"
434                    ),
435                });
436            }
437        }
438
439        let (width, height) = (first_width, first_height);
440        let channels = 3usize;
441        let img_size = Self::image_len(width, height, channels);
442        let mut result = vec![0.0f32; batch_size * img_size];
443
444        // Parallelism is gated by total batch output size (not per-image size)
445        // so tiny OCR crops stay serial unless the batch is large enough to
446        // amortize rayon overhead. Each image is normalized into its own
447        // contiguous `img_size` slice via the shared SIMD-backed kernel.
448        let use_parallel =
449            Self::should_parallelize(batch_size, result.len() * std::mem::size_of::<f32>());
450        if !use_parallel {
451            for (rgb_img, batch_slice) in rgb_imgs.iter().zip(result.chunks_mut(img_size)) {
452                self.normalize_rgb_into(rgb_img, batch_slice);
453            }
454        } else {
455            result
456                .par_chunks_mut(img_size)
457                .zip(rgb_imgs.par_iter())
458                .for_each(|(batch_slice, rgb_img)| {
459                    self.normalize_rgb_into(rgb_img, batch_slice);
460                });
461        }
462
463        let shape = match self.order {
464            TensorLayout::CHW => (batch_size, channels, height as usize, width as usize),
465            TensorLayout::HWC => (batch_size, height as usize, width as usize, channels),
466        };
467        ndarray::Array4::from_shape_vec(shape, result).map_err(|e| {
468            OCRError::tensor_operation("Failed to create batch normalization tensor", e)
469        })
470    }
471}
472
473fn into_rgb8_no_copy(img: DynamicImage) -> RgbImage {
474    match img {
475        DynamicImage::ImageRgb8(img) => img,
476        img => img.to_rgb8(),
477    }
478}
479
480#[cfg(test)]
481mod tests {
482    use super::*;
483    use image::{Rgb, RgbImage};
484    use ndarray::Axis;
485
486    #[test]
487    fn test_normalize_image_color_order_rgb_vs_bgr_chw() -> Result<(), OCRError> {
488        let mut img = RgbImage::new(1, 1);
489        img.put_pixel(0, 0, Rgb([10, 20, 30])); // R, G, B
490
491        let rgb = NormalizeImage::with_color_order(
492            Some(1.0),
493            Some(vec![0.0, 0.0, 0.0]),
494            Some(vec![1.0, 1.0, 1.0]),
495            Some(TensorLayout::CHW),
496            Some(ColorOrder::RGB),
497        )?;
498        let bgr = NormalizeImage::with_color_order(
499            Some(1.0),
500            Some(vec![0.0, 0.0, 0.0]),
501            Some(vec![1.0, 1.0, 1.0]),
502            Some(TensorLayout::CHW),
503            Some(ColorOrder::BGR),
504        )?;
505
506        let rgb_out = rgb.apply(vec![DynamicImage::ImageRgb8(img.clone())]);
507        let bgr_out = bgr.apply(vec![DynamicImage::ImageRgb8(img)]);
508
509        assert_eq!(rgb_out.len(), 1);
510        assert_eq!(bgr_out.len(), 1);
511        assert_eq!(rgb_out[0], vec![10.0, 20.0, 30.0]);
512        assert_eq!(bgr_out[0], vec![30.0, 20.0, 10.0]);
513        Ok(())
514    }
515
516    #[test]
517    fn test_normalize_image_mean_std_applied_in_output_channel_order() -> Result<(), OCRError> {
518        let mut img = RgbImage::new(1, 1);
519        img.put_pixel(0, 0, Rgb([11, 22, 33])); // R, G, B
520
521        let rgb = NormalizeImage::with_color_order(
522            Some(1.0),
523            Some(vec![1.0, 2.0, 3.0]), // RGB means
524            Some(vec![2.0, 4.0, 5.0]), // RGB stds
525            Some(TensorLayout::CHW),
526            Some(ColorOrder::RGB),
527        )?;
528        let bgr = NormalizeImage::with_color_order(
529            Some(1.0),
530            Some(vec![3.0, 2.0, 1.0]), // BGR means
531            Some(vec![5.0, 4.0, 2.0]), // BGR stds
532            Some(TensorLayout::CHW),
533            Some(ColorOrder::BGR),
534        )?;
535
536        let rgb_out = rgb.apply(vec![DynamicImage::ImageRgb8(img.clone())]);
537        let bgr_out = bgr.apply(vec![DynamicImage::ImageRgb8(img)]);
538
539        assert_eq!(rgb_out[0], vec![5.0, 5.0, 6.0]); // (R-1)/2, (G-2)/4, (B-3)/5
540        assert_eq!(bgr_out[0], vec![6.0, 5.0, 5.0]); // (B-3)/5, (G-2)/4, (R-1)/2
541        Ok(())
542    }
543
544    #[test]
545    fn test_should_parallelize_threshold_behavior() {
546        assert!(!NormalizeImage::should_parallelize(
547            1,
548            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES * 4,
549        ));
550        assert!(!NormalizeImage::should_parallelize(
551            4,
552            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES,
553        ));
554        assert!(NormalizeImage::should_parallelize(
555            4,
556            NormalizeImage::PARALLEL_NORMALIZE_MIN_BYTES + 1,
557        ));
558    }
559
560    #[test]
561    fn test_normalize_batch_to_matches_single_image_paths_for_serial_and_parallel()
562    -> Result<(), OCRError> {
563        let normalizer = NormalizeImage::with_color_order(
564            Some(1.0),
565            Some(vec![0.0, 0.0, 0.0]),
566            Some(vec![1.0, 1.0, 1.0]),
567            Some(TensorLayout::CHW),
568            Some(ColorOrder::RGB),
569        )?;
570
571        let mut small_a = RgbImage::new(1, 1);
572        small_a.put_pixel(0, 0, Rgb([1, 2, 3]));
573        let mut small_b = RgbImage::new(1, 1);
574        small_b.put_pixel(0, 0, Rgb([4, 5, 6]));
575        let small_batch = vec![
576            DynamicImage::ImageRgb8(small_a.clone()),
577            DynamicImage::ImageRgb8(small_b.clone()),
578        ];
579        let serial = normalizer.normalize_batch_to(small_batch)?;
580        let serial_expected = [
581            normalizer.normalize_to(DynamicImage::ImageRgb8(small_a))?,
582            normalizer.normalize_to(DynamicImage::ImageRgb8(small_b))?,
583        ];
584
585        assert_eq!(serial.len_of(Axis(0)), serial_expected.len());
586        for (idx, expected) in serial_expected.iter().enumerate() {
587            assert_eq!(
588                serial.index_axis(Axis(0), idx).to_owned(),
589                expected.index_axis(Axis(0), 0)
590            );
591        }
592
593        let large_a =
594            RgbImage::from_fn(512, 512, |x, y| Rgb([(x % 251) as u8, (y % 241) as u8, 7]));
595        let large_b =
596            RgbImage::from_fn(512, 512, |x, y| Rgb([11, (x % 239) as u8, (y % 233) as u8]));
597        let parallel_batch = vec![
598            DynamicImage::ImageRgb8(large_a.clone()),
599            DynamicImage::ImageRgb8(large_b.clone()),
600        ];
601        let parallel = normalizer.normalize_batch_to(parallel_batch)?;
602        let parallel_expected = [
603            normalizer.normalize_to(DynamicImage::ImageRgb8(large_a))?,
604            normalizer.normalize_to(DynamicImage::ImageRgb8(large_b))?,
605        ];
606
607        assert_eq!(parallel.len_of(Axis(0)), parallel_expected.len());
608        for (idx, expected) in parallel_expected.iter().enumerate() {
609            assert_eq!(
610                parallel.index_axis(Axis(0), idx).to_owned(),
611                expected.index_axis(Axis(0), 0)
612            );
613        }
614
615        Ok(())
616    }
617
618    #[test]
619    fn test_normalize_batch_to_preserves_batch_and_layout_semantics() -> Result<(), OCRError> {
620        let chw = NormalizeImage::with_color_order(
621            Some(1.0),
622            Some(vec![0.0, 0.0, 0.0]),
623            Some(vec![1.0, 1.0, 1.0]),
624            Some(TensorLayout::CHW),
625            Some(ColorOrder::RGB),
626        )?;
627        let hwc = NormalizeImage::with_color_order(
628            Some(1.0),
629            Some(vec![0.0, 0.0, 0.0]),
630            Some(vec![1.0, 1.0, 1.0]),
631            Some(TensorLayout::HWC),
632            Some(ColorOrder::RGB),
633        )?;
634
635        let img_a = RgbImage::from_fn(2, 2, |x, y| {
636            let base = (y * 2 + x) as u8 * 3 + 1;
637            Rgb([base, base + 1, base + 2])
638        });
639        let img_b = RgbImage::from_fn(2, 2, |x, y| {
640            let base = (y * 2 + x) as u8 * 3 + 21;
641            Rgb([base, base + 1, base + 2])
642        });
643
644        let chw_batch = chw.normalize_batch_to(vec![
645            DynamicImage::ImageRgb8(img_a.clone()),
646            DynamicImage::ImageRgb8(img_b.clone()),
647        ])?;
648        assert_eq!(chw_batch.shape(), &[2, 3, 2, 2]);
649        assert_eq!(
650            chw_batch.iter().copied().collect::<Vec<_>>(),
651            vec![
652                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,
653                30.0, 22.0, 25.0, 28.0, 31.0, 23.0, 26.0, 29.0, 32.0,
654            ]
655        );
656
657        let hwc_batch = hwc.normalize_batch_to(vec![
658            DynamicImage::ImageRgb8(img_a),
659            DynamicImage::ImageRgb8(img_b),
660        ])?;
661        assert_eq!(hwc_batch.shape(), &[2, 2, 2, 3]);
662        assert_eq!(
663            hwc_batch.iter().copied().collect::<Vec<_>>(),
664            vec![
665                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,
666                24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 32.0,
667            ]
668        );
669
670        Ok(())
671    }
672}