Implement more imgproc functionality (#152)

* expose to python

* implement add_weighted

* disable py crate

* fix doctest

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "crates/kornia-imgproc",
     "crates/kornia",
     "examples/*",
+    # "kornia-py",
 ]
 exclude = ["kornia-py", "kornia-serve"]
 

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

@@ -70,6 +70,93 @@ where
     Ok(())
 }
 
+/// Convert a grayscale image to an RGB image by replicating the grayscale value across all three channels.
+///
+/// # Arguments
+///
+/// * `src` - The input grayscale image.
+/// * `dst` - The output RGB image.
+///
+/// Precondition: the input image must have 1 channel.
+/// Precondition: the output image must have 3 channels.
+/// Precondition: the input and output images must have the same size.
+///
+/// # Example
+///
+/// ```
+/// use kornia_image::{Image, ImageSize};
+/// use kornia_imgproc::color::rgb_from_gray;
+///
+/// let image = Image::<f32, 1>::new(
+///     ImageSize {
+///         width: 4,
+///         height: 5,
+///     },
+///     vec![0f32; 4 * 5 * 1],
+/// )
+/// .unwrap();
+///
+/// let mut rgb = Image::<f32, 3>::from_size_val(image.size(), 0.0).unwrap();
+///
+/// rgb_from_gray(&image, &mut rgb).unwrap();
+/// ```
+pub fn rgb_from_gray<T>(src: &Image<T, 1>, dst: &mut Image<T, 3>) -> Result<(), ImageError>
+where
+    T: SafeTensorType,
+{
+    if src.size() != dst.size() {
+        return Err(ImageError::InvalidImageSize(
+            src.cols(),
+            src.rows(),
+            dst.cols(),
+            dst.rows(),
+        ));
+    }
+
+    // parallelize the grayscale conversion by rows
+    parallel::par_iter_rows(src, dst, |src_pixel, dst_pixel| {
+        let gray = src_pixel[0];
+        dst_pixel.iter_mut().for_each(|dst_pixel| {
+            *dst_pixel = gray;
+        });
+    });
+
+    Ok(())
+}
+
+/// Convert an RGB image to BGR by swapping the red and blue channels.
+///
+/// # Arguments
+///
+/// * `src` - The input RGB image.
+/// * `dst` - The output BGR image.
+///
+/// Precondition: the input and output images must have the same size.
+pub fn bgr_from_rgb<T>(src: &Image<T, 3>, dst: &mut Image<T, 3>) -> Result<(), ImageError>
+where
+    T: SafeTensorType,
+{
+    if src.size() != dst.size() {
+        return Err(ImageError::InvalidImageSize(
+            src.cols(),
+            src.rows(),
+            dst.cols(),
+            dst.rows(),
+        ));
+    }
+
+    parallel::par_iter_rows(src, dst, |src_pixel, dst_pixel| {
+        dst_pixel
+            .iter_mut()
+            .zip(src_pixel.iter().rev())
+            .for_each(|(d, s)| {
+                *d = *s;
+            });
+    });
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use kornia_image::{ops, Image, ImageSize};
@@ -94,14 +181,19 @@ mod tests {
 
     #[test]
     fn gray_from_rgb_regression() -> Result<(), Box<dyn std::error::Error>> {
+        #[rustfmt::skip]
         let image = Image::new(
             ImageSize {
                 width: 2,
                 height: 3,
             },
             vec![
-                1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
-                0.0, 0.0,
+                1.0, 0.0, 0.0,
+                0.0, 1.0, 0.0,
+                0.0, 0.0, 1.0,
+                0.0, 0.0, 0.0,
+                0.0, 0.0, 0.0,
+                0.0, 0.0, 0.0,
             ],
         )?;
 
@@ -123,4 +215,76 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn rgb_from_grayscale() -> Result<(), Box<dyn std::error::Error>> {
+        let image = Image::new(
+            ImageSize {
+                width: 2,
+                height: 3,
+            },
+            vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
+        )?;
+
+        let mut rgb = Image::<f32, 3>::from_size_val(image.size(), 0.0)?;
+
+        super::rgb_from_gray(&image, &mut rgb)?;
+
+        #[rustfmt::skip]
+        let expected: Image<f32, 3> = Image::new(
+            ImageSize {
+                width: 2,
+                height: 3,
+            },
+            vec![
+                0.0, 0.0, 0.0,
+                1.0, 1.0, 1.0,
+                2.0, 2.0, 2.0,
+                3.0, 3.0, 3.0,
+                4.0, 4.0, 4.0,
+                5.0, 5.0, 5.0,
+            ],
+        )?;
+
+        assert_eq!(rgb.as_slice(), expected.as_slice());
+
+        Ok(())
+    }
+
+    #[test]
+    fn bgr_from_rgb() -> Result<(), Box<dyn std::error::Error>> {
+        #[rustfmt::skip]
+        let image = Image::new(
+            ImageSize {
+                width: 1,
+                height: 3,
+            },
+            vec![
+                0.0, 1.0, 2.0,
+                3.0, 4.0, 5.0,
+                6.0, 7.0, 8.0,
+            ],
+        )?;
+
+        let mut bgr = Image::<f32, 3>::from_size_val(image.size(), 0.0)?;
+
+        super::bgr_from_rgb(&image, &mut bgr)?;
+
+        #[rustfmt::skip]
+        let expected: Image<f32, 3> = Image::new(
+            ImageSize {
+                width: 1,
+                height: 3,
+            },
+            vec![
+                2.0, 1.0, 0.0,
+                5.0, 4.0, 3.0,
+                8.0, 7.0, 6.0,
+            ],
+        )?;
+
+        assert_eq!(bgr.as_slice(), expected.as_slice());
+
+        Ok(())
+    }
 }

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

@@ -1,5 +1,5 @@
 mod gray;
 mod hsv;
 
-pub use gray::gray_from_rgb;
+pub use gray::{bgr_from_rgb, gray_from_rgb, rgb_from_gray};
 pub use hsv::hsv_from_rgb;

crates/kornia-imgproc/src/enhance.rs

@@ -28,9 +28,9 @@ pub fn add_weighted<T, const C: usize>(
     src1: &Image<T, C>,
     alpha: T,
     src2: &Image<T, C>,
-    dst: &mut Image<T, C>,
     beta: T,
     gamma: T,
+    dst: &mut Image<T, C>,
 ) -> Result<(), ImageError>
 where
     T: num_traits::Float
@@ -96,7 +96,7 @@ mod tests {
 
         let mut weighted = Image::<f32, 1>::from_size_val(src1.size(), 0.0)?;
 
-        super::add_weighted(&src1, alpha, &src2, &mut weighted, beta, gamma)?;
+        super::add_weighted(&src1, alpha, &src2, beta, gamma, &mut weighted)?;
 
         weighted
             .as_slice()

examples/onnx/Cargo.toml

@@ -7,7 +7,6 @@ edition.workspace = true
 homepage.workspace = true
 include.workspace = true
 license.workspace = true
-license-file.workspace = true
 readme.workspace = true
 repository.workspace = true
 rust-version.workspace = true

kornia-py/Cargo.toml

@@ -6,7 +6,6 @@ edition = "2021"
 homepage = "http://kornia.org"
 include = ["Cargo.toml"]
 license = "Apache-2.0"
-license-file = "LICENSE"
 repository = "https://github.com/kornia/kornia-rs"
 rust-version = "1.76"
 version = "0.1.6-rc.5"

kornia-py/src/color.rs

@@ -0,0 +1,58 @@
+use pyo3::prelude::*;
+
+use crate::image::{FromPyImage, PyImage, ToPyImage};
+use kornia_image::Image;
+use kornia_imgproc::color;
+
+#[pyfunction]
+pub fn rgb_from_gray(image: PyImage) -> PyResult<PyImage> {
+    let image_gray = Image::from_pyimage(image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("src image: {}", e)))?;
+
+    let mut image_rgb = Image::from_size_val(image_gray.size(), 0u8)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    color::rgb_from_gray(&image_gray, &mut image_rgb).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("failed to convert image: {}", e))
+    })?;
+
+    Ok(image_rgb.to_pyimage())
+}
+
+#[pyfunction]
+pub fn bgr_from_rgb(image: PyImage) -> PyResult<PyImage> {
+    let image_rgb = Image::from_pyimage(image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("src image: {}", e)))?;
+
+    let mut image_bgr = Image::from_size_val(image_rgb.size(), 0u8)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    color::bgr_from_rgb(&image_rgb, &mut image_bgr).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("failed to convert image: {}", e))
+    })?;
+
+    Ok(image_bgr.to_pyimage())
+}
+
+#[pyfunction]
+pub fn gray_from_rgb(image: PyImage) -> PyResult<PyImage> {
+    let image_rgb = Image::from_pyimage(image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("src image: {}", e)))?;
+
+    let image_rgb = image_rgb.cast::<f32>().map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("failed to convert image: {}", e))
+    })?;
+
+    let mut image_gray = Image::from_size_val(image_rgb.size(), 0f32)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    color::gray_from_rgb(&image_rgb, &mut image_gray).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("failed to convert image: {}", e))
+    })?;
+
+    let image_gray = image_gray.cast::<u8>().map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("failed to convert image: {}", e))
+    })?;
+
+    Ok(image_gray.to_pyimage())
+}

