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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use crate::ndarray as nd;
use crate::opencv::{core as cv, prelude::*};
use crate::with_opencv::MatExt as _;
use crate::with_opencv::OpenCvElement;
use crate::{common::*, TryFromCv, TryIntoCv};

impl<'a, A, D> TryFromCv<&'a cv::Mat> for nd::ArrayView<'a, A, D>
where
    A: OpenCvElement,
    D: nd::Dimension,
{
    type Error = anyhow::Error;

    fn try_from_cv(from: &'a cv::Mat) -> Result<Self, Self::Error> {
        let src_shape = from.size_with_depth();
        let array = nd::ArrayViewD::from_shape(src_shape, from.as_slice()?)?;
        let array = array.into_dimensionality()?;
        Ok(array)
    }
}

impl<A, D> TryFromCv<&cv::Mat> for nd::Array<A, D>
where
    A: OpenCvElement + Clone,
    D: nd::Dimension,
{
    type Error = anyhow::Error;

    fn try_from_cv(from: &cv::Mat) -> Result<Self, Self::Error> {
        let src_shape = from.size_with_depth();
        let array = nd::ArrayViewD::from_shape(src_shape, from.as_slice()?)?;
        let array = array.into_dimensionality()?;
        let array = array.into_owned();
        Ok(array)
    }
}

impl<A, D> TryFromCv<cv::Mat> for nd::Array<A, D>
where
    A: OpenCvElement + Clone,
    D: nd::Dimension,
{
    type Error = anyhow::Error;

    fn try_from_cv(from: cv::Mat) -> Result<Self, Self::Error> {
        (&from).try_into_cv()
    }
}

impl<A, S, D> TryFromCv<&nd::ArrayBase<S, D>> for cv::Mat
where
    A: cv::DataType,
    S: nd::RawData<Elem = A> + nd::Data,
    D: nd::Dimension,
{
    type Error = Error;

    fn try_from_cv(from: &nd::ArrayBase<S, D>) -> Result<Self> {
        let shape_with_channels: Vec<i32> = from.shape().iter().map(|&sz| sz as i32).collect();
        let (channels, shape) = match shape_with_channels.split_last() {
            Some(split) => split,
            None => {
                return Ok(Mat::default());
            }
        };
        let array = from.as_standard_layout();
        let slice = array.as_slice().unwrap();
        let mat = cv::Mat::from_slice(slice)?.reshape_nd(*channels, shape)?;
        Ok(mat)
    }
}

impl<A, S, D> TryFromCv<nd::ArrayBase<S, D>> for cv::Mat
where
    A: cv::DataType,
    S: nd::RawData<Elem = A> + nd::Data,
    D: nd::Dimension,
{
    type Error = Error;

    fn try_from_cv(from: nd::ArrayBase<S, D>) -> Result<Self> {
        (&from).try_into_cv()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use itertools::chain;
    use itertools::Itertools as _;
    use rand::prelude::*;

    #[test]
    fn opencv_ndarray_conversion() -> Result<()> {
        let mut rng = rand::thread_rng();

        for _ in 0..5 {
            // Generate a random shape
            let ndim: usize = rng.gen_range(2..=4);
            let shape: Vec<usize> = (0..ndim).map(|_| rng.gen_range(1..=32)).collect();

            let in_mat = cv::Mat::new_randn_nd::<f32>(&shape)?;
            let view: nd::ArrayViewD<f32> = (&in_mat).try_into_cv()?;
            let array: nd::ArrayD<f32> = (&in_mat).try_into_cv()?;
            let out_mat: cv::Mat = (&array).try_into_cv()?;

            shape
                .iter()
                .map(|&size| 0..size)
                .multi_cartesian_product()
                .try_for_each(|index| {
                    // OpenCV expects a &[i32] index.
                    let index_cv: Vec<_> = index.iter().map(|&size| size as i32).collect();
                    let e1: f32 = *in_mat.at_nd(&index_cv)?;

                    // It adds an extra dimension for Mat ->
                    // nd::ArrayView conversion.
                    let index_nd: Vec<_> = chain!(index, [0]).collect();
                    let e2 = view[index_nd.as_slice()];

                    // It adds an extra dimension for Mat -> nd::Array
                    // conversion.
                    let e3 = array[index_nd.as_slice()];

                    // Ensure the path Mat -> nd::Array -> Mat
                    // preserves the values.
                    let e4: f32 = *out_mat.at_nd(&index_cv)?;

                    ensure!(e1 == e2);
                    ensure!(e1 == e3);
                    ensure!(e1 == e4);
                    anyhow::Ok(())
                })?;
        }

        Ok(())
    }
}