fix python and decompres api

Cargo.toml

@@ -9,7 +9,7 @@ categories = ["image-processing", "computer-vision", "machine-learning"]
 license-file = "LICENSE"
 
 [dependencies]
-anyhow = "1.0.79"
+anyhow = "1.0.80"
 num-traits = "0.2.17"
 image = { version = "0.24.5" }
 turbojpeg = {version = "1.0.0"}

py-kornia/Cargo.toml

@@ -8,6 +8,7 @@ name = "kornia_rs"
 crate-type = ["cdylib"]
 
 [dependencies]
+anyhow = "1.0.80"
 kornia-rs = { path = ".." }
 dlpack-rs = {version = "^0.1.1"}
 pyo3 = { version = "0.20", features = ["extension-module"] }

py-kornia/src/dlpack.rs

@@ -1,3 +1,7 @@
+//! This module provides utilities to convert between PyTorch and DLPack tensors.
+//!
+//! NOTE: this is deprecated and will be removed in the future.
+//!
 use dlpack_rs as dlpack;
 use kornia_rs::tensor::Tensor;
 

py-kornia/src/image.rs

@@ -1,8 +1,45 @@
+use anyhow::Result;
 use numpy::{PyArray3, ToPyArray};
 
 use kornia_rs::image::{Image, ImageSize};
 use pyo3::prelude::*;
 
+// type alias for a 3D numpy array of u8
+pub type PyImage = Py<PyArray3<u8>>;
+
+/// Trait to convert an image to a PyImage (3D numpy array of u8)
+pub trait ToPyImage {
+    fn to_pyimage(&self) -> PyImage;
+}
+
+impl<const CHANNELS: usize> ToPyImage for kornia_rs::image::Image<u8, CHANNELS> {
+    fn to_pyimage(&self) -> PyImage {
+        Python::with_gil(|py| self.data.to_pyarray(py).to_owned())
+    }
+}
+
+/// Trait to convert a PyImage (3D numpy array of u8) to an image
+pub trait FromPyImage<const CHANNELS: usize> {
+    fn from_pyimage(image: PyImage) -> Result<Image<u8, CHANNELS>>;
+}
+
+impl<const CHANNELS: usize> FromPyImage<CHANNELS> for kornia_rs::image::Image<u8, CHANNELS> {
+    fn from_pyimage(image: PyImage) -> Result<Image<u8, CHANNELS>> {
+        Python::with_gil(|py| {
+            let array = image.as_ref(py).to_owned_array();
+            let data = match array.as_slice() {
+                Some(d) => d.to_vec(),
+                None => return Err(anyhow::anyhow!("Image data is not contiguous")),
+            };
+            let size = ImageSize {
+                width: array.shape()[1],
+                height: array.shape()[0],
+            };
+            Ok(Image::new(size, data)?)
+        })
+    }
+}
+
 #[pyclass(name = "ImageSize")]
 #[derive(Clone)]
 pub struct PyImageSize {
@@ -42,82 +79,14 @@ impl PyImageSize {
     }
 }
 
-#[pyclass(name = "Image")]
-#[derive(Clone)]
-pub struct PyImage {
-    pub inner: Image<u8, 3>,
-}
-
-#[pymethods]
-impl PyImage {
-    #[getter]
-    pub fn shape(&self) -> PyResult<(usize, usize, usize)> {
-        Ok((
-            self.inner.image_size().height,
-            self.inner.image_size().width,
-            self.inner.num_channels(),
-        ))
-    }
-
-    pub fn size(&self) -> PyImageSize {
-        self.inner.image_size().into()
-    }
-
-    pub fn wdith(&self) -> usize {
-        self.inner.image_size().width
-    }
-
-    pub fn height(&self) -> usize {
-        self.inner.image_size().height
-    }
-
-    pub fn num_channels(&self) -> usize {
-        self.inner.num_channels()
-    }
-
-    fn __str__(&self) -> PyResult<String> {
-        Ok(format!(
-            "Image(height: {}, width: {}, num_channels: {}, dtype: u8)",
-            self.inner.image_size().height,
-            self.inner.image_size().width,
-            self.inner.num_channels()
-        ))
-    }
-
-    fn __repr__(&self) -> PyResult<String> {
-        Ok(format!(
-            "Image(height: {}, width: {}, num_channels: {}, dtype: u8)",
-            self.inner.image_size().height,
-            self.inner.image_size().width,
-            self.inner.num_channels()
-        ))
-    }
-
-    fn numpy(&self, py: Python) -> PyResult<Py<PyArray3<u8>>> {
-        Ok(self.inner.data.to_pyarray(py).to_owned())
-    }
-}
-
 impl From<ImageSize> for PyImageSize {
     fn from(image_size: ImageSize) -> Self {
         PyImageSize { inner: image_size }
     }
 }
 
-impl From<Image<u8, 3>> for PyImage {
-    fn from(image: Image<u8, 3>) -> Self {
-        PyImage { inner: image }
-    }
-}
-
 impl From<PyImageSize> for ImageSize {
     fn from(image_size: PyImageSize) -> Self {
         image_size.inner
     }
 }
