surtgis_colormap/
encode.rs

1//! RGBA pixel buffer + PNG encoder.
2//!
3//! Adds an in-memory `RgbaImage` type with alpha-over and multiply blends,
4//! plus native PNG output via the `image` crate. The native PNG path is
5//! gated on `cfg(not(target_arch = "wasm32"))` so WASM builds skip the
6//! `image` dependency; on WASM the JS side encodes the raw RGBA buffer.
7//!
8//! This is the in-tree home for what was previously expected to live in
9//! `surtgis-relief`. Hosting it here lets every other crate that produces
10//! RGBA buffers (hypsometric, curvature previews, fluvial maps) save PNG
11//! without pulling a separate dep.
12
13use surtgis_core::raster::Raster;
14use thiserror::Error;
15
16#[cfg(not(target_arch = "wasm32"))]
17use std::path::Path;
18
19#[derive(Debug, Error)]
20pub enum EncodeError {
21    #[error("shape mismatch: pixels.len()={got}, expected={expected} ({width}x{height}x4)")]
22    Shape {
23        width: usize,
24        height: usize,
25        got: usize,
26        expected: usize,
27    },
28    #[cfg(not(target_arch = "wasm32"))]
29    #[error("png encode: {0}")]
30    Png(#[from] image::ImageError),
31    #[cfg(not(target_arch = "wasm32"))]
32    #[error("io: {0}")]
33    Io(#[from] std::io::Error),
34}
35
36/// Result alias scoped to this module.
37pub type Result<T> = std::result::Result<T, EncodeError>;
38
39/// Row-major 8-bit RGBA pixel buffer.
40///
41/// Layout: `pixels[(row * width + col) * 4 + channel]`, with `channel`
42/// 0=R, 1=G, 2=B, 3=A. Matches the convention produced by
43/// [`crate::raster_to_rgba`].
44#[derive(Debug, Clone)]
45pub struct RgbaImage {
46    pub width: usize,
47    pub height: usize,
48    pub pixels: Vec<u8>,
49}
50
51impl RgbaImage {
52    /// Construct from a pre-built RGBA buffer. Returns an error if
53    /// `pixels.len() != width * height * 4`.
54    pub fn from_rgba(width: usize, height: usize, pixels: Vec<u8>) -> Result<Self> {
55        let expected = width * height * 4;
56        if pixels.len() != expected {
57            return Err(EncodeError::Shape {
58                width,
59                height,
60                got: pixels.len(),
61                expected,
62            });
63        }
64        Ok(Self {
65            width,
66            height,
67            pixels,
68        })
69    }
70
71    /// Construct from a single-channel intensity raster in `[0, 1]`.
72    ///
73    /// Values outside `[0, 1]` are clamped. NaN cells render fully
74    /// transparent black. Output is greyscale (R = G = B = scaled value),
75    /// alpha 255 for finite cells.
76    pub fn from_intensity(intensity: &Raster<f64>) -> Self {
77        let (rows, cols) = intensity.shape();
78        let mut pixels = vec![0u8; rows * cols * 4];
79        for (i, v) in intensity.data().iter().enumerate() {
80            let off = i * 4;
81            if v.is_nan() {
82                // transparent black (already zeroed)
83                continue;
84            }
85            let clamped = v.clamp(0.0, 1.0);
86            let g = (clamped * 255.0).round() as u8;
87            pixels[off] = g;
88            pixels[off + 1] = g;
89            pixels[off + 2] = g;
90            pixels[off + 3] = 255;
91        }
92        Self {
93            width: cols,
94            height: rows,
95            pixels,
96        }
97    }
98
99    /// Alpha-over composite: paint `top` on top of `self`, modulated by
100    /// `opacity` in `[0, 1]`. `self` is mutated in place.
101    ///
102    /// `top.width` and `top.height` must match.
103    pub fn over(&mut self, top: &RgbaImage, opacity: f64) -> Result<()> {
104        if top.width != self.width || top.height != self.height {
105            return Err(EncodeError::Shape {
106                width: top.width,
107                height: top.height,
108                got: top.pixels.len(),
109                expected: self.width * self.height * 4,
110            });
111        }
112        let op = opacity.clamp(0.0, 1.0);
113        let n = self.pixels.len() / 4;
114        for i in 0..n {
115            let off = i * 4;
116            let ta = (top.pixels[off + 3] as f64 / 255.0) * op;
117            if ta <= 0.0 {
118                continue;
119            }
120            let inv = 1.0 - ta;
121            for c in 0..3 {
122                let s = self.pixels[off + c] as f64;
123                let t = top.pixels[off + c] as f64;
124                self.pixels[off + c] = (t * ta + s * inv).round().clamp(0.0, 255.0) as u8;
125            }
126            // Alpha: standard over compositing.
127            let sa = self.pixels[off + 3] as f64 / 255.0;
128            let out_a = ta + sa * inv;
129            self.pixels[off + 3] = (out_a * 255.0).round().clamp(0.0, 255.0) as u8;
130        }
131        Ok(())
132    }
133
134    /// Multiply blend: `out_rgb = self_rgb * lerp(white, top_rgb, opacity)`,
135    /// per channel, normalised to `[0, 1]`. Useful for laying shadows over a
136    /// colored base. Alpha of `self` is preserved.
137    pub fn multiply(&mut self, top: &RgbaImage, opacity: f64) -> Result<()> {
138        if top.width != self.width || top.height != self.height {
139            return Err(EncodeError::Shape {
140                width: top.width,
141                height: top.height,
142                got: top.pixels.len(),
143                expected: self.width * self.height * 4,
144            });
145        }
146        let op = opacity.clamp(0.0, 1.0);
147        let n = self.pixels.len() / 4;
148        for i in 0..n {
149            let off = i * 4;
150            for c in 0..3 {
151                let s = self.pixels[off + c] as f64 / 255.0;
152                let t = top.pixels[off + c] as f64 / 255.0;
153                // lerp(1.0, t, op) = (1-op) + op*t
154                let factor = (1.0 - op) + op * t;
155                self.pixels[off + c] = (s * factor * 255.0).round().clamp(0.0, 255.0) as u8;
156            }
157        }
158        Ok(())
159    }
160
161    /// Encode the buffer as PNG bytes. Native target only.
162    #[cfg(not(target_arch = "wasm32"))]
163    pub fn to_png_bytes(&self) -> Result<Vec<u8>> {
164        rgba_to_png_bytes(self.width as u32, self.height as u32, &self.pixels)
165    }
166
167    /// Write the buffer to `path` as a PNG. Native target only.
168    #[cfg(not(target_arch = "wasm32"))]
169    pub fn save_png<P: AsRef<Path>>(&self, path: P) -> Result<()> {
170        save_png(path, self.width as u32, self.height as u32, &self.pixels)
171    }
172}
173
174/// Encode an arbitrary RGBA byte buffer as PNG. Native target only.
175///
176/// `rgba.len()` must equal `width * height * 4`.
177#[cfg(not(target_arch = "wasm32"))]
178pub fn rgba_to_png_bytes(width: u32, height: u32, rgba: &[u8]) -> Result<Vec<u8>> {
179    use image::{ExtendedColorType, ImageEncoder, codecs::png::PngEncoder};
180
181    let expected = (width as usize) * (height as usize) * 4;
182    if rgba.len() != expected {
183        return Err(EncodeError::Shape {
184            width: width as usize,
185            height: height as usize,
186            got: rgba.len(),
187            expected,
188        });
189    }
190
191    let mut out = Vec::with_capacity(rgba.len() / 8);
192    PngEncoder::new(&mut out).write_image(rgba, width, height, ExtendedColorType::Rgba8)?;
193    Ok(out)
194}
195
196/// Write an RGBA buffer to `path` as a PNG. Native target only.
197#[cfg(not(target_arch = "wasm32"))]
198pub fn save_png<P: AsRef<Path>>(path: P, width: u32, height: u32, rgba: &[u8]) -> Result<()> {
199    let bytes = rgba_to_png_bytes(width, height, rgba)?;
200    std::fs::write(path, bytes)?;
201    Ok(())
202}
203
204// ---------------------------------------------------------------------------
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use surtgis_core::raster::Raster;
210
211    fn solid(w: usize, h: usize, rgba: [u8; 4]) -> RgbaImage {
212        let mut pixels = Vec::with_capacity(w * h * 4);
213        for _ in 0..(w * h) {
214            pixels.extend_from_slice(&rgba);
215        }
216        RgbaImage::from_rgba(w, h, pixels).unwrap()
217    }
218
219    #[test]
220    fn from_rgba_rejects_wrong_length() {
221        let err = RgbaImage::from_rgba(2, 2, vec![0u8; 15]).unwrap_err();
222        match err {
223            EncodeError::Shape { expected, got, .. } => {
224                assert_eq!(expected, 16);
225                assert_eq!(got, 15);
226            }
227            _ => panic!("wrong error kind"),
228        }
229    }
230
231    #[test]
232    fn from_intensity_clamps_and_handles_nan() {
233        // 2x1 raster: NaN, 0.5, 2.0 -> transparent, 128, 255
234        let mut r = Raster::<f64>::new(1, 3);
235        r.set(0, 0, f64::NAN).unwrap();
236        r.set(0, 1, 0.5).unwrap();
237        r.set(0, 2, 2.0).unwrap();
238        let img = RgbaImage::from_intensity(&r);
239        assert_eq!(img.width, 3);
240        assert_eq!(img.height, 1);
241        // Pixel 0: transparent black
242        assert_eq!(&img.pixels[0..4], &[0, 0, 0, 0]);
243        // Pixel 1: 128 grey
244        assert_eq!(&img.pixels[4..8], &[128, 128, 128, 255]);
245        // Pixel 2: 255 grey (clamped)
246        assert_eq!(&img.pixels[8..12], &[255, 255, 255, 255]);
247    }
248
249    #[test]
250    fn over_opaque_top_replaces_base() {
251        let mut base = solid(1, 1, [10, 20, 30, 255]);
252        let top = solid(1, 1, [100, 200, 50, 255]);
253        base.over(&top, 1.0).unwrap();
254        assert_eq!(&base.pixels[..], &[100, 200, 50, 255]);
255    }
256
257    #[test]
258    fn over_zero_opacity_keeps_base() {
259        let mut base = solid(1, 1, [10, 20, 30, 255]);
260        let top = solid(1, 1, [100, 200, 50, 255]);
261        base.over(&top, 0.0).unwrap();
262        assert_eq!(&base.pixels[..], &[10, 20, 30, 255]);
263    }
264
265    #[test]
266    fn over_half_opacity_lerps() {
267        let mut base = solid(1, 1, [0, 0, 0, 255]);
268        let top = solid(1, 1, [200, 200, 200, 255]);
269        base.over(&top, 0.5).unwrap();
270        // 200 * 0.5 + 0 * 0.5 = 100
271        assert_eq!(&base.pixels[..3], &[100, 100, 100]);
272    }
273
274    #[test]
275    fn multiply_with_white_top_full_opacity_preserves_base() {
276        let mut base = solid(1, 1, [128, 64, 32, 255]);
277        let top = solid(1, 1, [255, 255, 255, 255]);
278        base.multiply(&top, 1.0).unwrap();
279        // factor = 1.0; base unchanged
280        assert_eq!(&base.pixels[..3], &[128, 64, 32]);
281    }
282
283    #[test]
284    fn multiply_with_black_top_full_opacity_zeroes_rgb() {
285        let mut base = solid(1, 1, [128, 64, 32, 255]);
286        let top = solid(1, 1, [0, 0, 0, 255]);
287        base.multiply(&top, 1.0).unwrap();
288        assert_eq!(&base.pixels[..3], &[0, 0, 0]);
289    }
290
291    #[test]
292    fn over_rejects_shape_mismatch() {
293        let mut base = solid(2, 2, [0, 0, 0, 255]);
294        let top = solid(3, 3, [0, 0, 0, 255]);
295        assert!(base.over(&top, 1.0).is_err());
296    }
297
298    #[cfg(not(target_arch = "wasm32"))]
299    #[test]
300    fn png_roundtrip_size_and_signature() {
301        let img = solid(4, 3, [255, 0, 0, 255]);
302        let bytes = img.to_png_bytes().unwrap();
303        // PNG signature
304        assert_eq!(
305            &bytes[..8],
306            &[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]
307        );
308    }
309
310    #[cfg(not(target_arch = "wasm32"))]
311    #[test]
312    fn save_png_writes_a_valid_file() {
313        let img = solid(2, 2, [0, 128, 255, 255]);
314        let tmp = std::env::temp_dir().join("colormap_encode_test.png");
315        img.save_png(&tmp).unwrap();
316        let bytes = std::fs::read(&tmp).unwrap();
317        assert_eq!(
318            &bytes[..8],
319            &[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a]
320        );
321        let _ = std::fs::remove_file(&tmp);
322    }
323}
surtgis_colormap/encode.rs

surtgis_colormap/
encode.rs
