ndarray_vision/processing/
sobel.rs

1use crate::core::*;
2use crate::processing::*;
3use core::mem::MaybeUninit;
4use core::ops::Neg;
5use ndarray::{prelude::*, s, DataMut, OwnedRepr, Zip};
6use num_traits::{cast::FromPrimitive, real::Real, Num, NumAssignOps};
7use std::marker::Sized;
8
9/// Runs the sobel operator on an image
10pub trait SobelExt
11where
12    Self: Sized,
13{
14    /// Type to output
15    type Output;
16    /// Returns the magnitude output of the sobel - an image of only lines
17    fn apply_sobel(&self) -> Result<Self::Output, Error>;
18
19    /// Returns the magntitude and rotation outputs for use in other algorithms
20    /// like the Canny edge detector. Rotation is in radians
21    fn full_sobel(&self) -> Result<(Self::Output, Self::Output), Error>;
22}
23
24fn get_edge_images<T, U>(mat: &ArrayBase<U, Ix3>) -> Result<(Array3<T>, Array3<T>), Error>
25where
26    U: DataMut<Elem = T>,
27    T: Copy + Clone + Num + NumAssignOps + Neg<Output = T> + FromPrimitive + Real,
28{
29    let v_temp: Array3<T> = SobelFilter::build_with_params(Orientation::Vertical).unwrap();
30    let h_temp: Array3<T> = SobelFilter::build_with_params(Orientation::Horizontal).unwrap();
31    let shape = (v_temp.shape()[0], v_temp.shape()[1], mat.shape()[2]);
32    let mut h_kernel = Array3::<T>::uninit(shape);
33    let mut v_kernel = Array3::<T>::uninit(shape);
34    for i in 0..mat.dim().2 {
35        h_temp
36            .slice(s![.., .., 0])
37            .assign_to(h_kernel.slice_mut(s![.., .., i]));
38        v_temp
39            .slice(s![.., .., 0])
40            .assign_to(v_kernel.slice_mut(s![.., .., i]));
41    }
42    let h_kernel = unsafe { h_kernel.assume_init() };
43    let v_kernel = unsafe { v_kernel.assume_init() };
44    let h_deriv = mat.conv2d(h_kernel.view())?;
45    let v_deriv = mat.conv2d(v_kernel.view())?;
46
47    Ok((h_deriv, v_deriv))
48}
49
50impl<T, U> SobelExt for ArrayBase<U, Ix3>
51where
52    U: DataMut<Elem = T>,
53    T: Copy + Clone + Num + NumAssignOps + Neg<Output = T> + FromPrimitive + Real,
54{
55    type Output = ArrayBase<OwnedRepr<T>, Ix3>;
56
57    fn apply_sobel(&self) -> Result<Self::Output, Error> {
58        let (h_deriv, v_deriv) = get_edge_images(self)?;
59        let res_shape = h_deriv.dim();
60        let mut result = Self::Output::uninit(res_shape);
61        for r in 0..res_shape.0 {
62            for c in 0..res_shape.1 {
63                for channel in 0..res_shape.2 {
64                    let temp = (h_deriv[[r, c, channel]].powi(2)
65                        + v_deriv[[r, c, channel]].powi(2))
66                    .sqrt();
67                    unsafe {
68                        *result.uget_mut([r, c, channel]) = MaybeUninit::new(temp);
69                    }
70                }
71            }
72        }
73        Ok(unsafe { result.assume_init() })
74    }
75
76    fn full_sobel(&self) -> Result<(Self::Output, Self::Output), Error> {
77        let (h_deriv, v_deriv) = get_edge_images(self)?;
78
79        let mut magnitude = h_deriv.mapv(|x| x.powi(2)) + v_deriv.mapv(|x| x.powi(2));
80        magnitude.mapv_inplace(|x| x.sqrt());
81
82        let dim = h_deriv.dim();
83        let mut rotation = Array3::uninit((dim.0, dim.1, dim.2));
84        Zip::from(&mut rotation)
85            .and(&h_deriv)
86            .and(&v_deriv)
87            .for_each(|r, &h, &v| *r = MaybeUninit::new(h.atan2(v)));
88
89        let rotation = unsafe { rotation.assume_init() };
90
91        Ok((magnitude, rotation))
92    }
93}
94
95impl<T, U, C> SobelExt for ImageBase<U, C>
96where
97    U: DataMut<Elem = T>,
98    T: Copy + Clone + Num + NumAssignOps + Neg<Output = T> + FromPrimitive + Real,
99    C: ColourModel,
100{
101    type Output = Image<T, C>;
102
103    fn apply_sobel(&self) -> Result<Self::Output, Error> {
104        let data = self.data.apply_sobel()?;
105        Ok(Image::from_data(data))
106    }
107
108    fn full_sobel(&self) -> Result<(Self::Output, Self::Output), Error> {
109        self.data
110            .full_sobel()
111            .map(|(m, r)| (Image::from_data(m), Image::from_data(r)))
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use approx::*;
119
120    #[test]
121    fn simple() {
122        let mut image: Image<f64, Gray> = ImageBase::new(11, 11);
123        image.data.slice_mut(s![4..7, 4..7, ..]).fill(1.0);
124        image.data.slice_mut(s![3..8, 5, ..]).fill(1.0);
125        image.data.slice_mut(s![5, 3..8, ..]).fill(1.0);
126
127        let sobel = image.full_sobel().unwrap();
128
129        // Did a calculation of sobel_mag[1..9, 1..9, ..] in a spreadsheet
130        #[rustfmt::skip]
131        let mag = vec![0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
132            0.0, 0.0, 0.0, 1.41421356237301, 2.0, 1.41421356237301, 0.0, 0.0, 0.0,
133            0.0, 0.0, 1.41421356237301, 4.24264068711929, 4.0, 4.24264068711929, 1.4142135623731, 0.0, 0.0, 
134            0.0, 1.4142135623731, 4.24264068711929, 4.24264068711929, 2.0, 4.24264068711929, 4.24264068711929, 1.4142135623731, 0.0,
135            0.0, 2.0, 4.0, 2.0, 0.0, 2.0, 4.0, 2.0, 0.0,
136            0.0, 1.4142135623731, 4.24264068711929, 4.24264068711929, 2.0, 4.24264068711929, 4.24264068711929, 1.4142135623731, 0.0,
137            0.0, 0.0, 1.4142135623731, 4.24264068711929, 4.0, 4.24264068711929, 1.4142135623731, 0.0,
138            0.0, 0.0, 0.0, 0.0, 1.4142135623731, 2.0, 1.4142135623731, 0.0, 0.0,
139            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
140        ];
141
142        let mag = Array::from_shape_vec((9, 9), mag).unwrap();
143
144        assert_abs_diff_eq!(sobel.0.data.slice(s![1..10, 1..10, 0]), mag, epsilon = 1e-5);
145
146        let only_mag = image.apply_sobel().unwrap();
147        assert_abs_diff_eq!(sobel.0.data, only_mag.data);
148
149        // Did a calculation of sobel_rot[1..9, 1..9, ..] in a spreadsheet
150        #[rustfmt::skip]
151        let rot = vec![0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,
152                       0.00000000000000,0.00000000000000,0.00000000000000,-2.35619449019234,3.14159265358979,2.35619449019234,0.00000000000000,0.00000000000000,0.00000000000000,
153                       0.00000000000000,0.00000000000000,-2.35619449019234,-2.35619449019234,3.14159265358979,2.35619449019234,2.35619449019234,0.00000000000000,0.00000000000000,
154                       0.00000000000000,-2.35619449019234,-2.35619449019234,-2.35619449019234,3.14159265358979,2.35619449019234,2.35619449019234,2.35619449019234,0.00000000000000,
155                       0.00000000000000,-1.57079632679490,-1.57079632679490,-1.57079632679490,0.00000000000000,1.57079632679490,1.57079632679490,1.57079632679490,0.00000000000000,
156                       0.00000000000000,-0.78539816339745,-0.78539816339745,-0.78539816339745,0.00000000000000,0.78539816339745,0.78539816339745,0.78539816339745,0.00000000000000,
157                       0.00000000000000,0.00000000000000,-0.78539816339745,-0.78539816339745,0.00000000000000,0.78539816339745,0.78539816339745,0.00000000000000,0.00000000000000,
158                       0.00000000000000,0.00000000000000,0.00000000000000,-0.78539816339745,0.00000000000000,0.78539816339745,0.00000000000000,0.00000000000000,0.00000000000000,
159                       0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000,0.00000000000000];
160        let rot = Array::from_shape_vec((9, 9), rot).unwrap();
161
162        assert_abs_diff_eq!(sobel.1.data.slice(s![1..10, 1..10, 0]), rot, epsilon = 1e-5);
163    }
164}
ndarray_vision/processing/sobel.rs

ndarray_vision/processing/
sobel.rs
