cv_convert_fork/
with_opencv_ndarray.rs

1use crate::ndarray as nd;
2use crate::opencv::{core as cv, prelude::*};
3use crate::with_opencv::MatExt as _;
4use crate::with_opencv::OpenCvElement;
5use crate::{common::*, TryFromCv, TryIntoCv};
6
7impl<'a, A, D> TryFromCv<&'a cv::Mat> for nd::ArrayView<'a, A, D>
8where
9    A: OpenCvElement,
10    D: nd::Dimension,
11{
12    type Error = anyhow::Error;
13
14    fn try_from_cv(from: &'a cv::Mat) -> Result<Self, Self::Error> {
15        let src_shape = from.size_with_depth();
16        let array = nd::ArrayViewD::from_shape(src_shape, from.as_slice()?)?;
17        let array = array.into_dimensionality()?;
18        Ok(array)
19    }
20}
21
22impl<A, D> TryFromCv<&cv::Mat> for nd::Array<A, D>
23where
24    A: OpenCvElement + Clone,
25    D: nd::Dimension,
26{
27    type Error = anyhow::Error;
28
29    fn try_from_cv(from: &cv::Mat) -> Result<Self, Self::Error> {
30        let src_shape = from.size_with_depth();
31        let array = nd::ArrayViewD::from_shape(src_shape, from.as_slice()?)?;
32        let array = array.into_dimensionality()?;
33        let array = array.into_owned();
34        Ok(array)
35    }
36}
37
38impl<A, D> TryFromCv<cv::Mat> for nd::Array<A, D>
39where
40    A: OpenCvElement + Clone,
41    D: nd::Dimension,
42{
43    type Error = anyhow::Error;
44
45    fn try_from_cv(from: cv::Mat) -> Result<Self, Self::Error> {
46        (&from).try_into_cv()
47    }
48}
49
50impl<A, S, D> TryFromCv<&nd::ArrayBase<S, D>> for cv::Mat
51where
52    A: cv::DataType,
53    S: nd::RawData<Elem = A> + nd::Data,
54    D: nd::Dimension,
55{
56    type Error = Error;
57
58    fn try_from_cv(from: &nd::ArrayBase<S, D>) -> Result<Self> {
59        let shape_with_channels: Vec<i32> = from.shape().iter().map(|&sz| sz as i32).collect();
60        let (channels, shape) = match shape_with_channels.split_last() {
61            Some(split) => split,
62            None => {
63                return Ok(Mat::default());
64            }
65        };
66        let array = from.as_standard_layout();
67        let slice = array.as_slice().unwrap();
68        let mat = (cv::Mat::from_slice(slice)?.reshape_nd(*channels, shape)?).clone_pointee();
69        Ok(mat)
70    }
71}
72
73impl<A, S, D> TryFromCv<nd::ArrayBase<S, D>> for cv::Mat
74where
75    A: cv::DataType,
76    S: nd::RawData<Elem = A> + nd::Data,
77    D: nd::Dimension,
78{
79    type Error = Error;
80
81    fn try_from_cv(from: nd::ArrayBase<S, D>) -> Result<Self> {
82        (&from).try_into_cv()
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use itertools::chain;
90    use itertools::Itertools as _;
91    use rand::prelude::*;
92
93    #[test]
94    fn opencv_ndarray_conversion() -> Result<()> {
95        let mut rng = rand::thread_rng();
96
97        for _ in 0..5 {
98            // Generate a random shape
99            let ndim: usize = rng.gen_range(2..=4);
100            let shape: Vec<usize> = (0..ndim).map(|_| rng.gen_range(1..=32)).collect();
101
102            let in_mat = cv::Mat::new_randn_nd::<f32>(&shape)?;
103            let view: nd::ArrayViewD<f32> = (&in_mat).try_into_cv()?;
104            let array: nd::ArrayD<f32> = (&in_mat).try_into_cv()?;
105            let out_mat: cv::Mat = (&array).try_into_cv()?;
106
107            shape
108                .iter()
109                .map(|&size| 0..size)
110                .multi_cartesian_product()
111                .try_for_each(|index| {
112                    // OpenCV expects a &[i32] index.
113                    let index_cv: Vec<_> = index.iter().map(|&size| size as i32).collect();
114                    let e1: f32 = *in_mat.at_nd(&index_cv)?;
115
116                    // It adds an extra dimension for Mat ->
117                    // nd::ArrayView conversion.
118                    let index_nd: Vec<_> = chain!(index, [0]).collect();
119                    let e2 = view[index_nd.as_slice()];
120
121                    // It adds an extra dimension for Mat -> nd::Array
122                    // conversion.
123                    let e3 = array[index_nd.as_slice()];
124
125                    // Ensure the path Mat -> nd::Array -> Mat
126                    // preserves the values.
127                    let e4: f32 = *out_mat.at_nd(&index_cv)?;
128
129                    ensure!(e1 == e2);
130                    ensure!(e1 == e3);
131                    ensure!(e1 == e4);
132                    anyhow::Ok(())
133                })?;
134        }
135
136        Ok(())
137    }
138}