map_engine/raster/
pixels.rs

1//! Raw and Styled pixels.
2use crate::errors::MapEngineError;
3use crate::{
4    cmap::{Composite, HandleGet},
5    tiles::TILE_SIZE,
6};
7use gdal::raster::GdalType;
8use ndarray::{Array, Array3, Axis};
9use num_traits::{Num, NumCast};
10use std::fs::File;
11use std::io::{BufWriter, Write};
12use std::path::Path;
13
14/// Raw pixels read from a raster file.
15pub struct RawPixels<P>
16where
17    P: GdalType + Copy + Num + NumCast,
18    P: ndarray::ScalarOperand,
19{
20    data: Array3<P>,
21    driver: Box<dyn driver::Style<P>>,
22}
23
24impl<P> RawPixels<P>
25where
26    P: GdalType + Copy + Num + NumCast,
27    P: ndarray::ScalarOperand,
28{
29    pub(super) fn new(data: Array3<P>, driver: &str) -> Self {
30        match driver {
31            driver::MBTILES => {
32                let driver = Box::new(driver::Mbtile {});
33
34                Self { data, driver }
35            }
36            _ => {
37                let driver = Box::new(driver::Generic {});
38                Self { data, driver }
39            }
40        }
41    }
42
43    /// Apply a colour map to an array.
44    ///
45    /// # Arguments
46    ///
47    /// * `cmap` - A colour composite that maps pixel values to RGBA.
48    /// * `no_data_values` - Pixel values to be set as fully transparent.
49    pub fn style(
50        self,
51        cmap: Composite,
52        no_data_values: Vec<f64>,
53    ) -> Result<StyledPixels, MapEngineError>
54    where
55        P: std::fmt::Debug,
56    {
57        self.driver.style(&self, cmap, no_data_values)
58    }
59
60    pub fn as_array(&self) -> &Array3<P> {
61        &self.data
62    }
63}
64
65/// Pixels styled using the [`RawPixels::style`] method.
66pub struct StyledPixels {
67    data: Array3<u8>,
68    driver: driver::Driver,
69}
70
71impl StyledPixels {
72    pub fn new(data: Array3<u8>, driver: driver::Driver) -> Self {
73        Self { data, driver }
74    }
75
76    pub fn into_png(self) -> Result<Vec<u8>, MapEngineError> {
77        let mut buffer = Vec::<u8>::new();
78        let mut w: BufWriter<&mut Vec<u8>> = BufWriter::new(buffer.as_mut());
79        {
80            let mut encoder = png::Encoder::new(&mut w, TILE_SIZE as u32, TILE_SIZE as u32);
81            encoder.set_color(png::ColorType::Rgba);
82            encoder.set_depth(png::BitDepth::Eight);
83            let mut writer = encoder.write_header()?;
84            match self.driver {
85                driver::Driver::Generic => writer.write_image_data(&self.data.into_raw_vec())?,
86                driver::Driver::Mbtile => {
87                    writer.write_image_data(&self.data.into_iter().collect::<Vec<u8>>()[..])?
88                }
89            }
90        }
91        w.flush()?;
92        drop(w);
93        Ok(buffer)
94    }
95
96    /// Write a styled tile to disk.
97    ///
98    /// # Arguments
99    ///
100    /// * `out_path` - Path were to write the tile.
101    pub fn write_to_disk(self, out_path: &Path) -> Result<(), MapEngineError> {
102        let png_data = self.into_png()?;
103        let mut file = File::create(out_path)?;
104        file.write_all(&png_data[..])?;
105        Ok(())
106    }
107
108    #[allow(dead_code)]
109    fn as_array(&self) -> &Array3<u8> {
110        &self.data
111    }
112
113    #[allow(dead_code)]
114    fn into_array(self) -> Array3<u8> {
115        self.data
116    }
117}
118
119impl Default for StyledPixels {
120    fn default() -> Self {
121        Self::new(
122            Array::zeros((4, TILE_SIZE, TILE_SIZE)),
123            driver::Driver::Generic,
124        )
125    }
126}
127
128pub(crate) mod driver {
129    use super::*;
130
131    pub struct Mbtile;
132    pub struct Generic;
133    pub const MBTILES: &str = "MBTiles";
134    pub enum Driver {
135        Mbtile,
136        Generic,
137    }
138
139    pub trait Style<P>
140    where
141        P: GdalType + Copy + Num + NumCast,
142        P: ndarray::ScalarOperand,
143    {
144        fn style(
145            &self,
146            arr: &RawPixels<P>,
147            cmap: Composite,
148            no_data_values: Vec<f64>,
149        ) -> Result<StyledPixels, MapEngineError>;
150    }
151
152    impl<P> Style<P> for Generic
153    where
154        P: GdalType + Copy + Num + NumCast,
155        P: ndarray::ScalarOperand,
156    {
157        fn style(
158            &self,
159            raw: &RawPixels<P>,
160            cmap: Composite,
161            no_data_values: Vec<f64>,
162        ) -> Result<StyledPixels, MapEngineError> {
163            let arr_f64 = unsafe { raw.data.raw_view().cast::<f64>().deref_into_view() };
164            let v = if cmap.is_contiguous() {
165                arr_f64
166                    .lanes(Axis(0))
167                    .into_iter()
168                    .map(|v| cmap.get(v.as_slice().unwrap(), Some(&no_data_values)))
169                    .flatten()
170                    .collect()
171            } else {
172                arr_f64
173                    .lanes(Axis(0))
174                    .into_iter()
175                    .map(|v| cmap.get(&v.to_vec(), Some(&no_data_values)))
176                    .flatten()
177                    .collect()
178            };
179            let arr = unsafe {
180                Array::from_shape_vec_unchecked((raw.data.shape()[1], raw.data.shape()[2], 4), v)
181            };
182            Ok(StyledPixels::new(arr, driver::Driver::Generic))
183        }
184    }
185
186    impl<P> Style<P> for Mbtile
187    where
188        P: GdalType + Copy + Num + NumCast,
189        P: ndarray::ScalarOperand,
190    {
191        fn style(
192            &self,
193            raw: &RawPixels<P>,
194            _: Composite,
195            _: Vec<f64>,
196        ) -> Result<StyledPixels, MapEngineError> {
197            Ok(StyledPixels::new(
198                // TODO: Cast directly to u8 using unsafe rust
199                raw.data.mapv(|v| v.to_u8().expect("Cannot parse to u8")),
200                driver::Driver::Mbtile,
201            ))
202        }
203    }
204}
205
206#[cfg(test)]
207mod test {
208    use super::*;
209    use crate::cmap::viridis;
210    use crate::raster::Raster;
211    use crate::tiles::Tile;
212    use gdal::Dataset;
213    use ndarray::{arr3, s};
214    use std::path::PathBuf;
215
216    #[test]
217    fn test_write() {
218        let path = PathBuf::from("src/tests/data/chile_optimised.tif");
219        let src = Dataset::open(&path).unwrap();
220        let band = src.rasterband(1).unwrap();
221        let no_data_value = band.no_data_value().unwrap_or(0.0);
222        let (vmin, vmax) = (0.0, 27412.0); // Extracted from raster
223
224        let raster = Raster::new(path).unwrap();
225        let tile = Tile::new(304, 624, 10);
226        let arr: RawPixels<f64> = raster.read_tile(&tile, Some(&[1]), None).unwrap();
227        let styled = arr
228            .style(
229                Composite::new_gradient(vmin, vmax, &viridis),
230                vec![no_data_value],
231            )
232            .unwrap();
233        let out_path = format!("/tmp/tile_{}_{}_{}.png", tile.x, tile.y, tile.z);
234        let out_path = Path::new(&out_path);
235        styled.write_to_disk(out_path).unwrap();
236    }
237
238    #[test]
239    fn test_style_tile() {
240        // let arr = RawPixels::new(arr3(&[[[0.0, 0.25], [0.5, 1.]], [[0.0, 0.0], [0.0, 0.0]]]));
241        let arr = RawPixels::new(arr3(&[[[0.0, 0.25], [0.5, 1.]]]), "");
242        let styled = arr
243            .style(Composite::new_gradient(0.0, 1., &viridis), vec![0.25])
244            .unwrap();
245        assert_eq!(
246            styled.as_array().slice(s![0, 0, ..]).to_vec(),
247            [68, 1, 84, 255]
248        );
249        assert_eq!(styled.as_array()[[0, 1, 3]], 0); // 0.25 is transparent
250    }
251
252    #[test]
253    fn test_style_rgb_tile() {
254        let arr = RawPixels::new(
255            arr3(&[
256                [[1.0, 0.0], [0.0, 0.25]],
257                [[0.0, 1.0], [0.0, 0.0]],
258                [[0.0, 0.0], [1.0, 0.0]],
259            ]),
260            "",
261        );
262        let styled = arr
263            .style(
264                Composite::new_rgb(vec![0.0, 0.0, 0.0], vec![1.0, 1.0, 1.0]),
265                vec![0.25, 0.25, 0.25],
266            )
267            .unwrap();
268        assert_eq!(
269            styled.as_array().slice(s![0, 0, ..]).to_vec(),
270            [255, 0, 0, 255]
271        ); // Red pixel
272        assert_eq!(
273            styled.as_array().slice(s![0, 1, ..]).to_vec(),
274            [0, 255, 0, 255]
275        ); // Green pixel
276        assert_eq!(
277            styled.as_array().slice(s![1, 0, ..]).to_vec(),
278            [0, 0, 255, 255]
279        ); // Blue pixel
280        assert_eq!(styled.as_array().slice(s![1, 1, ..]).to_vec()[3], 0); // Transparent pixel
281    }
282
283    #[test]
284    #[should_panic]
285    fn test_style_tile_with_gradient_fails() {
286        let arr = RawPixels::new(arr3(&[[[1.0, 0.0], [0.0, 0.25]]]), "");
287        // TODO: Should style panic, return a Result, or provide a good default as a fallback?
288        arr.style(
289            Composite::new_gradient(0.0, 1., &viridis),
290            vec![0.25, 0.25, 0.25], // We should provide 1 value for a Gradient
291        )
292        .unwrap();
293    }
294
295    #[test]
296    #[should_panic]
297    fn test_style_rgb_tile_fails() {
298        let arr = RawPixels::new(
299            arr3(&[
300                [[1.0, 0.0], [0.0, 0.25]],
301                [[0.0, 1.0], [0.0, 0.0]],
302                [[0.0, 0.0], [1.0, 0.0]],
303            ]),
304            "",
305        );
306        arr.style(
307            Composite::new_rgb(vec![0.0, 0.0, 0.0], vec![1.0, 1.0, 1.0]),
308            vec![0.25], // We should provide 3 values for a RBG composite
309        )
310        .unwrap();
311    }
312}
map_engine/raster/pixels.rs

map_engine/raster/
pixels.rs
