From dec0dd45eb944f4233daf5c271dd080512617ec4 Mon Sep 17 00:00:00 2001 From: jianjunz Date: Wed, 24 Jan 2024 06:26:49 +0800 Subject: [PATCH] image2tensor supports planar format output. (#92) * image2tensor supports planar format output. * Add enum MemoryLayout for planar output. --- image2tensor/src/lib.rs | 128 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 118 insertions(+), 10 deletions(-) diff --git a/image2tensor/src/lib.rs b/image2tensor/src/lib.rs index f56df3f..1e3e9f7 100644 --- a/image2tensor/src/lib.rs +++ b/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, 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, 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 = 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 { + let mut out: Vec = 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]); + } +}