ndarray_vision/transform/
mod.rs

1use crate::core::{ColourModel, Image, ImageBase};
2use ndarray::{prelude::*, s, Data};
3use num_traits::{Num, NumAssignOps};
4use std::fmt::Display;
5
6pub mod affine;
7
8#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash)]
9pub enum TransformError {
10    InvalidTransform,
11    NonInvertibleTransform,
12}
13
14impl std::error::Error for TransformError {}
15
16impl Display for TransformError {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        match self {
19            TransformError::InvalidTransform => write!(f, "invalid transform"),
20            TransformError::NonInvertibleTransform => {
21                write!(
22                    f,
23                    "Non Invertible Transform, Forward transform not yet implemented "
24                )
25            }
26        }
27    }
28}
29
30pub trait Transform {
31    fn apply(&self, p: (f64, f64)) -> (f64, f64);
32    fn apply_inverse(&self, p: (f64, f64)) -> (f64, f64);
33    fn inverse_exists(&self) -> bool;
34}
35
36/// Composition of two transforms.  Specifically, derives transform2(transform1(image)).
37/// this is not equivalent to running the transforms separately, since the composition of the
38/// transforms occurs before sampling.  IE, running transforms separately incur a resample per
39/// transform, whereas composed Transforms only incur a single image resample.
40pub struct ComposedTransform<T: Transform> {
41    transform1: T,
42    transform2: T,
43}
44
45impl<T: Transform> Transform for ComposedTransform<T> {
46    fn apply(&self, p: (f64, f64)) -> (f64, f64) {
47        self.transform2.apply(self.transform1.apply(p))
48    }
49
50    fn apply_inverse(&self, p: (f64, f64)) -> (f64, f64) {
51        self.transform1
52            .apply_inverse(self.transform2.apply_inverse(p))
53    }
54
55    fn inverse_exists(&self) -> bool {
56        self.transform1.inverse_exists() && self.transform2.inverse_exists()
57    }
58}
59
60pub trait TransformExt<T: Transform>
61where
62    Self: Sized,
63{
64    /// Output type for the operation
65    type Output;
66
67    /// Transforms an image given the transformation matrix and output size.
68    /// Uses the source index coordinate space
69    /// Assume nearest-neighbour interpolation
70    fn transform(
71        &self,
72        transform: &T,
73        output_size: Option<(usize, usize)>,
74    ) -> Result<Self::Output, TransformError>;
75}
76
77#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
78struct Rect {
79    x: isize,
80    y: isize,
81    w: usize,
82    h: usize,
83}
84
85impl<T, U, V> TransformExt<V> for ArrayBase<U, Ix3>
86where
87    T: Copy + Clone + Num + NumAssignOps,
88    U: Data<Elem = T>,
89    V: Transform,
90{
91    type Output = Array<T, Ix3>;
92
93    fn transform(
94        &self,
95        transform: &V,
96        output_size: Option<(usize, usize)>,
97    ) -> Result<Self::Output, TransformError> {
98        let mut output = match output_size {
99            Some((r, c)) => Self::Output::zeros((r, c, self.shape()[2])),
100            None => Self::Output::zeros(self.raw_dim()),
101        };
102
103        for r in 0..output.shape()[0] {
104            for c in 0..output.shape()[1] {
105                let (x, y) = transform.apply_inverse((c as f64, r as f64));
106                let x = x.round() as isize;
107                let y = y.round() as isize;
108                if x >= 0
109                    && y >= 0
110                    && (x as usize) < self.shape()[1]
111                    && (y as usize) < self.shape()[0]
112                {
113                    output
114                        .slice_mut(s![r, c, ..])
115                        .assign(&self.slice(s![y, x, ..]));
116                }
117            }
118        }
119
120        Ok(output)
121    }
122}
123
124impl<T, U, C, V> TransformExt<V> for ImageBase<U, C>
125where
126    U: Data<Elem = T>,
127    T: Copy + Clone + Num + NumAssignOps,
128    C: ColourModel,
129    V: Transform,
130{
131    type Output = Image<T, C>;
132
133    fn transform(
134        &self,
135        transform: &V,
136        output_size: Option<(usize, usize)>,
137    ) -> Result<Self::Output, TransformError> {
138        let data = self.data.transform(transform, output_size)?;
139        let result = Self::Output::from_data(data).to_owned();
140        Ok(result)
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::affine;
147    use super::*;
148    use crate::core::colour_models::Gray;
149    use std::f64::consts::PI;
150
151    #[test]
152    fn translation() {
153        let src_data = vec![2.0, 0.0, 1.0, 0.0, 5.0, 0.0, 1.0, 2.0, 3.0];
154        let src = Image::<f64, Gray>::from_shape_data(3, 3, src_data);
155
156        let trans = affine::transform_from_2dmatrix(affine::translation(2.0, 1.0));
157
158        let res = src.transform(&trans, Some((3, 3)));
159        assert!(res.is_ok());
160        let res = res.unwrap();
161
162        let expected = vec![0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0];
163        let expected = Image::<f64, Gray>::from_shape_data(3, 3, expected);
164
165        assert_eq!(expected, res)
166    }
167
168    #[test]
169    fn rotate() {
170        let src = Image::<u8, Gray>::from_shape_data(5, 5, (0..25).collect());
171        let trans = affine::transform_from_2dmatrix(affine::rotate_around_centre(PI, (2.0, 2.0)));
172        let upside_down = src.transform(&trans, Some((5, 5))).unwrap();
173
174        let res = upside_down.transform(&trans, Some((5, 5))).unwrap();
175
176        assert_eq!(src, res);
177
178        let trans_2 =
179            affine::transform_from_2dmatrix(affine::rotate_around_centre(PI / 2.0, (2.0, 2.0)));
180        let trans_3 =
181            affine::transform_from_2dmatrix(affine::rotate_around_centre(-PI / 2.0, (2.0, 2.0)));
182
183        let upside_down_sideways = upside_down.transform(&trans_2, Some((5, 5))).unwrap();
184        let src_sideways = src.transform(&trans_3, Some((5, 5))).unwrap();
185
186        assert_eq!(upside_down_sideways, src_sideways);
187    }
188
189    #[test]
190    fn scale() {
191        let src = Image::<u8, Gray>::from_shape_data(4, 4, (0..16).collect());
192        let trans = affine::transform_from_2dmatrix(affine::scale(0.5, 2.0));
193        let res = src.transform(&trans, None).unwrap();
194
195        assert_eq!(res.rows(), 4);
196        assert_eq!(res.cols(), 4);
197    }
198}