diff --git a/ndarray-image/src/lib.rs b/ndarray-image/src/lib.rs index 79d24e9..9f34542 100644 --- a/ndarray-image/src/lib.rs +++ b/ndarray-image/src/lib.rs @@ -1,4 +1,8 @@ type Result = core::result::Result; +fn shape_error() -> ndarray::ShapeError { + ndarray::ShapeError::from_kind(ndarray::ErrorKind::IncompatibleShape) +} + mod rgb8 { use super::Result; pub(super) fn image_as_ndarray(image: &image::RgbImage) -> Result> { @@ -11,6 +15,18 @@ mod rgb8 { let data = image.into_raw(); ndarray::Array3::from_shape_vec((height as usize, width as usize, 3), data) } + pub(super) fn ndarray_to_image(array: &ndarray::ArrayView3) -> Result { + let (height, width, channels) = array.dim(); + let data = array.as_slice().ok_or_else(super::shape_error)?; + if channels != 3 { + return Err(super::shape_error()); + } + Ok( + image::RgbImage::from_raw(width as u32, height as u32, data.to_vec()).ok_or( + ndarray::ShapeError::from_kind(ndarray::ErrorKind::IncompatibleShape), + )?, + ) + } } mod rgba8 { @@ -25,6 +41,18 @@ mod rgba8 { let data = image.into_raw(); ndarray::Array3::from_shape_vec((height as usize, width as usize, 4), data) } + pub(super) fn ndarray_to_image(array: &ndarray::ArrayView3) -> Result { + let (height, width, channels) = array.dim(); + let data = array.as_slice().ok_or_else(super::shape_error)?; + if channels != 4 { + return Err(super::shape_error()); + } + Ok( + image::RgbaImage::from_raw(width as u32, height as u32, data.to_vec()).ok_or( + ndarray::ShapeError::from_kind(ndarray::ErrorKind::IncompatibleShape), + )?, + ) + } } mod gray8 { @@ -39,6 +67,15 @@ mod gray8 { let data = image.into_raw(); ndarray::Array2::from_shape_vec((height as usize, width as usize), data) } + pub(super) fn ndarray_to_image(array: &ndarray::ArrayView2) -> Result { + let (height, width) = array.dim(); + let data = array.as_slice().ok_or_else(super::shape_error)?; + Ok( + image::GrayImage::from_raw(width as u32, height as u32, data.to_vec()).ok_or( + ndarray::ShapeError::from_kind(ndarray::ErrorKind::IncompatibleShape), + )?, + ) + } } mod gray_alpha8 { @@ -55,6 +92,20 @@ mod gray_alpha8 { let data = image.into_raw(); ndarray::Array3::from_shape_vec((height as usize, width as usize, 2), data) } + pub(super) fn ndarray_to_image( + array: &ndarray::ArrayView3, + ) -> Result { + let (height, width, channels) = array.dim(); + let data = array.as_slice().ok_or_else(super::shape_error)?; + if channels != 2 { + return Err(super::shape_error()); + } + Ok( + image::GrayAlphaImage::from_raw(width as u32, height as u32, data.to_vec()).ok_or( + ndarray::ShapeError::from_kind(ndarray::ErrorKind::IncompatibleShape), + )?, + ) + } } mod dynamic_image { @@ -95,6 +146,34 @@ pub trait ImageToNdarray { fn into_ndarray(self) -> Result; } +pub trait NdarrayToImage { + fn to_image(&self) -> Result; +} + +impl NdarrayToImage for ndarray::ArrayView3<'_, u8> { + fn to_image(&self) -> Result { + rgb8::ndarray_to_image(self) + } +} + +impl NdarrayToImage for ndarray::ArrayView3<'_, u8> { + fn to_image(&self) -> Result { + rgba8::ndarray_to_image(self) + } +} + +impl NdarrayToImage for ndarray::ArrayView2<'_, u8> { + fn to_image(&self) -> Result { + gray8::ndarray_to_image(self) + } +} + +impl NdarrayToImage for ndarray::ArrayView3<'_, u8> { + fn to_image(&self) -> Result { + gray_alpha8::ndarray_to_image(self) + } +} + impl ImageToNdarray for image::RgbImage { type OwnedOutput = ndarray::Array3; type RefOutput<'a> = ndarray::ArrayView3<'a, u8>;