Add nchw conversion option (#173)

* add nchw conversion option

* fix test, cargo fmt

Author: rahulchaphalkar

crates/openvino-tensor-converter/src/lib.rs

@@ -14,6 +14,34 @@ use opencv::core::{MatTraitConst, Scalar_};
 use std::convert::TryInto;
 use std::{num::ParseIntError, path::Path, str::FromStr};
 
+/// Convert an image from NHWC format to NCHW format.
+fn nhwc_to_nchw(data: &[u8], dimensions: &Dimensions) -> Vec<u8> {
+    let mut nchw_data = vec![0; data.len()];
+    let (height, width, channels) = (
+        dimensions.height as usize,
+        dimensions.width as usize,
+        dimensions.channels as usize,
+    );
+    assert_eq!(
+        data.len(),
+        height * width * channels * dimensions.precision.bytes()
+    );
+    for h in 0..height {
+        for w in 0..width {
+            for c in 0..channels {
+                let nhwc_index =
+                    (h * width * channels + w * channels + c) * dimensions.precision.bytes();
+                let nchw_index =
+                    (c * height * width + h * width + w) * dimensions.precision.bytes();
+                for b in 0..dimensions.precision.bytes() {
+                    nchw_data[nchw_index + b] = data[nhwc_index + b];
+                }
+            }
+        }
+    }
+    nchw_data
+}
+
 /// Convert an image a path to a resized sequence of bytes.
 ///
 /// # Errors
@@ -23,6 +51,7 @@ use std::{num::ParseIntError, path::Path, str::FromStr};
 pub fn convert<P: AsRef<Path>>(
     path: P,
     dimensions: &Dimensions,
+    format: &str,
 ) -> Result<Vec<u8>, ConversionError> {
     let path = path.as_ref();
     info!("Converting {} to {:?}", path.display(), dimensions);
@@ -77,7 +106,12 @@ pub fn convert<P: AsRef<Path>>(
 
     // Copy the bytes of the Mat out to a Vec<u8>.
     let dst_slice = unsafe { slice::from_raw_parts(dst.data(), dimensions.bytes()) };
-    Ok(dst_slice.to_vec())
+    let nhwc_data = dst_slice.to_vec();
+    match format {
+        "nchw" => Ok(nhwc_to_nchw(&nhwc_data, dimensions)),
+        "nhwc" => Ok(nhwc_data),
+        _ => Err(ConversionError("Invalid format specified.".to_string())),
+    }
 }
 
 /// Container for the reasons a conversion can fail.

crates/openvino-tensor-converter/src/main.rs

@@ -6,7 +6,8 @@ fn main() {
     env_logger::init();
     let options = Options::from_args();
     let dimensions = Dimensions::from_str(&options.dimensions).expect("Failed to parse dimensions");
-    let tensor_data = convert(options.input, &dimensions).expect("Failed to convert image");
+    let tensor_data =
+        convert(options.input, &dimensions, &options.format).expect("Failed to convert image");
     fs::write(options.output, tensor_data).expect("Failed to write tensor")
 }
 
@@ -27,4 +28,8 @@ struct Options {
     /// The dimensions of the output file as "[height]x[width]x[channels]x[precision]"; e.g. 300x300x3xfp32.
     #[structopt(name = "OUTPUT DIMENSIONS")]
     dimensions: String,
+
+    /// Format of the output tensor: "nchw" or "nhwc".
+    #[structopt(name = "OUTPUT FORMAT", default_value = "nchw")]
+    format: String,
 }

crates/openvino-tensor-converter/tests/conversion.rs

@@ -5,9 +5,10 @@ use openvino_tensor_converter::{convert, Dimensions, Precision};
 fn same_result_twice_u8() {
     let input = "tests/test.jpg";
     let dimensions = Dimensions::new(227, 227, 3, Precision::U8);
+    let format = "nchw";
 
-    let first = convert(input, &dimensions).unwrap();
-    let second = convert(input, &dimensions).unwrap();
+    let first = convert(input, &dimensions, &format).unwrap();
+    let second = convert(input, &dimensions, &format).unwrap();
     assert_same_bytes(&first, &second);
 }
 
@@ -16,9 +17,10 @@ fn same_result_twice_fp32() {
     env_logger::init();
     let input = "tests/test.jpg";
     let dimensions = Dimensions::new(227, 227, 3, Precision::FP32);
+    let format = "nhwc";
 
-    let first = convert(input, &dimensions).unwrap();
-    let second = convert(input, &dimensions).unwrap();
+    let first = convert(input, &dimensions, &format).unwrap();
+    let second = convert(input, &dimensions, &format).unwrap();
     assert_same_bytes(&first, &second);
 }