-
-impl From<PyImage> for Image<u8, 3> {
-    fn from(image: PyImage) -> Self {
-        image.inner
-    }
-}

py-kornia/src/io/functional.rs

@@ -1,23 +1,28 @@
 use pyo3::prelude::*;
 use std::path::Path;
 
-use crate::image::PyImage;
-use kornia_rs::io::functional as F;
+use crate::image::{FromPyImage, PyImage, ToPyImage};
+use kornia_rs::{image::Image, io::functional as F};
 
 #[pyfunction]
-pub fn read_image_jpeg(file_path: String) -> PyResult<PyImage> {
-    let image = F::read_image_jpeg(Path::new(&file_path)).unwrap();
-    Ok(image.into())
+pub fn read_image_jpeg(file_path: &str) -> PyResult<PyImage> {
+    let image = F::read_image_jpeg(Path::new(file_path))
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyFileExistsError, _>(format!("{}", e)))?;
+    Ok(image.to_pyimage())
 }
 
 #[pyfunction]
 pub fn write_image_jpeg(file_path: String, image: PyImage) -> PyResult<()> {
-    F::write_image_jpeg(Path::new(&file_path), image.into()).unwrap();
+    let image = Image::from_pyimage(image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+    F::write_image_jpeg(Path::new(&file_path), image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
     Ok(())
 }
 
 #[pyfunction]
 pub fn read_image_any(file_path: String) -> PyResult<PyImage> {
-    let image = F::read_image_any(Path::new(&file_path)).unwrap();
-    Ok(image.into())
+    let image = F::read_image_any(Path::new(&file_path))
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyFileExistsError, _>(format!("{}", e)))?;
+    Ok(image.to_pyimage())
 }

py-kornia/src/io/jpeg.rs

@@ -1,7 +1,10 @@
-use kornia_rs::io::jpeg::{ImageDecoder, ImageEncoder};
+use kornia_rs::{
+    image::Image,
+    io::jpeg::{ImageDecoder, ImageEncoder},
+};
 use pyo3::prelude::*;
 
-use crate::image::{PyImage, PyImageSize};
+use crate::image::{FromPyImage, PyImage, PyImageSize, ToPyImage};
 
 #[pyclass(name = "ImageDecoder")]
 pub struct PyImageDecoder {
@@ -12,18 +15,25 @@ pub struct PyImageDecoder {
 impl PyImageDecoder {
     #[new]
     pub fn new() -> PyResult<PyImageDecoder> {
-        let inner = ImageDecoder::new();
+        let inner = ImageDecoder::new()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
         Ok(PyImageDecoder { inner })
     }
 
     pub fn read_header(&mut self, jpeg_data: &[u8]) -> PyResult<PyImageSize> {
-        let image_size = self.inner.read_header(jpeg_data);
+        let image_size = self
+            .inner
+            .read_header(jpeg_data)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
         Ok(image_size.into())
     }
 
     pub fn decode(&mut self, jpeg_data: &[u8]) -> PyResult<PyImage> {
-        let image = self.inner.decode(jpeg_data);
-        Ok(image.into())
+        let image = self
+            .inner
+            .decode(jpeg_data)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+        Ok(image.to_pyimage())
     }
 }
 
@@ -36,16 +46,25 @@ pub struct PyImageEncoder {
 impl PyImageEncoder {
     #[new]
     pub fn new() -> PyResult<PyImageEncoder> {
-        let inner = ImageEncoder::new();
+        let inner = ImageEncoder::new()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
         Ok(PyImageEncoder { inner })
     }
 
     pub fn encode(&mut self, image: PyImage) -> PyResult<Vec<u8>> {
-        let jpeg_data = self.inner.encode(image.into());
+        let image = Image::from_pyimage(image)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+        let jpeg_data = self
+            .inner
+            .encode(image)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
         Ok(jpeg_data)
     }
 
-    pub fn set_quality(&mut self, quality: i32) {
-        self.inner.set_quality(quality)
+    pub fn set_quality(&mut self, quality: i32) -> PyResult<()> {
+        self.inner
+            .set_quality(quality)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+        Ok(())
     }
 }

py-kornia/src/lib.rs

@@ -1,12 +1,9 @@
-mod dlpack;
 mod image;
 mod io;
-mod tensor;
 
 use crate::image::PyImageSize;
 use crate::io::functional::{read_image_any, read_image_jpeg, write_image_jpeg};
 use crate::io::jpeg::{PyImageDecoder, PyImageEncoder};
-use crate::tensor::PyTensor;
 use pyo3::prelude::*;
 
 pub fn get_version() -> String {
@@ -25,7 +22,6 @@ pub fn kornia_rs(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(read_image_jpeg, m)?)?;
     m.add_function(wrap_pyfunction!(write_image_jpeg, m)?)?;
     m.add_function(wrap_pyfunction!(read_image_any, m)?)?;
-    m.add_class::<PyTensor>()?;
     m.add_class::<PyImageSize>()?;
     m.add_class::<PyImageDecoder>()?;
     m.add_class::<PyImageEncoder>()?;