Advise for zero copy between numpy / arrow-rs

[edgarriba]

HI! I'm running kornia-rs a Rust computer vision library in which i recently moved away from the ndarray crate to represent our internal Tensor and Image structures towards supporting the arrow-rs project.

I have implemented the following approach in order to provide minimal python support to interact with numpy. And I'm looking for adivse wether this could be most efficient / elegant way to achieve zero-copy conversions. Thanks in advance !

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 {
    fn to_pyimage(self) -> PyImage;
}


impl<const C: usize> ToPyImage for Image<u8, C> {
    fn to_pyimage(self) -> PyImage {
        Python::with_gil(|py| unsafe {
            let array = PyArray::<u8, _>::new_bound(py, [self.height(), self.width(), C], false);
            // TODO: verify that the data is contiguous, otherwise iterate over the image and copy
            std::ptr::copy_nonoverlapping(self.as_ptr(), array.data(), self.numel());
            array.unbind()
        })
    }
}


/// Trait to convert a PyImage (3D numpy array of u8) to an image
pub trait FromPyImage<const C: usize> {
    fn from_pyimage(image: PyImage) -> Result<Image<u8, C>, ImageError>;
}


impl<const C: usize> FromPyImage<C> for Image<u8, C> {
    fn from_pyimage(image: PyImage) -> Result<Image<u8, C>, ImageError> {
        Python::with_gil(|py| {
            let pyarray = image.bind(py);


            // 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),
                }
            };


            let size = ImageSize {
                width: pyarray.shape()[1],
                height: pyarray.shape()[0],
            };


            Image::new(size, data)
        })
    }
}