1use crate::bounds::{Bbox, GeoTransform};
2use crate::crs::Crs;
3use ndarray::{Array2, Array3};
4
5#[derive(Debug, Clone)]
10pub enum PixelData {
11 U8(Array3<u8>),
13 U16(Array3<u16>),
15 I16(Array3<i16>),
17 U32(Array3<u32>),
19 I32(Array3<i32>),
21 F32(Array3<f32>),
23 F64(Array3<f64>),
25}
26
27impl PixelData {
28 pub fn bands(&self) -> usize {
30 self.shape()[0]
31 }
32
33 pub fn height(&self) -> usize {
35 self.shape()[1]
36 }
37
38 pub fn width(&self) -> usize {
40 self.shape()[2]
41 }
42
43 fn shape(&self) -> &[usize] {
44 match self {
45 Self::U8(a) => a.shape(),
46 Self::U16(a) => a.shape(),
47 Self::I16(a) => a.shape(),
48 Self::U32(a) => a.shape(),
49 Self::I32(a) => a.shape(),
50 Self::F32(a) => a.shape(),
51 Self::F64(a) => a.shape(),
52 }
53 }
54
55 pub fn to_f64(&self) -> Array3<f64> {
58 match self {
59 Self::U8(a) => a.mapv(|v| v as f64),
60 Self::U16(a) => a.mapv(|v| v as f64),
61 Self::I16(a) => a.mapv(|v| v as f64),
62 Self::U32(a) => a.mapv(|v| v as f64),
63 Self::I32(a) => a.mapv(|v| v as f64),
64 Self::F32(a) => a.mapv(|v| v as f64),
65 Self::F64(a) => a.clone(),
66 }
67 }
68
69 pub fn band_to_f64(&self, b: usize) -> Array2<f64> {
74 use ndarray::Axis;
75 macro_rules! conv {
76 ($a:expr) => {
77 $a.index_axis(Axis(0), b).mapv(|v| v as f64)
78 };
79 }
80 match self {
81 Self::U8(a) => conv!(a),
82 Self::U16(a) => conv!(a),
83 Self::I16(a) => conv!(a),
84 Self::U32(a) => conv!(a),
85 Self::I32(a) => conv!(a),
86 Self::F32(a) => conv!(a),
87 Self::F64(a) => a.index_axis(Axis(0), b).to_owned(),
88 }
89 }
90
91 pub fn as_f64(&self) -> std::borrow::Cow<'_, Array3<f64>> {
96 match self {
97 Self::F64(a) => std::borrow::Cow::Borrowed(a),
98 other => std::borrow::Cow::Owned(other.to_f64()),
99 }
100 }
101
102 pub fn into_f64(self) -> Array3<f64> {
104 match self {
105 Self::F64(a) => a,
106 other => other.to_f64(),
107 }
108 }
109
110 pub fn split_off_band(self, b: usize) -> Option<(Self, ndarray::Array2<bool>)> {
115 use ndarray::{s, Axis};
116 macro_rules! split {
117 ($arr:expr, $wrap:ident, $nz:expr) => {{
118 let bands = $arr.shape()[0];
119 if b >= bands {
120 return None;
121 }
122 let alpha = $arr.index_axis(Axis(0), b).mapv($nz);
123 let h = $arr.shape()[1];
124 let w = $arr.shape()[2];
125 let mut keep = Array3::default((bands - 1, h, w));
127 let mut k = 0usize;
128 for src_b in 0..bands {
129 if src_b == b {
130 continue;
131 }
132 keep.slice_mut(s![k, .., ..])
133 .assign(&$arr.index_axis(Axis(0), src_b));
134 k += 1;
135 }
136 Some((Self::$wrap(keep), alpha))
137 }};
138 }
139 match self {
140 Self::U8(a) => split!(a, U8, |v: u8| v != 0),
141 Self::U16(a) => split!(a, U16, |v: u16| v != 0),
142 Self::I16(a) => split!(a, I16, |v: i16| v != 0),
143 Self::U32(a) => split!(a, U32, |v: u32| v != 0),
144 Self::I32(a) => split!(a, I32, |v: i32| v != 0),
145 Self::F32(a) => split!(a, F32, |v: f32| v != 0.0),
146 Self::F64(a) => split!(a, F64, |v: f64| v != 0.0),
147 }
148 }
149}
150
151#[derive(Debug, Clone)]
153pub struct ImageData {
154 pub pixels: PixelData,
156 pub mask: Option<Array2<bool>>,
159 pub bbox: Bbox,
161 pub crs: Crs,
163 pub transform: GeoTransform,
165 pub nodata: Option<f64>,
167}
168
169#[cfg(test)]
170mod tests {
171 use super::*;
172 use ndarray::Array3;
173
174 #[test]
175 fn f64_views_borrow_and_consuming_conversion_reuses_strided_storage() {
176 let array = Array3::from_shape_vec((1, 2, 3), (0..6).map(f64::from).collect())
177 .unwrap()
178 .permuted_axes([0, 2, 1]);
179 let pointer = array.as_ptr();
180 let pixels = PixelData::F64(array);
181 assert!(matches!(pixels.as_f64(), std::borrow::Cow::Borrowed(_)));
182 assert_eq!(pixels.as_f64().as_ptr(), pointer);
183 let converted = pixels.into_f64();
184 assert_eq!(converted.as_ptr(), pointer);
185 assert_eq!(converted[[0, 2, 1]], 5.0);
186
187 let integer = PixelData::U16(Array3::from_elem((1, 2, 3), 1234));
188 assert!(matches!(integer.as_f64(), std::borrow::Cow::Owned(_)));
189 assert!(integer.into_f64().iter().all(|&v| v == 1234.0));
190 }
191
192 #[test]
193 fn pixel_data_shape_helpers() {
194 let arr = Array3::<u8>::zeros((3, 256, 256));
195 let p = PixelData::U8(arr);
196 assert_eq!(p.bands(), 3);
197 assert_eq!(p.height(), 256);
198 assert_eq!(p.width(), 256);
199 }
200
201 #[test]
202 fn to_f64_preserves_values() {
203 let arr = Array3::<u16>::from_shape_vec((1, 1, 2), vec![100, 200]).unwrap();
204 let p = PixelData::U16(arr);
205 let f = p.to_f64();
206 assert_eq!(f[[0, 0, 0]], 100.0);
207 assert_eq!(f[[0, 0, 1]], 200.0);
208 }
209}