Add Warp affine

Cargo.toml

@@ -52,3 +52,7 @@ harness = false
 [[bench]]
 name = "bench_io"
 harness = false
+
+[[bench]]
+name = "bench_warp"
+harness = false

benches/bench_warp.rs

@@ -0,0 +1,25 @@
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+
+use kornia_rs::image::{Image, ImageSize};
+use kornia_rs::resize::InterpolationMode;
+use kornia_rs::warp::{get_rotation_matrix2d, warp_affine};
+
+fn bench_warp_affine(c: &mut Criterion) {
+    let mut group = c.benchmark_group("warp_affine");
+    let image_sizes = vec![(256, 224), (512, 448), (1024, 896)];
+
+    for (width, height) in image_sizes {
+        let image_size = ImageSize { width, height };
+        let id = format!("{}x{}", width, height);
+        let image = Image::<u8, 3>::new(image_size, vec![0u8; width * height * 3]).unwrap();
+        let image_f32 = image.clone().cast::<f32>().unwrap();
+        let m = get_rotation_matrix2d((width as f32 / 2.0, height as f32 / 2.0), 45.0, 1.0);
+        group.bench_with_input(BenchmarkId::new("native", &id), &image_f32, |b, i| {
+            b.iter(|| warp_affine(black_box(i), m, image_size, InterpolationMode::Bilinear))
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_warp_affine);
+criterion_main!(benches);

kornia-py/benchmark/bench_warp_affine.py

@@ -0,0 +1,52 @@
+import timeit
+
+import cv2
+from PIL import Image
+import kornia_rs
+import numpy as np
+# import tensorflow as tf
+
+image_path = "tests/data/dog.jpeg"
+img = kornia_rs.read_image_jpeg(image_path)
+img_pil = Image.open(image_path)
+height, width, _ = img.shape
+M = cv2.getRotationMatrix2D((width / 2, height / 2), 45.0, 1.0)
+M_tuple = tuple(M.flatten())
+N = 5000  # number of iterations
+
+
+def warp_affine_opencv(img: np.ndarray, M: np.ndarray) -> None:
+    return cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), flags=cv2.INTER_LINEAR)
+
+def warp_affine_pil(img: Image.Image, M_tuple: tuple) -> None:
+    return img.transform(img.size, Image.Transform.AFFINE, M_tuple, Image.BILINEAR)
+
+def warp_affine_kornia(img: np.ndarray, M_tuple: tuple) -> None:
+    return kornia_rs.warp_affine(img, M_tuple, img.shape[:2], "bilinear")
+
+tests = [
+    {
+        "name": "OpenCV",
+        "stmt": "warp_affine_opencv(img, M)",
+        "setup": "from __main__ import warp_affine_opencv, img, M",
+        "globals": {"img": img, "M": M},
+    },
+    {
+        "name": "PIL",
+        "stmt": "warp_affine_pil(img_pil, M_tuple)",
+        "setup": "from __main__ import warp_affine_pil, img_pil, M_tuple",
+        "globals": {"img_pil": img_pil, "M_tuple": M_tuple},
+    },
+    {
+        "name": "Kornia",
+        "stmt": "warp_affine_kornia(img, M_tuple)",
+        "setup": "from __main__ import warp_affine_kornia, img, M_tuple",
+        "globals": {"img": img, "M_tuple": M_tuple},
+    },
+]
+
+for test in tests:
+    timer = timeit.Timer(
+        stmt=test["stmt"], setup=test["setup"], globals=test["globals"]
+    )
+    print(f"{test['name']}: {timer.timeit(N)/ N * 1e3:.2f} ms")

kornia-py/src/lib.rs

@@ -2,6 +2,7 @@ mod histogram;
 mod image;
 mod io;
 mod resize;
+mod warp;
 
 use crate::image::PyImageSize;
 use crate::io::functional::{read_image_any, read_image_jpeg, write_image_jpeg};
@@ -25,6 +26,7 @@ pub fn kornia_rs(_py: Python, m: &PyModule) -> PyResult<()> {
     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!(histogram::compute_histogram, m)?)?;
     m.add_class::<PyImageSize>()?;
     m.add_class::<PyImageDecoder>()?;

kornia-py/src/warp.rs

@@ -0,0 +1,47 @@
+use pyo3::prelude::*;
+
+use crate::image::{FromPyImage, PyImage, ToPyImage};
+use kornia_rs::image::Image;
+
+#[pyfunction]
+pub fn warp_affine(
+    image: PyImage,
+    m: (f32, f32, f32, f32, f32, f32),
+    new_size: (usize, usize),
+    interpolation: &str,
+) -> PyResult<PyImage> {
+    // have to add annotation Image<u8, 3>, otherwise the compiler will complain
+    // NOTE: do we support images with channels != 3?
+    let image: Image<u8, 3> = Image::from_pyimage(image)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+
+    let new_size = kornia_rs::image::ImageSize {
+        height: new_size.0,
+        width: new_size.1,
+    };
+
+    let interpolation = match interpolation.to_lowercase().as_str() {
+        "nearest" => kornia_rs::resize::InterpolationMode::Nearest,
+        "bilinear" => kornia_rs::resize::InterpolationMode::Bilinear,
+        _ => {
+            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                "Invalid interpolation mode",
+            ))
+        }
+    };
+
+    // we need to cast to f32 for now since kornia-rs interpolation function only works with f32
+    let image = image
+        .cast::<f32>()
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+
+    let image = kornia_rs::warp::warp_affine(&image, m, new_size, interpolation)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+
+    // NOTE: for bicubic interpolation (not implemented yet), f32 may overshoot 255
+    let image = image
+        .cast::<u8>()
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", e)))?;
+
+    Ok(image.to_pyimage())
+}

