Skip to content

Commit

Permalink
Implement working with image bytes directly (no file required). (#88)
Browse files Browse the repository at this point in the history
* Implement working with image bytes directly (no file required).

Co-authored-by: Matthew Tamayo <matthew@getmethodic.com>

* fix: `covert` to `convert`

* fix: un-`pub` new conversion function, improve documentation

* fix: broken doc link

---------

Co-authored-by: geekbeast <geekbeast@users.noreply.github.com>
Co-authored-by: Matthew Tamayo <matthew@getmethodic.com>
Co-authored-by: Andrew Brown <andrew.brown@intel.com>
  • Loading branch information
4 people authored Aug 25, 2023
1 parent 436d50e commit 6d726e6
Showing 1 changed file with 33 additions and 7 deletions.
40 changes: 33 additions & 7 deletions image2tensor/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,10 @@ pub enum ColorOrder {
BGR,
}

// Take the image located at 'path', open it, resize it to height x width, and then converts
// the pixel precision to the type requested. NOTE: this function assumes the image is in
// standard 8bit RGB format. It should work with standard image formats such as .jpg, .png, etc
pub fn convert_image_to_bytes(
// 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.
pub fn convert_image_to_tensor_bytes(
path: &str,
width: u32,
height: u32,
Expand All @@ -35,8 +35,32 @@ pub fn convert_image_to_bytes(
.map_err(|_| format!("Failed to decode the file: {:?}", path))
.unwrap();

convert_dynamic_image_to_tensor_bytes(decoded, width, height, precision, order)
}

/// Same as [convert_image_to_tensor_bytes] but accepts a `bytes` slice instead.
pub fn convert_image_bytes_to_tensor_bytes(
bytes: &[u8],
width: u32,
height: u32,
precision: TensorType,
order: ColorOrder,
) -> Result<Vec<u8>, String> {
// 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)
}

fn convert_dynamic_image_to_tensor_bytes(
image: DynamicImage,
width: u32,
height: u32,
precision: TensorType,
order: ColorOrder,
) -> Result<Vec<u8>, String> {
// Resize the image to the specified W/H and get an array of u8 RGB values.
let dyn_img: DynamicImage = decoded.resize_exact(width, height, image::imageops::Triangle);
let dyn_img: DynamicImage = image.resize_exact(width, height, image::imageops::Triangle);
let mut img_bytes = dyn_img.into_bytes();

// Get an array of the pixel values and return it.
Expand All @@ -46,13 +70,15 @@ pub fn convert_image_to_bytes(
}
}

/// Calculate the expected tensor data size of an image of `width` x `height` with the given
/// `precision`.
pub fn calculate_buffer_size(width: u32, height: u32, precision: TensorType) -> usize {
let bytes_per_pixel = get_bytes_per_pixel(precision);
let pixels: u32 = width * height * 3;
pixels as usize * bytes_per_pixel
}

// Save the bytes into the specified TensorType format.
/// Save the bytes into the specified TensorType format.
fn save_bytes(buffer: &[u8], tt: TensorType) -> Vec<u8> {
let mut out: Vec<u8> = vec![];

Expand Down Expand Up @@ -80,7 +106,7 @@ fn get_bytes_per_pixel(precision: TensorType) -> usize {
}
}

// Converts an RGB array to BGR
/// Converts an RGB array to BGR.
fn rgb_to_bgr(buffer: &mut [u8]) -> &[u8] {
for i in (0..buffer.len()).step_by(3) {
buffer.swap(i + 2, i);
Expand Down

0 comments on commit 6d726e6

Please sign in to comment.