Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement more imgproc functionality #152

Merged
merged 4 commits into from
Sep 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ members = [
"crates/kornia-imgproc",
"crates/kornia",
"examples/*",
# "kornia-py",
]
exclude = ["kornia-py", "kornia-serve"]

Expand Down
168 changes: 166 additions & 2 deletions crates/kornia-imgproc/src/color/gray.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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,
],
)?;

Expand All @@ -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(())
}
}
2 changes: 1 addition & 1 deletion crates/kornia-imgproc/src/color/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
4 changes: 2 additions & 2 deletions crates/kornia-imgproc/src/enhance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
1 change: 0 additions & 1 deletion examples/onnx/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion kornia-py/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
58 changes: 58 additions & 0 deletions kornia-py/src/color.rs
Original file line number Diff line number Diff line change
@@ -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())
}
44 changes: 44 additions & 0 deletions kornia-py/src/enhance.rs
Original file line number Diff line number Diff line change
@@ -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())
}
11 changes: 4 additions & 7 deletions kornia-py/src/image.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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() {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@emilmgeorge tangential to #149 -- ideally we would like in python too to have zerop copies, especially that in the python case we need to allocate everytime the outputs, or find an anti-pattern for python to allow passing preallocated images :)

Ok(d) => d,
Err(_) => return Err(ImageError::ImageDataNotContiguous),
};

let size = ImageSize {
Expand Down
7 changes: 7 additions & 0 deletions kornia-py/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
mod color;
mod enhance;
mod histogram;
mod image;
mod io;
Expand All @@ -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>()?;
Expand Down
Loading
Loading