kornia-py/src/enhance.rs

@@ -0,0 +1,44 @@
+use pyo3::prelude::*;
+
+use crate::image::{FromPyImage, PyImage, ToPyImage};
+use kornia_image::Image;
+use kornia_imgproc::enhance;
+
+#[pyfunction]
+pub fn add_weighted(
+    src1: PyImage,
+    alpha: f32,
+    src2: PyImage,
+    beta: f32,
+    gamma: f32,
+) -> PyResult<PyImage> {
+    let image1: Image<u8, 3> = Image::from_pyimage(src1).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("src1 image: {}", e))
+    })?;
+
+    let image2: Image<u8, 3> = Image::from_pyimage(src2).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("src2 image: {}", e))
+    })?;
+
+    // cast input images to f32
+    let image1 = image1.cast::<f32>().map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("src1 image: {}", e))
+    })?;
+
+    let image2 = image2.cast::<f32>().map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyException, _>(format!("src2 image: {}", e))
+    })?;
+
+    let mut dst: Image<f32, 3> = Image::from_size_val(image1.size(), 0.0f32)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    enhance::add_weighted(&image1, alpha, &image2, beta, gamma, &mut dst)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    // cast dst image to u8
+    let dst = dst
+        .cast::<u8>()
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("dst image: {}", e)))?;
+
+    Ok(dst.to_pyimage())
+}