kornia-py/tests/test_warp_affine.py

@@ -0,0 +1,23 @@
+from pathlib import Path
+import kornia_rs as K
+
+import numpy as np
+
+# TODO: inject this from elsewhere
+DATA_DIR = Path(__file__).parents[2] / "tests" / "data"
+
+
+def test_warp_affine():
+    # load an image with libjpeg-turbo
+    img_path: Path = DATA_DIR / "dog.jpeg"
+    img: np.ndarray = K.read_image_jpeg(str(img_path.absolute()))
+
+    # check the image properties
+    assert img.shape == (195, 258, 3)
+
+    affine_matrix = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0)
+
+    img_transformed: np.ndarray = K.warp_affine(
+        img, affine_matrix, img.shape[:2], "bilinear"
+    )
+    assert (img_transformed == img).all()

src/lib.rs

@@ -12,3 +12,4 @@ pub mod resize;
 // NOTE: not ready yet
 // pub mod tensor;
 pub mod threshold;
+pub mod warp;

src/resize.rs

@@ -71,6 +71,7 @@ impl ImageDtype for u8 {
 // TODO: add support for other data types. Maybe use a trait? or template?
 fn bilinear_interpolation<T: ImageDtype>(image: &Array3<T>, u: f32, v: f32, c: usize) -> T {
     let (height, width, _) = image.dim();
+
     let iu = u.trunc() as usize;
     let iv = v.trunc() as usize;
 
@@ -135,6 +136,19 @@ pub enum InterpolationMode {
     Nearest,
 }
 
+pub(crate) fn interpolate_pixel<T: ImageDtype>(
+    image: &Array3<T>,
+    u: f32,
+    v: f32,
+    c: usize,
+    interpolation: InterpolationMode,
+) -> T {
+    match interpolation {
+        InterpolationMode::Bilinear => bilinear_interpolation(image, u, v, c),
+        InterpolationMode::Nearest => nearest_neighbor_interpolation(image, u, v, c),
+    }
+}
+
 /// Resize an image to a new size.
 ///
 /// The function resizes an image to a new size using the specified interpolation mode.
@@ -206,10 +220,8 @@ pub fn resize_native<T: ImageDtype, const CHANNELS: usize>(
             let (u, v) = (uv[0], uv[1]);
 
             // compute the pixel values for each channel
-            let pixels = (0..image.num_channels()).map(|k| match interpolation {
-                InterpolationMode::Bilinear => bilinear_interpolation(&image.data, u, v, k),
-                InterpolationMode::Nearest => nearest_neighbor_interpolation(&image.data, u, v, k),
-            });
+            let pixels = (0..image.num_channels())
+                .map(|k| interpolate_pixel(&image.data, u, v, k, interpolation));
 
             // write the pixel values to the output image
             for (k, pixel) in pixels.enumerate() {