use crate::with_opencv::MatExt as _;
use crate::with_opencv::OpenCvElement;
use crate::{TryAsRefCv, TryToCv};
use anyhow::{Error, Result};
use ndarray as nd;
use opencv::{core as cv, prelude::*};
impl<'a, A, D> TryAsRefCv<'a, nd::ArrayView<'a, A, D>> for cv::Mat
where
A: OpenCvElement,
D: nd::Dimension + 'a,
{
type Error = anyhow::Error;
fn try_as_ref_cv(&'a self) -> Result<nd::ArrayView<'a, A, D>, Self::Error> {
let src_shape = self.size_with_depth();
let array = nd::ArrayViewD::from_shape(src_shape, self.as_slice()?)?;
let array = array.into_dimensionality()?;
Ok(array)
}
}
impl<A, D> TryToCv<nd::Array<A, D>> for cv::Mat
where
A: OpenCvElement + Clone,
D: nd::Dimension,
{
type Error = anyhow::Error;
fn try_to_cv(&self) -> Result<nd::Array<A, D>, Self::Error> {
let src_shape = self.size_with_depth();
let array = nd::ArrayViewD::from_shape(src_shape, self.as_slice()?)?;
let array = array.into_dimensionality()?;
let array = array.into_owned();
Ok(array)
}
}
impl<A, S, D> TryToCv<cv::Mat> for nd::ArrayBase<S, D>
where
A: cv::DataType,
S: nd::RawData<Elem = A> + nd::Data,
D: nd::Dimension,
{
type Error = Error;
fn try_to_cv(&self) -> Result<cv::Mat, Self::Error> {
let shape_with_channels: Vec<i32> = self.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 = self.as_standard_layout();
let slice = array.as_slice().unwrap();
let mat = cv::Mat::from_slice(slice)?.reshape_nd(*channels, shape)?;
Ok(mat)
}
}
#[cfg(test)]
mod tests {
use crate::TryToCv;
use super::*;
use anyhow::ensure;
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 {
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_as_ref_cv()?;
let array: nd::ArrayD<f32> = in_mat.try_to_cv()?;
let out_mat: cv::Mat = array.try_to_cv()?;
shape
.iter()
.map(|&size| 0..size)
.multi_cartesian_product()
.try_for_each(|index| {
let index_cv: Vec<_> = index.iter().map(|&size| size as i32).collect();
let e1: f32 = *in_mat.at_nd(&index_cv)?;
let index_nd: Vec<_> = chain!(index, [0]).collect();
let e2 = view[index_nd.as_slice()];
let e3 = array[index_nd.as_slice()];
let e4: f32 = *out_mat.at_nd(&index_cv)?;
ensure!(e1 == e2);
ensure!(e1 == e3);
ensure!(e1 == e4);
anyhow::Ok(())
})?;
}
Ok(())
}
}