kornia-py/src/image.rs

@@ -5,7 +5,6 @@ use pyo3::prelude::*;
 
 // type alias for a 3D numpy array of u8
 pub type PyImage = Py<PyArray3<u8>>;
-//pub type PyImage<'a> = Bound<'a, PyArray3<u8>>;
 
 /// Trait to convert an image to a PyImage (3D numpy array of u8)
 pub trait ToPyImage {
@@ -36,12 +35,10 @@ impl<const C: usize> FromPyImage<C> for Image<u8, C> {
             // TODO: we should find a way to avoid copying the data
             // Possible solutions:
             // - Use a custom ndarray wrapper that does not copy the data
-            // - Return direectly pyarray and use it in the Rust code
-            let data = unsafe {
-                match pyarray.as_slice() {
-                    Ok(d) => d.to_vec(),
-                    Err(_) => return Err(ImageError::ImageDataNotContiguous),
-                }
+            // - Return directly pyarray and use it in the Rust code
+            let data = match pyarray.to_vec() {
+                Ok(d) => d,
+                Err(_) => return Err(ImageError::ImageDataNotContiguous),
             };
 
             let size = ImageSize {

kornia-py/src/lib.rs

@@ -1,3 +1,5 @@
+mod color;
+mod enhance;
 mod histogram;
 mod image;
 mod io;
@@ -22,11 +24,16 @@ pub fn get_version() -> String {
 #[pymodule]
 pub fn kornia_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", get_version())?;
+    m.add_function(wrap_pyfunction!(color::rgb_from_gray, m)?)?;
+    m.add_function(wrap_pyfunction!(color::bgr_from_rgb, m)?)?;
+    m.add_function(wrap_pyfunction!(color::gray_from_rgb, m)?)?;
+    m.add_function(wrap_pyfunction!(enhance::add_weighted, m)?)?;
     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_function(wrap_pyfunction!(resize::resize, m)?)?;
     m.add_function(wrap_pyfunction!(warp::warp_affine, m)?)?;
+    m.add_function(wrap_pyfunction!(warp::warp_perspective, m)?)?;
     m.add_function(wrap_pyfunction!(histogram::compute_histogram, m)?)?;
     m.add_class::<PyImageSize>()?;
     m.add_class::<PyImageDecoder>()?;