diff --git a/crates/kornia-io/Cargo.toml b/crates/kornia-io/Cargo.toml index 891875df..f74c3c2a 100644 --- a/crates/kornia-io/Cargo.toml +++ b/crates/kornia-io/Cargo.toml @@ -16,6 +16,7 @@ all-features = true [dependencies] image = "0.25" kornia-image = { workspace = true } +png = "0.17" log = { workspace = true } thiserror = { workspace = true } diff --git a/crates/kornia-io/src/error.rs b/crates/kornia-io/src/error.rs index efdbee9c..096951c8 100644 --- a/crates/kornia-io/src/error.rs +++ b/crates/kornia-io/src/error.rs @@ -25,4 +25,8 @@ pub enum IoError { /// Error to decode the image. #[error("Failed to decode the image")] ImageDecodeError(#[from] image::ImageError), + + /// Error to decode the PNG image. + #[error("Failed to decode the image")] + PngDecodeError(String), } diff --git a/crates/kornia-io/src/lib.rs b/crates/kornia-io/src/lib.rs index 24656d70..bc30c435 100644 --- a/crates/kornia-io/src/lib.rs +++ b/crates/kornia-io/src/lib.rs @@ -14,6 +14,9 @@ pub mod functional; #[cfg(feature = "jpegturbo")] pub mod jpeg; +/// PNG image encoding and decoding. +pub mod png; + /// GStreamer video module for real-time video processing. #[cfg(feature = "gstreamer")] pub mod stream; diff --git a/crates/kornia-io/src/png.rs b/crates/kornia-io/src/png.rs new file mode 100644 index 00000000..c274385d --- /dev/null +++ b/crates/kornia-io/src/png.rs @@ -0,0 +1,92 @@ +use std::{fs::File, path::Path}; + +use kornia_image::Image; +use png::Decoder; + +use crate::error::IoError; + +/// Read a PNG image with a single channel (mono8). +/// +/// # Arguments +/// +/// * `file_path` - The path to the PNG file. +/// +/// # Returns +/// +/// A grayscale image with a single channel (mono8). +pub fn read_image_png_mono8(file_path: impl AsRef) -> Result, IoError> { + let (buf, size) = read_png_impl(file_path)?; + Ok(Image::new(size.into(), buf)?) +} + +/// Read a PNG image with a three channels (rgb8). +/// +/// # Arguments +/// +/// * `file_path` - The path to the PNG file. +/// +/// # Returns +/// +/// A RGB image with three channels (rgb8). +pub fn read_image_png_rgb8(file_path: impl AsRef) -> Result, IoError> { + let (buf, size) = read_png_impl(file_path)?; + Ok(Image::new(size.into(), buf)?) +} + +/// Read a PNG image with a four channels (rgba8). +/// +/// # Arguments +/// +/// * `file_path` - The path to the PNG file. +/// +/// # Returns +/// +/// A RGBA image with four channels (rgba8). +pub fn read_image_png_rgba8(file_path: impl AsRef) -> Result, IoError> { + let (buf, size) = read_png_impl(file_path)?; + Ok(Image::new(size.into(), buf)?) +} + +// utility function to read the png file +fn read_png_impl(file_path: impl AsRef) -> Result<(Vec, [usize; 2]), IoError> { + // verify the file exists + let file_path = file_path.as_ref(); + if !file_path.exists() { + return Err(IoError::FileDoesNotExist(file_path.to_path_buf())); + } + + // verify the file extension + if let Some(extension) = file_path.extension() { + if extension != "png" { + return Err(IoError::InvalidFileExtension(file_path.to_path_buf())); + } + } else { + return Err(IoError::InvalidFileExtension(file_path.to_path_buf())); + } + + let file = File::open(file_path)?; + let mut reader = Decoder::new(file) + .read_info() + .map_err(|e| IoError::PngDecodeError(e.to_string()))?; + + let mut buf = vec![0; reader.output_buffer_size()]; + let info = reader + .next_frame(&mut buf) + .map_err(|e| IoError::PngDecodeError(e.to_string()))?; + + Ok((buf, [info.width as usize, info.height as usize])) +} + +#[cfg(test)] +mod tests { + use crate::error::IoError; + use crate::png::read_image_png_mono8; + + #[test] + fn read_png_mono8() -> Result<(), IoError> { + let image = read_image_png_mono8("../../tests/data/dog.png")?; + assert_eq!(image.size().width, 258); + assert_eq!(image.size().height, 195); + Ok(()) + } +} diff --git a/tests/data/dog.png b/tests/data/dog.png new file mode 100644 index 00000000..f46e463e Binary files /dev/null and b/tests/data/dog.png differ