kornia_io: refactor jpegturbo api and specify image types (#233)

Cargo.toml

@@ -3,8 +3,8 @@ resolver = "2"
 members = [
     "crates/*",
     "examples/*",
-    # "kornia-py",
     "kornia-viz",
+    # "kornia-py",
 ]
 exclude = ["kornia-py"]
 

README.md

@@ -20,7 +20,7 @@ use kornia::io::functional as F;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // read the image
-    let image: Image<u8, 3> = F::read_image_any("tests/data/dog.jpeg")?;
+    let image: Image<u8, 3> = F::read_image_any_rgb8("tests/data/dog.jpeg")?;
 
     println!("Hello, world! 🦀");
     println!("Loaded Image size: {:?}", image.size());
@@ -84,17 +84,18 @@ Add the following to your `Cargo.toml`:
 
 ```toml
 [dependencies]
-kornia = "v0.1.7"
+kornia = "v0.1.8"
 ```
 
 Alternatively, you can use each sub-crate separately:
 
 ```toml
 [dependencies]
-kornia-core = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.7" }
-kornia-io = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.7" }
-kornia-image = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.7" }
-kornia-imgproc = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.7" }
+kornia-tensor = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.8" }
+kornia-io = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.8" }
+kornia-image = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.8" }
+kornia-imgproc = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.8" }
+kornia-3d = { git = "https://github.com/kornia/kornia-rs", tag = "v0.1.8" }
 ```
 
 ### 🐍 Python
@@ -115,7 +116,7 @@ use kornia::io::functional as F;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     // read the image
-    let image: Image<u8, 3> = F::read_image_any("tests/data/dog.jpeg")?;
+    let image: Image<u8, 3> = F::read_image_any_rgb8("tests/data/dog.jpeg")?;
     let image_viz = image.clone();
 
     let image_f32: Image<f32, 3> = image.cast_and_scale::<f32>(1.0 / 255.0)?;

crates/kornia-imgproc/src/color/gray.rs

@@ -162,7 +162,7 @@ mod tests {
 
     #[test]
     fn gray_from_rgb() -> Result<(), Box<dyn std::error::Error>> {
-        let image = F::read_image_any("../../tests/data/dog.jpeg")?;
+        let image = F::read_image_any_rgb8("../../tests/data/dog.jpeg")?;
 
         let mut image_norm = Image::from_size_val(image.size(), 0.0)?;
         ops::cast_and_scale(&image, &mut image_norm, 1. / 255.0)?;
@@ -171,8 +171,8 @@ mod tests {
         super::gray_from_rgb(&image_norm, &mut gray)?;
 
         assert_eq!(gray.num_channels(), 1);
-        assert_eq!(gray.size().width, 258);
-        assert_eq!(gray.size().height, 195);
+        assert_eq!(gray.cols(), 258);
+        assert_eq!(gray.rows(), 195);
 
         Ok(())
     }

crates/kornia-io/Cargo.toml

@@ -23,16 +23,15 @@ thiserror = { workspace = true }
 # optional dependencies
 gst = { version = "0.23.4", package = "gstreamer", optional = true }
 gst-app = { version = "0.23.4", package = "gstreamer-app", optional = true }
-memmap2 = "0.9.4"
-turbojpeg = { version = "1.0.0", optional = true }
+turbojpeg = { version = "1.2", optional = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
 tempfile = { workspace = true }
 
 [features]
 gstreamer = ["gst", "gst-app"]
-jpegturbo = ["turbojpeg"]
+turbojpeg = ["dep:turbojpeg"]
 
 [[bench]]
 name = "bench_io"

crates/kornia-io/benches/bench_io.rs

@@ -1,6 +1,6 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
-use kornia_io::functional::{read_image_any, read_image_jpeg};
+use kornia_io::functional::{read_image_any_rgb8, read_image_jpegturbo_rgb8};
 
 fn bench_read_jpeg(c: &mut Criterion) {
     let mut group = c.benchmark_group("JpegReader");
@@ -9,11 +9,11 @@ fn bench_read_jpeg(c: &mut Criterion) {
 
     // NOTE: this is the fastest method
     group.bench_function("jpegturbo", |b| {
-        b.iter(|| black_box(read_image_jpeg(img_path)).unwrap())
+        b.iter(|| black_box(read_image_jpegturbo_rgb8(img_path)).unwrap())
     });
 
     group.bench_function("image_any", |b| {
-        b.iter(|| black_box(read_image_any(img_path).unwrap()))
+        b.iter(|| black_box(read_image_any_rgb8(img_path)).unwrap())
     });
 
     group.finish();

crates/kornia-io/src/error.rs

@@ -14,9 +14,9 @@ pub enum IoError {
     FileError(#[from] std::io::Error),
 
     /// Error to map the file to memory.
-    #[cfg(feature = "jpegturbo")]
+    #[cfg(feature = "turbojpeg")]
     #[error("Error with Jpeg encoding/decoding")]
-    JpegError(#[from] crate::jpeg::JpegError),
+    JpegTurboError(#[from] crate::jpegturbo::JpegTurboError),
 
     /// Error to create the image.
     #[error("Failed to create image")]

crates/kornia-io/src/functional.rs

@@ -4,11 +4,11 @@ use kornia_image::{Image, ImageSize};
 
 use crate::error::IoError;
 
-#[cfg(feature = "jpegturbo")]
-use super::jpeg::{ImageDecoder, ImageEncoder};
+#[cfg(feature = "turbojpeg")]
+use super::jpegturbo::{JpegTurboDecoder, JpegTurboEncoder};
 
-#[cfg(feature = "jpegturbo")]
-/// Reads a JPEG image from the given file path.
+#[cfg(feature = "turbojpeg")]
+/// Reads a JPEG image in `RGB8` format from the given file path.
 ///
 /// The method reads the JPEG image data directly from a file leveraging the libjpeg-turbo library.
 ///
@@ -26,13 +26,13 @@ use super::jpeg::{ImageDecoder, ImageEncoder};
 /// use kornia_image::Image;
 /// use kornia_io::functional as F;
 ///
-/// let image: Image<u8, 3> = F::read_image_jpeg("../../tests/data/dog.jpeg").unwrap();
+/// let image: Image<u8, 3> = F::read_image_jpegturbo_rgb8("../../tests/data/dog.jpeg").unwrap();
 ///
-/// assert_eq!(image.size().width, 258);
-/// assert_eq!(image.size().height, 195);
+/// assert_eq!(image.cols(), 258);
+/// assert_eq!(image.rows(), 195);
 /// assert_eq!(image.num_channels(), 3);
 /// ```
-pub fn read_image_jpeg(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoError> {
+pub fn read_image_jpegturbo_rgb8(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoError> {
     let file_path = file_path.as_ref().to_owned();
     // verify the file exists and is a JPEG
     if !file_path.exists() {
@@ -46,38 +46,40 @@ pub fn read_image_jpeg(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoEr
     }
 
     // open the file and map it to memory
-    let file = std::fs::File::open(file_path)?;
-    let mmap = unsafe { memmap2::Mmap::map(&file)? };
+    let jpeg_data = std::fs::read(file_path)?;
 
     // decode the data directly from memory
     let image: Image<u8, 3> = {
-        let mut decoder = ImageDecoder::new()?;
-        decoder.decode(&mmap)?
+        let mut decoder = JpegTurboDecoder::new()?;
+        decoder.decode_rgb8(&jpeg_data)?
     };
 
     Ok(image)
 }
 
-#[cfg(feature = "jpegturbo")]
+#[cfg(feature = "turbojpeg")]
 /// Writes the given JPEG data to the given file path.
 ///
 /// # Arguments
 ///
 /// * `file_path` - The path to the JPEG image.
 /// * `image` - The tensor containing the JPEG image data.
-pub fn write_image_jpeg(file_path: impl AsRef<Path>, image: &Image<u8, 3>) -> Result<(), IoError> {
+pub fn write_image_jpegturbo_rgb8(
+    file_path: impl AsRef<Path>,
+    image: &Image<u8, 3>,
+) -> Result<(), IoError> {
     let file_path = file_path.as_ref().to_owned();
 
     // compress the image
-    let jpeg_data = ImageEncoder::new()?.encode(image)?;
+    let jpeg_data = JpegTurboEncoder::new()?.encode_rgb8(image)?;
 
     // write the data directly to a file
     std::fs::write(file_path, jpeg_data)?;
 
     Ok(())
 }
 
-/// Reads an image from the given file path.
+/// Reads a RGB8 image from the given file path.
 ///
 /// The method tries to read from any image format supported by the image crate.
 ///
@@ -87,21 +89,21 @@ pub fn write_image_jpeg(file_path: impl AsRef<Path>, image: &Image<u8, 3>) -> Re
 ///
 /// # Returns
 ///
-/// A tensor containing the image data.
+/// A tensor image containing the image data in RGB8 format with shape (H, W, 3).
 ///
 /// # Example
 ///
 /// ```
 /// use kornia_image::Image;
 /// use kornia_io::functional as F;
 ///
-/// let image: Image<u8, 3> = F::read_image_any("../../tests/data/dog.jpeg").unwrap();
+/// let image: Image<u8, 3> = F::read_image_any_rgb8("../../tests/data/dog.jpeg").unwrap();
 ///
-/// assert_eq!(image.size().width, 258);
-/// assert_eq!(image.size().height, 195);
+/// assert_eq!(image.cols(), 258);
+/// assert_eq!(image.rows(), 195);
 /// assert_eq!(image.num_channels(), 3);
 /// ```
-pub fn read_image_any(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoError> {
+pub fn read_image_any_rgb8(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoError> {
     let file_path = file_path.as_ref().to_owned();
 
     // verify the file exists
@@ -110,13 +112,10 @@ pub fn read_image_any(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoErr
     }
 
     // open the file and map it to memory
-    let file = std::fs::File::open(file_path)?;
-    let mmap = unsafe { memmap2::Mmap::map(&file)? };
+    let jpeg_data = std::fs::read(file_path)?;
 
     // decode the data directly from memory
-    // TODO: update the image crate
-    #[allow(deprecated)]
-    let img = image::io::Reader::new(std::io::Cursor::new(&mmap))
+    let img = image::ImageReader::new(std::io::Cursor::new(&jpeg_data))
         .with_guessed_format()?
         .decode()?;
 
@@ -136,43 +135,43 @@ pub fn read_image_any(file_path: impl AsRef<Path>) -> Result<Image<u8, 3>, IoErr
 #[cfg(test)]
 mod tests {
     use crate::error::IoError;
-    use crate::functional::read_image_any;
+    use crate::functional::read_image_any_rgb8;
 
-    #[cfg(feature = "jpegturbo")]
-    use crate::functional::{read_image_jpeg, write_image_jpeg};
+    #[cfg(feature = "turbojpeg")]
+    use crate::functional::{read_image_jpegturbo_rgb8, write_image_jpegturbo_rgb8};
 
     #[test]
     fn read_any() -> Result<(), IoError> {
-        let image = read_image_any("../../tests/data/dog.jpeg")?;
-        assert_eq!(image.size().width, 258);
-        assert_eq!(image.size().height, 195);
+        let image = read_image_any_rgb8("../../tests/data/dog.jpeg")?;
+        assert_eq!(image.cols(), 258);
+        assert_eq!(image.rows(), 195);
         Ok(())
     }
 
     #[test]
-    #[cfg(feature = "jpegturbo")]
+    #[cfg(feature = "turbojpeg")]
     fn read_jpeg() -> Result<(), IoError> {
-        let image = read_image_jpeg("../../tests/data/dog.jpeg")?;
-        assert_eq!(image.size().width, 258);
-        assert_eq!(image.size().height, 195);
+        let image = read_image_jpegturbo_rgb8("../../tests/data/dog.jpeg")?;
+        assert_eq!(image.cols(), 258);
+        assert_eq!(image.rows(), 195);
         Ok(())
     }
 
     #[test]
-    #[cfg(feature = "jpegturbo")]
+    #[cfg(feature = "turbojpeg")]
     fn read_write_jpeg() -> Result<(), IoError> {
         let tmp_dir = tempfile::tempdir()?;
         std::fs::create_dir_all(tmp_dir.path())?;
 
         let file_path = tmp_dir.path().join("dog.jpeg");
-        let image_data = read_image_jpeg("../../tests/data/dog.jpeg")?;
-        write_image_jpeg(&file_path, &image_data)?;
+        let image_data = read_image_jpegturbo_rgb8("../../tests/data/dog.jpeg")?;
+        write_image_jpegturbo_rgb8(&file_path, &image_data)?;
 
-        let image_data_back = read_image_jpeg(&file_path)?;
+        let image_data_back = read_image_jpegturbo_rgb8(&file_path)?;
         assert!(file_path.exists(), "File does not exist: {:?}", file_path);
 
-        assert_eq!(image_data_back.size().width, 258);
-        assert_eq!(image_data_back.size().height, 195);
+        assert_eq!(image_data_back.cols(), 258);
+        assert_eq!(image_data_back.rows(), 195);
         assert_eq!(image_data_back.num_channels(), 3);
 
         Ok(())