From d8eea973470e1fcd9f535dbd42b07f4387a4f00e Mon Sep 17 00:00:00 2001
From: edgar <edgar.riba@gmail.com>
Date: Wed, 31 Jan 2024 12:35:56 +0100
Subject: [PATCH] small improvements

---
 src/color.rs  | 13 +++++++++----
 src/image.rs  | 13 +++++++++++++
 src/resize.rs | 41 ++++++++++++++++++++++++++++++-----------
 3 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/src/color.rs b/src/color.rs
index 7b71e7bc..162b82d6 100644
--- a/src/color.rs
+++ b/src/color.rs
@@ -7,9 +7,16 @@ use ndarray::{Array3, Zip};
 // or automatically:
 // let gray = Image<u8, 1>::try_from(rgb);
 pub fn gray_from_rgb(image: &Image) -> Image {
+    assert_eq!(image.num_channels(), 3);
+
     // TODO: implement this using a map or cast
     // let image_f32 = image.cast::<f32>();
-    let mut output = Array3::<u8>::zeros(image.data.dim());
+    //let mut output = Array3::<u8>::zeros(image.data.dim());
+    let mut output = Array3::<u8>::zeros((
+        image.image_size().height as usize,
+        image.image_size().width as usize,
+        1,
+    ));
 
     Zip::from(output.rows_mut())
         .and(image.data.rows())
@@ -21,8 +28,6 @@ pub fn gray_from_rgb(image: &Image) -> Image {
             let gray = (76. * r + 150. * g + 29. * b) / 255.;
 
             out[0] = gray as u8;
-            out[1] = gray as u8;
-            out[2] = gray as u8;
         });
 
     Image { data: output }
@@ -37,7 +42,7 @@ mod tests {
         let image_path = std::path::Path::new("tests/data/dog.jpeg");
         let image = F::read_image_jpeg(image_path);
         let gray = super::gray_from_rgb(&image);
-        assert_eq!(gray.num_channels(), 3);
+        assert_eq!(gray.num_channels(), 1);
         assert_eq!(gray.image_size().width, 258);
         assert_eq!(gray.image_size().height, 195);
     }
diff --git a/src/image.rs b/src/image.rs
index 7388bef7..ce566528 100644
--- a/src/image.rs
+++ b/src/image.rs
@@ -60,6 +60,7 @@ impl<T> Image<T> {
         Image { data: image }
     }
 
+    // TODO: implement without `num_traits`
     pub fn cast<U>(&self) -> Image<U>
     where
         T: num_traits::cast::AsPrimitive<U>,
@@ -96,6 +97,9 @@ impl<T> Image<T> {
             _ => io::functions::read_image_any(image_path),
         }
     }
+
+    // TODO: implement from bytes
+    // pub fn from_bytes(bytes: &[u8]) -> Image {
 }
 
 #[cfg(test)]
@@ -158,4 +162,13 @@ mod tests {
         let image_i32: Image<i32> = image_u8.cast();
         assert_eq!(image_i32.data.get((1, 0, 2)).unwrap(), &5i32);
     }
+
+    #[test]
+    fn image_rgbd() {
+        use crate::image::Image;
+        let image = Image::from_shape_vec([2, 2, 4], vec![0f32; 2 * 2 * 4]);
+        assert_eq!(image.image_size().width, 2);
+        assert_eq!(image.image_size().height, 2);
+        assert_eq!(image.num_channels(), 4);
+    }
 }
diff --git a/src/resize.rs b/src/resize.rs
index 9363862a..fcd6c063 100644
--- a/src/resize.rs
+++ b/src/resize.rs
@@ -81,7 +81,7 @@ pub fn resize(image: &Image, new_size: ImageSize, optional_args: ResizeOptions)
     let image_size = image.image_size();
 
     // create the output image
-    let mut output = Array3::<u8>::zeros((new_size.height, new_size.width, 3));
+    let mut output = Array3::<u8>::zeros((new_size.height, new_size.width, image.num_channels()));
 
     // create a grid of x and y coordinates for the output image
     // and interpolate the values from the input image.
@@ -120,15 +120,17 @@ pub fn resize(image: &Image, new_size: ImageSize, optional_args: ResizeOptions)
             assert_eq!(uv.len(), 2);
             let (u, v) = (uv[0], uv[1]);
 
-            // TODO: this assumes 3 channels
-            for k in [0, 1, 2].iter() {
-                let pixel = match optional_args.interpolation {
-                    InterpolationMode::Bilinear => bilinear_interpolation(&image.data, u, v, *k),
-                    InterpolationMode::NearestNeighbor => {
-                        nearest_neighbor_interpolation(&image.data, u, v, *k)
-                    }
-                };
-                out[*k] = pixel as u8;
+            // compute the pixel values for each channel
+            let pixels = (0..image.num_channels()).map(|k| match optional_args.interpolation {
+                InterpolationMode::Bilinear => bilinear_interpolation(&image.data, u, v, k),
+                InterpolationMode::NearestNeighbor => {
+                    nearest_neighbor_interpolation(&image.data, u, v, k)
+                }
+            });
+
+            // write the pixel values to the output image
+            for (k, pixel) in pixels.enumerate() {
+                out[k] = pixel as u8;
             }
         });
 
@@ -139,7 +141,7 @@ pub fn resize(image: &Image, new_size: ImageSize, optional_args: ResizeOptions)
 mod tests {
 
     #[test]
-    fn resize_smoke() {
+    fn resize_smoke_ch3() {
         use crate::image::{Image, ImageSize};
         let image = Image::from_shape_vec([4, 5, 3], vec![0; 4 * 5 * 3]);
         let image_resized = super::resize(
@@ -154,4 +156,21 @@ mod tests {
         assert_eq!(image_resized.image_size().width, 2);
         assert_eq!(image_resized.image_size().height, 3);
     }
+
+    #[test]
+    fn resize_smoke_ch1() {
+        use crate::image::{Image, ImageSize};
+        let image = Image::from_shape_vec([4, 5, 1], vec![0; 4 * 5 * 1]);
+        let image_resized = super::resize(
+            &image,
+            ImageSize {
+                width: 2,
+                height: 3,
+            },
+            super::ResizeOptions::default(),
+        );
+        assert_eq!(image_resized.num_channels(), 1);
+        assert_eq!(image_resized.image_size().width, 2);
+        assert_eq!(image_resized.image_size().height, 3);
+    }
 }