image2tensor supports planar format output. (#92)

* image2tensor supports planar format output.

* Add enum MemoryLayout for planar output.

image2tensor/src/lib.rs

@@ -1,4 +1,4 @@
-use image::{io::Reader, DynamicImage};
+use image::{io::Reader, ColorType, DynamicImage};
 
 #[derive(Debug, Copy, Clone)]
 pub enum TensorType {
@@ -14,6 +14,12 @@ pub enum ColorOrder {
     BGR,
 }
 
+#[derive(Debug, Copy, Clone)]
+pub enum MemoryLayout {
+    Planar,
+    Interleaved,
+}
+
 // Take the image located at 'path', open it, resize it to `height` x `width`, and then convert the
 // pixel precision to the type requested. NOTE: this function assumes the image is in standard 8-bit
 // RGB format. It should work with standard image formats such as `.jpg`, `.png`, etc.
@@ -34,8 +40,51 @@ pub fn convert_image_to_tensor_bytes(
         .decode()
         .map_err(|_| format!("Failed to decode the file: {:?}", path))
         .unwrap();
+    let color_type = decoded.color();
 
-    convert_dynamic_image_to_tensor_bytes(decoded, width, height, precision, order)
+    convert_dynamic_image_to_tensor_bytes(
+        decoded,
+        width,
+        height,
+        color_type,
+        precision,
+        order,
+        MemoryLayout::Interleaved,
+    )
+}
+
+// Take the image located at 'path', open it, resize it to `height` x `width`,  then convert the
+// pixel precision to the type requested in planar (not interleave) format. NOTE: this function
+// assumes the image is in standard 8-bit RGB format. It should work with standard image formats
+// such as `.jpg`, `.png`, etc.
+pub fn convert_image_to_planar_tensor_bytes(
+    path: &str,
+    width: u32,
+    height: u32,
+    precision: TensorType,
+    order: ColorOrder,
+) -> Result<Vec<u8>, String> {
+    // Open the file and create the reader.
+    let raw_file = Reader::open(path)
+        .map_err(|_| format!("Failed to open the file: {:?}", path))
+        .unwrap();
+
+    // Create the DynamicImage by decoding the image.
+    let decoded = raw_file
+        .decode()
+        .map_err(|_| format!("Failed to decode the file: {:?}", path))
+        .unwrap();
+    let color_type = decoded.color();
+
+    convert_dynamic_image_to_tensor_bytes(
+        decoded,
+        width,
+        height,
+        color_type,
+        precision,
+        order,
+        MemoryLayout::Planar,
+    )
 }
 
 /// Same as [convert_image_to_tensor_bytes] but accepts a `bytes` slice instead.
@@ -49,25 +98,50 @@ pub fn convert_image_bytes_to_tensor_bytes(
     // Create the DynamicImage by decoding the image.
     let decoded = image::load_from_memory(bytes).expect("Unable to load image from bytes.");
 
-    convert_dynamic_image_to_tensor_bytes(decoded, width, height, precision, order)
+    convert_dynamic_image_to_tensor_bytes(
+        decoded,
+        width,
+        height,
+        ColorType::Rgb8,
+        precision,
+        order,
+        MemoryLayout::Interleaved,
+    )
 }
 
 fn convert_dynamic_image_to_tensor_bytes(
     image: DynamicImage,
     width: u32,
     height: u32,
+    src_color: ColorType,
     precision: TensorType,
     order: ColorOrder,
+    layout: MemoryLayout,
 ) -> Result<Vec<u8>, String> {
     // Resize the image to the specified W/H and get an array of u8 RGB values.
     let dyn_img: DynamicImage = image.resize_exact(width, height, image::imageops::Triangle);
-    let mut img_bytes = dyn_img.into_bytes();
+    let mut img_bytes: Vec<u8> = dyn_img.into_bytes();
 
-    // Get an array of the pixel values and return it.
-    match order {
-        ColorOrder::RGB => Ok(save_bytes(&img_bytes, precision)),
-        ColorOrder::BGR => Ok(save_bytes(rgb_to_bgr(&mut img_bytes), precision)),
+    let src_channels = match src_color {
+        ColorType::Rgba8 | ColorType::Rgba16 | ColorType::Rgba32F => 4,
+        ColorType::Rgb8 | ColorType::Rgb16 | ColorType::Rgb32F => 3,
+        _ => unimplemented!(),
+    };
+
+    if matches!(order, ColorOrder::BGR) {
+        rgb_to_bgr(&mut img_bytes, src_channels);
     }
+
+    if matches!(layout, MemoryLayout::Planar) {
+        img_bytes =
+            interleave_to_planar(&img_bytes, src_channels, width as usize * height as usize);
+    }
+
+    // Output channel number is always 3 at this moment because we only support RGB and BGR.
+    Ok(save_bytes(
+        &img_bytes[..width as usize * height as usize * 3],
+        precision,
+    ))
 }
 
 /// Calculate the expected tensor data size of an image of `width` x `height` with the given
@@ -107,9 +181,43 @@ fn get_bytes_per_pixel(precision: TensorType) -> usize {
 }
 
 /// Converts an RGB array to BGR.
-fn rgb_to_bgr(buffer: &mut [u8]) -> &[u8] {
-    for i in (0..buffer.len()).step_by(3) {
+fn rgb_to_bgr(buffer: &mut [u8], channels: usize) -> &[u8] {
+    for i in (0..buffer.len()).step_by(channels) {
         buffer.swap(i + 2, i);
     }
     buffer
 }
+
+fn interleave_to_planar(buffer: &[u8], channels: usize, pixels: usize) -> Vec<u8> {
+    let mut out: Vec<u8> = Vec::with_capacity(buffer.len());
+    unsafe {
+        out.set_len(buffer.len());
+    }
+    let mut offset = 0;
+    for i in (0..buffer.len()).step_by(channels) {
+        for c in 0..channels {
+            out[pixels * c + offset] = buffer[i + c];
+        }
+        offset += 1;
+    }
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn convert_image() {
+        let path = std::path::Path::new("./../rust/examples/images/0.jpg");
+        let tensor_bytes = convert_image_to_tensor_bytes(
+            path.canonicalize().unwrap().to_str().unwrap(),
+            224,
+            224,
+            TensorType::F32,
+            ColorOrder::RGB,
+        )
+        .unwrap();
+        assert_eq!(&tensor_bytes[..8], [0, 0, 16, 65, 0, 0, 96, 65]);
+    }
+}