ndarray_vision/core/
image.rs

1use crate::core::colour_models::*;
2use crate::core::traits::PixelBound;
3use ndarray::prelude::*;
4use ndarray::{s, Data, DataMut, OwnedRepr, RawDataClone, ViewRepr};
5use num_traits::cast::{FromPrimitive, NumCast};
6use num_traits::Num;
7use std::{fmt, hash, marker::PhantomData};
8
9pub type Image<T, C> = ImageBase<OwnedRepr<T>, C>;
10pub type ImageView<'a, T, C> = ImageBase<ViewRepr<&'a T>, C>;
11
12/// Basic structure containing an image.
13pub struct ImageBase<T, C>
14where
15    C: ColourModel,
16    T: Data,
17{
18    /// Images are always going to be 3D to handle rows, columns and colour
19    /// channels
20    ///
21    /// This should allow for max compatibility with maths ops in ndarray.
22    /// Caution should be taken if performing any operations that change the
23    /// number of channels in an image as this may cause other functionality to
24    /// perform incorrectly. Use conversions to one of the `Generic` colour models
25    /// instead.
26    pub data: ArrayBase<T, Ix3>,
27    /// Representation of how colour is encoded in the image
28    pub(crate) model: PhantomData<C>,
29}
30
31impl<T, U, C> ImageBase<U, C>
32where
33    U: Data<Elem = T>,
34    T: Copy + Clone + FromPrimitive + Num + NumCast + PixelBound,
35    C: ColourModel,
36{
37    /// Converts image into a different type - doesn't scale to new pixel bounds
38    pub fn into_type<T2>(self) -> Image<T2, C>
39    where
40        T2: Copy + Clone + FromPrimitive + Num + NumCast + PixelBound,
41    {
42        let rescale = |x: &T| {
43            let scaled = normalise_pixel_value(*x)
44                * (T2::max_pixel() - T2::min_pixel())
45                    .to_f64()
46                    .unwrap_or(0.0f64);
47            T2::from_f64(scaled).unwrap_or_else(T2::zero) + T2::min_pixel()
48        };
49        let data = self.data.map(rescale);
50        Image::<_, C>::from_data(data)
51    }
52}
53
54impl<S, T, C> ImageBase<S, C>
55where
56    S: Data<Elem = T>,
57    T: Clone,
58    C: ColourModel,
59{
60    pub fn to_owned(&self) -> Image<T, C> {
61        Image {
62            data: self.data.to_owned(),
63            model: PhantomData,
64        }
65    }
66
67    /// Given the shape of the image and a data vector create an image. If
68    /// the data sizes don't match a zero filled image will be returned instead
69    /// of panicking
70    pub fn from_shape_data(rows: usize, cols: usize, data: Vec<T>) -> Image<T, C> {
71        let data = Array3::from_shape_vec((rows, cols, C::channels()), data).unwrap();
72
73        Image {
74            data,
75            model: PhantomData,
76        }
77    }
78}
79
80impl<T, C> Image<T, C>
81where
82    T: Clone + Num,
83    C: ColourModel,
84{
85    /// Construct a new image filled with zeros using the given dimensions and
86    /// a colour model
87    pub fn new(rows: usize, columns: usize) -> Self {
88        Image {
89            data: Array3::zeros((rows, columns, C::channels())),
90            model: PhantomData,
91        }
92    }
93}
94
95impl<T, U, C> ImageBase<T, C>
96where
97    T: Data<Elem = U>,
98    C: ColourModel,
99{
100    /// Construct the image from a given Array3
101    pub fn from_data(data: ArrayBase<T, Ix3>) -> Self {
102        Self {
103            data,
104            model: PhantomData,
105        }
106    }
107    /// Returns the number of rows in an image
108    pub fn rows(&self) -> usize {
109        self.data.shape()[0]
110    }
111    /// Returns the number of channels in an image
112    pub fn cols(&self) -> usize {
113        self.data.shape()[1]
114    }
115
116    /// Convenience method to get number of channels
117    pub fn channels(&self) -> usize {
118        C::channels()
119    }
120
121    /// Get a view of all colour channels at a pixels location
122    pub fn pixel(&self, row: usize, col: usize) -> ArrayView<U, Ix1> {
123        self.data.slice(s![row, col, ..])
124    }
125
126    pub fn into_type_raw<C2>(self) -> ImageBase<T, C2>
127    where
128        C2: ColourModel,
129    {
130        assert_eq!(C2::channels(), C::channels());
131        ImageBase::<T, C2>::from_data(self.data)
132    }
133}
134
135impl<T, U, C> ImageBase<T, C>
136where
137    T: DataMut<Elem = U>,
138    C: ColourModel,
139{
140    /// Get a mutable view of a pixels colour channels given a location
141    pub fn pixel_mut(&mut self, row: usize, col: usize) -> ArrayViewMut<U, Ix1> {
142        self.data.slice_mut(s![row, col, ..])
143    }
144}
145
146impl<T, U, C> fmt::Debug for ImageBase<U, C>
147where
148    U: Data<Elem = T>,
149    T: fmt::Debug,
150    C: ColourModel,
151{
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        write!(f, "ColourModel={:?} Data={:?}", self.model, self.data)?;
154        Ok(())
155    }
156}
157
158impl<T, U, C> PartialEq<ImageBase<U, C>> for ImageBase<U, C>
159where
160    U: Data<Elem = T>,
161    T: PartialEq,
162    C: ColourModel,
163{
164    fn eq(&self, other: &Self) -> bool {
165        self.model == other.model && self.data == other.data
166    }
167}
168
169impl<S, C> Clone for ImageBase<S, C>
170where
171    S: RawDataClone + Data,
172    C: ColourModel,
173{
174    fn clone(&self) -> Self {
175        Self {
176            data: self.data.clone(),
177            model: PhantomData,
178        }
179    }
180
181    fn clone_from(&mut self, other: &Self) {
182        self.data.clone_from(&other.data)
183    }
184}
185
186impl<'a, S, C> hash::Hash for ImageBase<S, C>
187where
188    S: Data,
189    S::Elem: hash::Hash,
190    C: ColourModel,
191{
192    fn hash<H: hash::Hasher>(&self, state: &mut H) {
193        self.model.hash(state);
194        self.data.hash(state);
195    }
196}
197
198/// Returns a normalised pixel value or 0 if it can't convert the types.
199/// This should never fail if your types are good.
200pub fn normalise_pixel_value<T>(t: T) -> f64
201where
202    T: PixelBound + Num + NumCast,
203{
204    let numerator = (t + T::min_pixel()).to_f64();
205    let denominator = (T::max_pixel() - T::min_pixel()).to_f64();
206
207    let numerator = numerator.unwrap_or(0.0f64);
208    let denominator = denominator.unwrap_or(1.0f64);
209
210    numerator / denominator
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use ndarray::arr1;
217
218    #[test]
219    fn image_consistency_checks() {
220        let i = Image::<u8, RGB>::new(1, 2);
221        assert_eq!(i.rows(), 1);
222        assert_eq!(i.cols(), 2);
223        assert_eq!(i.channels(), 3);
224        assert_eq!(i.channels(), i.data.shape()[2]);
225    }
226
227    #[test]
228    fn image_type_conversion() {
229        let mut i = Image::<u8, RGB>::new(1, 1);
230        i.pixel_mut(0, 0)
231            .assign(&arr1(&[u8::max_value(), 0, u8::max_value() / 3]));
232        let t: Image<u16, RGB> = i.into_type();
233        assert_eq!(
234            t.pixel(0, 0),
235            arr1(&[u16::max_value(), 0, u16::max_value() / 3])
236        );
237    }
238}