-
Notifications
You must be signed in to change notification settings - Fork 16
/
image.rs
106 lines (89 loc) · 2.92 KB
/
image.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
use numpy::{PyArray, PyArray3, PyArrayMethods, PyUntypedArrayMethods};
use kornia::image::{Image, ImageError, ImageSize};
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 {
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)
})
}
}
#[pyclass(name = "ImageSize")]
#[derive(Clone)]
pub struct PyImageSize {
inner: ImageSize,
}
#[pymethods]
impl PyImageSize {
#[new]
pub fn new(width: usize, height: usize) -> PyResult<PyImageSize> {
let inner = ImageSize { width, height };
Ok(PyImageSize { inner })
}
#[getter]
pub fn width(&self) -> usize {
self.inner.width
}
#[getter]
pub fn height(&self) -> usize {
self.inner.height
}
fn __str__(&self) -> PyResult<String> {
Ok(format!(
"ImageSize(width: {}, height: {})",
self.inner.width, self.inner.height
))
}
fn __repr__(&self) -> PyResult<String> {
Ok(format!(
"ImageSize(width: {}, height: {})",
self.inner.width, self.inner.height
))
}
}
impl From<ImageSize> for PyImageSize {
fn from(image_size: ImageSize) -> Self {
PyImageSize { inner: image_size }
}
}
impl From<PyImageSize> for ImageSize {
fn from(image_size: PyImageSize) -> Self {
image_size.inner
}
}