1use crate::core::OCRError;
8use crate::processors::types::{ColorOrder, TensorLayout};
9use image::{DynamicImage, RgbImage};
10use rayon::prelude::*;
11
12#[derive(Debug)]
18pub struct NormalizeImage {
19 pub alpha: Vec<f32>,
21 pub beta: Vec<f32>,
23 pub order: TensorLayout,
25 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 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 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 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 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 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 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 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 pub fn apply(&self, imgs: Vec<DynamicImage>) -> Vec<Vec<f32>> {
285 imgs.into_iter().map(|img| self.normalize(img)).collect()
286 }
287
288 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 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 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 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 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])); 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])); let rgb = NormalizeImage::with_color_order(
522 Some(1.0),
523 Some(vec![1.0, 2.0, 3.0]), Some(vec![2.0, 4.0, 5.0]), 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]), Some(vec![5.0, 4.0, 2.0]), 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]); assert_eq!(bgr_out[0], vec![6.0, 5.0, 5.0]); 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}