Skip to main content

async_geotiff/
pixel.rs

1use crate::bounds::{Bbox, GeoTransform};
2use crate::crs::Crs;
3use ndarray::{Array2, Array3};
4
5/// Decoded pixel buffer in `(bands, height, width)` planar layout.
6///
7/// Reading, warping, and mosaic reduction retain this planar layout. Renderers
8/// interleave selected channels when producing encoded image bytes.
9#[derive(Debug, Clone)]
10pub enum PixelData {
11    /// Unsigned 8-bit samples in `(bands, height, width)` order.
12    U8(Array3<u8>),
13    /// Unsigned 16-bit samples in `(bands, height, width)` order.
14    U16(Array3<u16>),
15    /// Signed 16-bit samples in `(bands, height, width)` order.
16    I16(Array3<i16>),
17    /// Unsigned 32-bit samples in `(bands, height, width)` order.
18    U32(Array3<u32>),
19    /// Signed 32-bit samples in `(bands, height, width)` order.
20    I32(Array3<i32>),
21    /// 32-bit floating-point samples in `(bands, height, width)` order.
22    F32(Array3<f32>),
23    /// 64-bit floating-point samples in `(bands, height, width)` order.
24    F64(Array3<f64>),
25}
26
27impl PixelData {
28    /// Number of sample bands (the first array axis).
29    pub fn bands(&self) -> usize {
30        self.shape()[0]
31    }
32
33    /// Number of pixel rows (the second array axis).
34    pub fn height(&self) -> usize {
35        self.shape()[1]
36    }
37
38    /// Number of pixel columns (the third array axis).
39    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    /// View every pixel as f64 — used by the render hot path.
56    /// Allocates one f64 array; caller should cache if iterating.
57    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    /// Convert a single band `b` to an owned `Array2<f64>`. Cheaper than
70    /// [`to_f64`](Self::to_f64) when only a few of many bands are rendered: it
71    /// touches one band, not the whole `(bands, H, W)` cube. Panics if `b` is
72    /// out of range — callers validate `bidx` first.
73    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    /// Like [`to_f64`](Self::to_f64) but borrows the buffer for the
92    /// already-`f64` case (`Cow::Borrowed`) instead of cloning a full
93    /// `bands×H×W` array. The scalar cross-CRS warp only reads the
94    /// promoted array, so f64 COGs pay zero extra allocation.
95    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    /// Consume the pixels, reusing the allocation when already F64.
103    pub fn into_f64(self) -> Array3<f64> {
104        match self {
105            Self::F64(a) => a,
106            other => other.to_f64(),
107        }
108    }
109
110    /// Split off band `b` as an `Array2<bool>` (truthy = non-zero) and
111    /// return the remaining planar `(bands-1, H, W)` buffer. Used by the
112    /// CogReader to promote an `ExtraSamples=Alpha` band to a validity
113    /// mask after stitching, mirroring rio-tiler's alpha-band priority.
114    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                // Drop band b: select all rows except b.
126                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/// The single data-bearing type passed between io → warp → render.
152#[derive(Debug, Clone)]
153pub struct ImageData {
154    /// Planar sample array. Readers preserve raw values; SCALE/OFFSET is not applied.
155    pub pixels: PixelData,
156    /// Optional `(height, width)` validity mask on the same grid as `pixels`.
157    /// `true` means valid; `None` imposes no mask, though `nodata` still applies.
158    pub mask: Option<Array2<bool>>,
159    /// Axis-aligned extent of this image in `crs` coordinates.
160    pub bbox: Bbox,
161    /// Coordinate reference system for `bbox` and `transform`.
162    pub crs: Crs,
163    /// Maps pixel-corner coordinates on this image's grid into `crs` coordinates.
164    pub transform: GeoTransform,
165    /// Optional raw-sample marker for invalid pixels, before SCALE/OFFSET.
166    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}