-
Couldn't load subscription status.
- Fork 128
-
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) }) } }
Beta Was this translation helpful? Give feedback.
All reactions
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment