1use crate::error::{AgError, AgResult};
13use bytes::Bytes;
14use image::codecs::png::{CompressionType, FilterType, PngEncoder};
15use image::{ColorType, ImageEncoder};
16use ndarray::Array3;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
24pub enum TileFormat {
25 #[default]
27 Png,
28 Webp,
30 Jpeg,
32}
33
34impl TileFormat {
35 pub fn content_type(&self) -> &'static str {
37 match self {
38 TileFormat::Png => "image/png",
39 TileFormat::Webp => "image/webp",
40 TileFormat::Jpeg => "image/jpeg",
41 }
42 }
43
44 pub fn extension(&self) -> &'static str {
46 match self {
47 TileFormat::Png => "png",
48 TileFormat::Webp => "webp",
49 TileFormat::Jpeg => "jpg",
50 }
51 }
52
53 pub fn feature_name(&self) -> Option<&'static str> {
55 match self {
56 TileFormat::Png => None,
57 TileFormat::Webp => Some("webp"),
58 TileFormat::Jpeg => Some("jpeg"),
59 }
60 }
61
62 pub fn is_available(&self) -> bool {
64 match self {
65 TileFormat::Png => true,
66 TileFormat::Webp => cfg!(feature = "webp"),
67 TileFormat::Jpeg => cfg!(feature = "jpeg"),
68 }
69 }
70
71 pub fn ensure_available(&self) -> AgResult<()> {
74 match self.feature_name() {
75 Some(feature) if !self.is_available() => Err(AgError::UnsupportedPixel(format!(
76 "{feature} output requires building with the \"{feature}\" feature"
77 ))),
78 _ => Ok(()),
79 }
80 }
81}
82
83#[cfg(feature = "jpeg")]
86const JPEG_QUALITY: u8 = 75;
87
88pub fn encode(rgba: &Array3<u8>, format: TileFormat) -> AgResult<Bytes> {
93 format.ensure_available()?;
94 match format {
95 TileFormat::Png => encode_png(rgba),
96 TileFormat::Webp => encode_webp(rgba),
97 TileFormat::Jpeg => encode_jpeg(rgba),
98 }
99}
100
101#[cfg(feature = "webp")]
103fn encode_webp(rgba: &Array3<u8>) -> AgResult<Bytes> {
104 use image::codecs::webp::WebPEncoder;
105
106 let (h, w, c) = (
107 rgba.shape()[0] as u32,
108 rgba.shape()[1] as u32,
109 rgba.shape()[2],
110 );
111 if c != 4 {
112 return Err(AgError::UnsupportedPixel(format!(
113 "expected RGBA channels, got {c}"
114 )));
115 }
116 let standard = rgba.as_standard_layout();
117 let raw = standard
118 .as_slice()
119 .expect("rgba array must be contiguous after as_standard_layout");
120 let mut buf = Vec::new();
121 WebPEncoder::new_lossless(&mut buf)
122 .write_image(raw, w, h, ColorType::Rgba8.into())
123 .map_err(|e| AgError::Tiff(format!("webp encode: {e}")))?;
124 Ok(Bytes::from(buf))
125}
126
127#[cfg(not(feature = "webp"))]
128fn encode_webp(_rgba: &Array3<u8>) -> AgResult<Bytes> {
129 Err(AgError::UnsupportedPixel(
130 "webp output requires building with the \"webp\" feature".into(),
131 ))
132}
133
134#[cfg(feature = "jpeg")]
137fn encode_jpeg(rgba: &Array3<u8>) -> AgResult<Bytes> {
138 use image::codecs::jpeg::JpegEncoder;
139
140 let rgb = flatten_to_rgb_black(rgba)?;
141 let (h, w) = (rgb.shape()[0] as u32, rgb.shape()[1] as u32);
142 let raw = rgb
143 .as_slice()
144 .expect("rgb array is contiguous after flatten");
145 let mut buf = Vec::new();
146 JpegEncoder::new_with_quality(&mut buf, JPEG_QUALITY)
147 .write_image(raw, w, h, ColorType::Rgb8.into())
148 .map_err(|e| AgError::Tiff(format!("jpeg encode: {e}")))?;
149 Ok(Bytes::from(buf))
150}
151
152#[cfg(not(feature = "jpeg"))]
153fn encode_jpeg(_rgba: &Array3<u8>) -> AgResult<Bytes> {
154 Err(AgError::UnsupportedPixel(
155 "jpeg output requires building with the \"jpeg\" feature".into(),
156 ))
157}
158
159#[cfg(feature = "jpeg")]
162fn flatten_to_rgb_black(rgba: &Array3<u8>) -> AgResult<Array3<u8>> {
163 let (h, w, c) = (rgba.shape()[0], rgba.shape()[1], rgba.shape()[2]);
164 if c != 4 {
165 return Err(AgError::UnsupportedPixel(format!(
166 "expected RGBA channels, got {c}"
167 )));
168 }
169 let standard = rgba.as_standard_layout();
170 let raw = standard
171 .as_slice()
172 .expect("rgba array must be contiguous after as_standard_layout");
173 let mut rgb = Vec::with_capacity(h * w * 3);
174 for pixel in raw.chunks_exact(4) {
175 let alpha = u32::from(pixel[3]);
176 rgb.extend_from_slice(&[
177 ((u32::from(pixel[0]) * alpha + 127) / 255) as u8,
178 ((u32::from(pixel[1]) * alpha + 127) / 255) as u8,
179 ((u32::from(pixel[2]) * alpha + 127) / 255) as u8,
180 ]);
181 }
182 Ok(Array3::from_shape_vec((h, w, 3), rgb)
183 .expect("flattened RGB length must match the output shape"))
184}
185
186pub fn encode_png(rgba: &Array3<u8>) -> AgResult<Bytes> {
191 let (h, w, c) = (
192 rgba.shape()[0] as u32,
193 rgba.shape()[1] as u32,
194 rgba.shape()[2],
195 );
196 if c != 4 {
197 return Err(AgError::UnsupportedPixel(format!(
198 "expected RGBA channels, got {c}"
199 )));
200 }
201 let standard = rgba.as_standard_layout();
202 let raw = standard
203 .as_slice()
204 .expect("rgba array must be contiguous after as_standard_layout");
205 let mut buf = Vec::with_capacity((w * h * 4) as usize);
206 let encoder =
207 PngEncoder::new_with_quality(&mut buf, CompressionType::Fast, FilterType::NoFilter);
208 encoder
209 .write_image(raw, w, h, ColorType::Rgba8.into())
210 .map_err(|e| crate::error::AgError::Tiff(format!("png encode: {e}")))?;
211 Ok(Bytes::from(buf))
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 use ndarray::Array3;
218
219 #[test]
220 fn encodes_solid_red_tile() {
221 let mut arr = Array3::<u8>::zeros((4, 4, 4));
222 for r in 0..4 {
223 for c in 0..4 {
224 arr[[r, c, 0]] = 255;
225 arr[[r, c, 3]] = 255;
226 }
227 }
228 let bytes = encode_png(&arr).unwrap();
229 assert!(bytes.starts_with(b"\x89PNG"));
230 }
231
232 #[test]
233 fn rejects_non_rgba_without_panicking() {
234 let arr = Array3::<u8>::zeros((4, 4, 3));
235 let err = encode_png(&arr).unwrap_err();
236 assert!(matches!(err, AgError::UnsupportedPixel(_)));
237 }
238
239 fn solid_rgba(r: u8, g: u8, b: u8, a: u8) -> Array3<u8> {
240 let mut arr = Array3::<u8>::zeros((4, 4, 4));
241 for row in 0..4 {
242 for col in 0..4 {
243 arr[[row, col, 0]] = r;
244 arr[[row, col, 1]] = g;
245 arr[[row, col, 2]] = b;
246 arr[[row, col, 3]] = a;
247 }
248 }
249 arr
250 }
251
252 #[test]
253 fn content_type_matches_format() {
254 assert_eq!(TileFormat::Png.content_type(), "image/png");
255 assert_eq!(TileFormat::Webp.content_type(), "image/webp");
256 assert_eq!(TileFormat::Jpeg.content_type(), "image/jpeg");
257 }
258
259 #[test]
260 fn extension_matches_format() {
261 assert_eq!(TileFormat::Png.extension(), "png");
262 assert_eq!(TileFormat::Webp.extension(), "webp");
263 assert_eq!(TileFormat::Jpeg.extension(), "jpg");
264 }
265
266 #[test]
267 fn png_is_always_available() {
268 assert!(TileFormat::Png.is_available());
269 assert_eq!(TileFormat::Png.feature_name(), None);
270 }
271
272 #[test]
273 fn encode_png_format_produces_png() {
274 let arr = solid_rgba(255, 0, 0, 255);
275 let bytes = encode(&arr, TileFormat::Png).unwrap();
276 assert!(bytes.starts_with(b"\x89PNG"));
277 }
278
279 #[cfg(feature = "jpeg")]
280 #[test]
281 fn flatten_to_rgb_black_premultiplies_alpha() {
282 let half = flatten_to_rgb_black(&solid_rgba(255, 0, 0, 128)).unwrap();
283 assert_eq!(half.shape(), &[4, 4, 3]);
284 assert_eq!(half[[0, 0, 0]], 128);
286 assert_eq!(half[[0, 0, 1]], 0);
287 assert_eq!(half[[0, 0, 2]], 0);
288
289 let transparent = flatten_to_rgb_black(&solid_rgba(255, 255, 255, 0)).unwrap();
290 assert_eq!(transparent[[0, 0, 0]], 0);
291 assert_eq!(transparent[[0, 0, 1]], 0);
292 assert_eq!(transparent[[0, 0, 2]], 0);
293
294 let opaque = flatten_to_rgb_black(&solid_rgba(10, 20, 30, 255)).unwrap();
295 assert_eq!(opaque[[1, 1, 0]], 10);
296 assert_eq!(opaque[[1, 1, 1]], 20);
297 assert_eq!(opaque[[1, 1, 2]], 30);
298 }
299
300 #[cfg(feature = "jpeg")]
301 #[test]
302 fn flatten_to_rgb_black_accepts_non_standard_layout() {
303 let mut rgba = Array3::<u8>::zeros((2, 3, 4));
304 rgba[[1, 2, 0]] = 200;
305 rgba[[1, 2, 1]] = 100;
306 rgba[[1, 2, 2]] = 50;
307 rgba[[1, 2, 3]] = 128;
308 let transposed = rgba.permuted_axes([1, 0, 2]);
309
310 let rgb = flatten_to_rgb_black(&transposed).unwrap();
311
312 assert_eq!(rgb.shape(), &[3, 2, 3]);
313 assert_eq!(rgb[[2, 1, 0]], 100);
314 assert_eq!(rgb[[2, 1, 1]], 50);
315 assert_eq!(rgb[[2, 1, 2]], 25);
316 }
317
318 #[cfg(feature = "webp")]
319 #[test]
320 fn encode_webp_format_produces_webp() {
321 let arr = solid_rgba(0, 255, 0, 255);
322 let bytes = encode(&arr, TileFormat::Webp).unwrap();
323 assert!(bytes.starts_with(b"RIFF"));
324 assert_eq!(&bytes[8..12], b"WEBP");
325 }
326
327 #[cfg(not(feature = "webp"))]
328 #[test]
329 fn encode_webp_without_feature_errors() {
330 let arr = solid_rgba(0, 255, 0, 255);
331 let err = encode(&arr, TileFormat::Webp).unwrap_err();
332 assert!(matches!(err, AgError::UnsupportedPixel(_)));
333 }
334
335 #[cfg(feature = "jpeg")]
336 #[test]
337 fn encode_jpeg_format_produces_jpeg() {
338 let arr = solid_rgba(0, 0, 255, 255);
339 let bytes = encode(&arr, TileFormat::Jpeg).unwrap();
340 assert!(bytes.starts_with(b"\xFF\xD8\xFF"));
341 }
342
343 #[cfg(not(feature = "jpeg"))]
344 #[test]
345 fn encode_jpeg_without_feature_errors() {
346 let arr = solid_rgba(0, 0, 255, 255);
347 let err = encode(&arr, TileFormat::Jpeg).unwrap_err();
348 assert!(matches!(err, AgError::UnsupportedPixel(_)));
349 }
350}