[wip] python: implement image to pyarrow

Cargo.toml

@@ -8,7 +8,7 @@ members = [
     "examples/dora/imgproc",
     "examples/dora/video-sink",
     "examples/dora/image-utils",
-    # "kornia-py",
+    "kornia-py",
 ]
 exclude = ["kornia-py", "examples/dora"]
 
@@ -58,6 +58,6 @@ thiserror = "2"
 # temporary fixes:
 # https://github.com/rerun-io/rerun/issues/9159
 # https://github.com/apache/arrow-rs/issues/7201#issuecomment-2685998489
-arrow = "=54.2.1"
-arrow-arith = "=54.2.1"
-chrono = "=0.4.39"
+#arrow = "=54.2.1"
+#arrow-arith = "=54.2.1"
+#chrono = "=0.4.39"

kornia-py/Cargo.toml

@@ -10,6 +10,8 @@ repository = "https://github.com/kornia/kornia-rs"
 rust-version = "1.76"
 version = "0.1.9-rc.2"
 
+[workspace]
+
 [lib]
 name = "kornia_rs"
 crate-type = ["cdylib"]
@@ -26,3 +28,7 @@ kornia-3d = { path = "../crates/kornia-3d" }
 # external
 pyo3 = { version = "0.24.0", features = ["extension-module"] }
 numpy = { version = "0.24.0" }
+
+arrow = { version = "55.0.0", features = ["pyarrow"] }
+rand = { version = "0.9.0", features = ["std_rng"] }
+rand_distr = { version = "0.5.1" }

kornia-py/requirements-dev.txt

@@ -7,6 +7,7 @@ pre-commit
 pytest
 pytest-run-parallel
 numpy
+pyarrow
 
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch

kornia-py/src/image.rs

@@ -3,9 +3,30 @@ use numpy::{PyArray, PyArray3, PyArrayMethods, PyUntypedArrayMethods};
 use kornia_image::{Image, ImageError, ImageSize};
 use pyo3::prelude::*;
 
+use arrow::array::{Array, UInt8Array};
+use arrow::pyarrow::{IntoPyArrow, PyArrowType};
+
+pub struct PyImageWrapper<const C: usize>(pub Image<u8, C>);
+
 // type alias for a 3D numpy array of u8
 pub type PyImage = Py<PyArray3<u8>>;
 
+impl<const C: usize> IntoPyArrow for PyImageWrapper<C> {
+    fn into_pyarrow(self, py: Python<'_>) -> PyResult<PyObject> {
+        let array = UInt8Array::from_iter(self.0.as_slice().iter().copied());
+        let data = array.into_data();
+        let pyarrow = PyArrowType(data);
+        let obj = pyarrow.into_pyobject(py)?;
+        Ok(obj.into())
+    }
+}
+
+impl<const C: usize> From<Image<u8, C>> for PyImageWrapper<C> {
+    fn from(image: Image<u8, C>) -> Self {
+        PyImageWrapper(image)
+    }
+}
+
 /// Trait to convert an image to a PyImage (3D numpy array of u8)
 pub trait ToPyImage {
     fn to_pyimage(self) -> PyImage;

kornia-py/src/lib.rs

@@ -14,6 +14,9 @@ use crate::io::functional::{decode_image_jpeg, read_image_any, read_image_jpeg,
 use crate::io::jpeg::decode_image_raw_jpeg;
 use crate::io::png::decode_image_png;
 use crate::io::turbojpeg::{PyImageDecoder, PyImageEncoder};
+use arrow::pyarrow::IntoPyArrow;
+use image::PyImageWrapper;
+use kornia_image::Image;
 use pyo3::prelude::*;
 
 pub fn get_version() -> String {
@@ -26,6 +29,19 @@ pub fn get_version() -> String {
     version.replace("-alpha", "a").replace("-beta", "b")
 }
 
+/// Generate a random boolean array from a Bernoulli distribution.
+use arrow::{array::ArrayData, pyarrow::PyArrowType};
+use pyo3::exceptions::PyValueError;
+
+#[pyfunction]
+fn image_to_pyarrow(py: Python<'_>) -> PyResult<PyArrowType<ArrayData>> {
+    let image = Image::<u8, 1>::new([1, 3].into(), vec![4u8, 5, 6])
+        .map_err(|err| PyValueError::new_err(err.to_string()))?;
+    let data = PyImageWrapper::from(image);
+    let obj = data.into_pyarrow(py)?;
+    Ok(obj.extract::<PyArrowType<ArrayData>>(py)?)
+}
+
 #[pymodule(gil_used = false)]
 pub fn kornia_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", get_version())?;
@@ -49,5 +65,6 @@ pub fn kornia_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyImageEncoder>()?;
     m.add_class::<PyICPConvergenceCriteria>()?;
     m.add_class::<PyICPResult>()?;
+    m.add_function(wrap_pyfunction!(image_to_pyarrow, m)?)?;
     Ok(())
 }

kornia-py/tests/test_pyarrow.py

@@ -0,0 +1,9 @@
+import kornia_rs as K
+
+import numpy as np
+
+
+def test_bernoulli():
+    breakpoint()
+    arr = K.image_to_pyarrow()
+    pass