agx/encode/
mod.rs

1//! Image encoding: writing rendered output to JPEG, PNG, TIFF.
2
3use std::io::Cursor;
4use std::path::PathBuf;
5
6use image::codecs::jpeg::JpegEncoder;
7use image::codecs::png::PngEncoder;
8use image::codecs::tiff::TiffEncoder;
9use image::{DynamicImage, Rgb, Rgb32FImage};
10use palette::{LinSrgb, Srgb};
11
12use crate::error::Result;
13use crate::metadata::ImageMetadata;
14
15/// Supported output image formats.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum OutputFormat {
18    /// JPEG with quality control.
19    Jpeg,
20    /// PNG (lossless).
21    Png,
22    /// TIFF (lossless, 16-bit support).
23    Tiff,
24}
25
26impl OutputFormat {
27    /// The canonical file extension for this format.
28    pub fn extension(&self) -> &'static str {
29        match self {
30            OutputFormat::Jpeg => "jpeg",
31            OutputFormat::Png => "png",
32            OutputFormat::Tiff => "tiff",
33        }
34    }
35
36    /// Try to infer format from a file extension string.
37    pub fn from_extension(ext: &str) -> Option<Self> {
38        match ext.to_ascii_lowercase().as_str() {
39            "jpg" | "jpeg" => Some(OutputFormat::Jpeg),
40            "png" => Some(OutputFormat::Png),
41            "tif" | "tiff" => Some(OutputFormat::Tiff),
42            _ => None,
43        }
44    }
45}
46
47/// Options controlling image encoding.
48pub struct EncodeOptions {
49    /// JPEG quality (1-100). Only applies to JPEG output. Default: 92.
50    pub jpeg_quality: u8,
51    /// Explicit output format. If `None`, inferred from file extension.
52    pub format: Option<OutputFormat>,
53}
54
55impl Default for EncodeOptions {
56    fn default() -> Self {
57        Self {
58            jpeg_quality: 92,
59            format: None,
60        }
61    }
62}
63
64/// Resolve the output file path and format.
65///
66/// Rules:
67/// 1. If `format` is specified and the extension matches, use as-is.
68/// 2. If `format` is specified and the extension doesn't match, append the correct extension.
69/// 3. If `format` is `None`, infer from extension.
70/// 4. If the extension is unknown, default to JPEG and append `.jpeg`.
71pub fn resolve_output(
72    path: &std::path::Path,
73    format: Option<OutputFormat>,
74) -> (std::path::PathBuf, OutputFormat) {
75    let ext_format = path
76        .extension()
77        .and_then(|e| e.to_str())
78        .and_then(OutputFormat::from_extension);
79
80    match (format, ext_format) {
81        // Explicit format, extension matches
82        (Some(fmt), Some(ext_fmt)) if fmt == ext_fmt => (path.to_path_buf(), fmt),
83        // Explicit format, extension doesn't match — append correct extension
84        (Some(fmt), _) => {
85            let mut new_path = path.as_os_str().to_owned();
86            new_path.push(".");
87            new_path.push(fmt.extension());
88            (std::path::PathBuf::from(new_path), fmt)
89        }
90        // No explicit format, known extension — infer
91        (None, Some(ext_fmt)) => (path.to_path_buf(), ext_fmt),
92        // No explicit format, unknown/missing extension — default JPEG, append
93        (None, None) => {
94            let mut new_path = path.as_os_str().to_owned();
95            new_path.push(".jpeg");
96            (std::path::PathBuf::from(new_path), OutputFormat::Jpeg)
97        }
98    }
99}
100
101/// Convert a linear sRGB f32 image buffer to a DynamicImage in sRGB gamma space.
102pub fn linear_to_srgb_dynamic(linear: &Rgb32FImage) -> DynamicImage {
103    let (w, h) = linear.dimensions();
104    let srgb = Rgb32FImage::from_fn(w, h, |x, y| {
105        let p = linear.get_pixel(x, y);
106        let srgb: Srgb<f32> = LinSrgb::new(p.0[0], p.0[1], p.0[2]).into_encoding();
107        Rgb([srgb.red, srgb.green, srgb.blue])
108    });
109    DynamicImage::ImageRgb32F(srgb)
110}
111
112/// Encode a linear sRGB f32 image to a file with full options.
113///
114/// Resolves the output format and path, encodes with the appropriate encoder,
115/// and optionally injects metadata. Returns the final output path (which may
116/// differ from the input path if an extension was appended).
117pub fn encode_to_file_with_options(
118    linear: &Rgb32FImage,
119    path: &std::path::Path,
120    options: &EncodeOptions,
121    metadata: Option<&ImageMetadata>,
122) -> Result<PathBuf> {
123    let (final_path, format) = resolve_output(path, options.format);
124
125    let dynamic = linear_to_srgb_dynamic(linear);
126    let rgb8 = dynamic.to_rgb8();
127
128    // Encode to in-memory buffer with format-specific encoder
129    let buf = match format {
130        OutputFormat::Jpeg => {
131            let mut buf = Vec::new();
132            let encoder = JpegEncoder::new_with_quality(&mut buf, options.jpeg_quality);
133            rgb8.write_with_encoder(encoder)
134                .map_err(|e| crate::error::AgxError::Encode(e.to_string()))?;
135            buf
136        }
137        OutputFormat::Png => {
138            let mut buf = Vec::new();
139            let encoder = PngEncoder::new(&mut buf);
140            rgb8.write_with_encoder(encoder)
141                .map_err(|e| crate::error::AgxError::Encode(e.to_string()))?;
142            buf
143        }
144        OutputFormat::Tiff => {
145            let mut buf = Vec::new();
146            let cursor = Cursor::new(&mut buf);
147            let encoder = TiffEncoder::new(cursor);
148            rgb8.write_with_encoder(encoder)
149                .map_err(|e| crate::error::AgxError::Encode(e.to_string()))?;
150            buf
151        }
152    };
153
154    // Inject metadata if available
155    let buf = if let Some(meta) = metadata {
156        inject_metadata(buf, format, meta)?
157    } else {
158        buf
159    };
160
161    std::fs::write(&final_path, &buf)?;
162
163    // For TIFF output, inject metadata via little_exif after writing
164    if format == OutputFormat::Tiff {
165        if let Some(meta) = metadata {
166            inject_metadata_tiff(&final_path, meta);
167        }
168    }
169
170    Ok(final_path)
171}
172
173/// Encode a linear sRGB f32 image to a file, converting to sRGB gamma space.
174///
175/// Uses default options (JPEG quality 92, format inferred from extension).
176/// For more control, use `encode_to_file_with_options`.
177pub fn encode_to_file(linear: &Rgb32FImage, path: &std::path::Path) -> Result<()> {
178    encode_to_file_with_options(linear, path, &EncodeOptions::default(), None)?;
179    Ok(())
180}
181
182/// Inject metadata into an encoded JPEG or PNG buffer.
183fn inject_metadata(
184    buf: Vec<u8>,
185    format: OutputFormat,
186    metadata: &ImageMetadata,
187) -> Result<Vec<u8>> {
188    use img_parts::{ImageEXIF, ImageICC};
189
190    match format {
191        OutputFormat::Jpeg => {
192            let mut jpeg = img_parts::jpeg::Jpeg::from_bytes(buf.into())
193                .map_err(|e| crate::error::AgxError::Encode(format!("metadata injection: {e}")))?;
194            if let Some(exif) = &metadata.exif {
195                jpeg.set_exif(Some(exif.clone().into()));
196            }
197            if let Some(icc) = &metadata.icc_profile {
198                jpeg.set_icc_profile(Some(icc.clone().into()));
199            }
200            let mut out = Vec::new();
201            jpeg.encoder()
202                .write_to(&mut out)
203                .map_err(|e| crate::error::AgxError::Encode(format!("metadata write: {e}")))?;
204            Ok(out)
205        }
206        OutputFormat::Png => {
207            let mut png = img_parts::png::Png::from_bytes(buf.into())
208                .map_err(|e| crate::error::AgxError::Encode(format!("metadata injection: {e}")))?;
209            if let Some(exif) = &metadata.exif {
210                png.set_exif(Some(exif.clone().into()));
211            }
212            if let Some(icc) = &metadata.icc_profile {
213                png.set_icc_profile(Some(icc.clone().into()));
214            }
215            let mut out = Vec::new();
216            png.encoder()
217                .write_to(&mut out)
218                .map_err(|e| crate::error::AgxError::Encode(format!("metadata write: {e}")))?;
219            Ok(out)
220        }
221        OutputFormat::Tiff => Ok(buf), // Handled separately via inject_metadata_tiff
222    }
223}
224
225/// Inject metadata into an existing TIFF file via little_exif. Best-effort — failures are silent.
226fn inject_metadata_tiff(path: &std::path::Path, metadata: &ImageMetadata) {
227    if let Some(exif_bytes) = &metadata.exif {
228        let file_ext = little_exif::filetype::FileExtension::TIFF;
229        if let Ok(exif_meta) = little_exif::metadata::Metadata::new_from_vec(exif_bytes, file_ext) {
230            let _ = exif_meta.write_to_file(path);
231        }
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use image::ImageBuffer;
239    use std::path::PathBuf;
240
241    #[test]
242    fn roundtrip_linear_to_srgb_pixel_values() {
243        // linear 0.2159 should round-trip to sRGB ~128
244        let linear: Rgb32FImage = ImageBuffer::from_pixel(1, 1, Rgb([0.2159f32, 0.2159, 0.2159]));
245        let dynamic = linear_to_srgb_dynamic(&linear);
246        let rgb8 = dynamic.to_rgb8();
247        let pixel = rgb8.get_pixel(0, 0);
248        assert!(
249            (pixel.0[0] as i32 - 128).unsigned_abs() <= 1,
250            "Expected ~128, got {}",
251            pixel.0[0]
252        );
253    }
254
255    #[test]
256    fn encode_saves_file() {
257        let temp_path = std::env::temp_dir().join("agx_test_encode.png");
258        let linear: Rgb32FImage = ImageBuffer::from_pixel(2, 2, Rgb([0.5f32, 0.5, 0.5]));
259        encode_to_file(&linear, &temp_path).unwrap();
260        assert!(temp_path.exists());
261        let _ = std::fs::remove_file(&temp_path);
262    }
263
264    #[test]
265    fn encode_options_default_quality_is_92() {
266        let opts = EncodeOptions::default();
267        assert_eq!(opts.jpeg_quality, 92);
268        assert!(opts.format.is_none());
269    }
270
271    #[test]
272    fn output_format_extensions() {
273        assert_eq!(OutputFormat::Jpeg.extension(), "jpeg");
274        assert_eq!(OutputFormat::Png.extension(), "png");
275        assert_eq!(OutputFormat::Tiff.extension(), "tiff");
276    }
277
278    #[test]
279    fn resolve_output_infers_jpeg_from_jpg() {
280        let (path, fmt) = resolve_output(std::path::Path::new("out.jpg"), None);
281        assert_eq!(fmt, OutputFormat::Jpeg);
282        assert_eq!(path, PathBuf::from("out.jpg"));
283    }
284
285    #[test]
286    fn resolve_output_infers_png() {
287        let (path, fmt) = resolve_output(std::path::Path::new("out.png"), None);
288        assert_eq!(fmt, OutputFormat::Png);
289        assert_eq!(path, PathBuf::from("out.png"));
290    }
291
292    #[test]
293    fn resolve_output_infers_tiff() {
294        let (path, fmt) = resolve_output(std::path::Path::new("out.tif"), None);
295        assert_eq!(fmt, OutputFormat::Tiff);
296        assert_eq!(path, PathBuf::from("out.tif"));
297    }
298
299    #[test]
300    fn resolve_output_format_override_matching_ext() {
301        let (path, fmt) = resolve_output(std::path::Path::new("out.jpg"), Some(OutputFormat::Jpeg));
302        assert_eq!(fmt, OutputFormat::Jpeg);
303        assert_eq!(path, PathBuf::from("out.jpg"));
304    }
305
306    #[test]
307    fn resolve_output_format_override_mismatched_ext_appends() {
308        let (path, fmt) = resolve_output(std::path::Path::new("out.png"), Some(OutputFormat::Jpeg));
309        assert_eq!(fmt, OutputFormat::Jpeg);
310        assert_eq!(path, PathBuf::from("out.png.jpeg"));
311    }
312
313    #[test]
314    fn resolve_output_unknown_ext_defaults_to_jpeg() {
315        let (path, fmt) = resolve_output(std::path::Path::new("out.xyz"), None);
316        assert_eq!(fmt, OutputFormat::Jpeg);
317        assert_eq!(path, PathBuf::from("out.xyz.jpeg"));
318    }
319
320    #[test]
321    fn resolve_output_no_extension_defaults_to_jpeg() {
322        let (path, fmt) = resolve_output(std::path::Path::new("output"), None);
323        assert_eq!(fmt, OutputFormat::Jpeg);
324        assert_eq!(path, PathBuf::from("output.jpeg"));
325    }
326
327    #[test]
328    fn encode_jpeg_with_quality_produces_file() {
329        let temp_path = std::env::temp_dir().join("agx_test_quality.jpg");
330        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
331        let opts = EncodeOptions {
332            jpeg_quality: 95,
333            format: None,
334        };
335        let result = encode_to_file_with_options(&linear, &temp_path, &opts, None);
336        assert!(result.is_ok());
337        let final_path = result.unwrap();
338        assert!(final_path.exists());
339        let _ = std::fs::remove_file(&final_path);
340    }
341
342    #[test]
343    fn encode_jpeg_quality_affects_file_size() {
344        let linear: Rgb32FImage = ImageBuffer::from_pixel(64, 64, Rgb([0.5f32, 0.3, 0.1]));
345
346        let path_low = std::env::temp_dir().join("agx_test_q50.jpg");
347        let path_high = std::env::temp_dir().join("agx_test_q95.jpg");
348
349        let opts_low = EncodeOptions {
350            jpeg_quality: 50,
351            format: None,
352        };
353        let opts_high = EncodeOptions {
354            jpeg_quality: 95,
355            format: None,
356        };
357
358        encode_to_file_with_options(&linear, &path_low, &opts_low, None).unwrap();
359        encode_to_file_with_options(&linear, &path_high, &opts_high, None).unwrap();
360
361        let size_low = std::fs::metadata(&path_low).unwrap().len();
362        let size_high = std::fs::metadata(&path_high).unwrap().len();
363        assert!(
364            size_high > size_low,
365            "Higher quality should produce larger file: q95={size_high} vs q50={size_low}"
366        );
367
368        let _ = std::fs::remove_file(&path_low);
369        let _ = std::fs::remove_file(&path_high);
370    }
371
372    #[test]
373    fn encode_png_format() {
374        let temp_path = std::env::temp_dir().join("agx_test_fmt.png");
375        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
376        let opts = EncodeOptions {
377            jpeg_quality: 92,
378            format: None,
379        };
380        let final_path = encode_to_file_with_options(&linear, &temp_path, &opts, None).unwrap();
381        assert!(final_path.exists());
382        let img = image::open(&final_path).unwrap();
383        assert_eq!(img.width(), 4);
384        let _ = std::fs::remove_file(&final_path);
385    }
386
387    #[test]
388    fn encode_tiff_format() {
389        let temp_path = std::env::temp_dir().join("agx_test_fmt.tiff");
390        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
391        let opts = EncodeOptions {
392            jpeg_quality: 92,
393            format: None,
394        };
395        let final_path = encode_to_file_with_options(&linear, &temp_path, &opts, None).unwrap();
396        assert!(final_path.exists());
397        let img = image::open(&final_path).unwrap();
398        assert_eq!(img.width(), 4);
399        let _ = std::fs::remove_file(&final_path);
400    }
401
402    #[test]
403    fn encode_format_override_appends_extension() {
404        let temp_path = std::env::temp_dir().join("agx_test_override.png");
405        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
406        let opts = EncodeOptions {
407            jpeg_quality: 92,
408            format: Some(OutputFormat::Jpeg),
409        };
410        let final_path = encode_to_file_with_options(&linear, &temp_path, &opts, None).unwrap();
411        assert_eq!(
412            final_path,
413            std::env::temp_dir().join("agx_test_override.png.jpeg")
414        );
415        assert!(final_path.exists());
416        let _ = std::fs::remove_file(&final_path);
417    }
418
419    #[test]
420    fn metadata_roundtrip_jpeg() {
421        let exif_bytes = vec![
422            0x45, 0x78, 0x69, 0x66, 0x00, 0x00, // "Exif\0\0"
423            0x4D, 0x4D, // Big-endian TIFF header
424            0x00, 0x2A, // TIFF magic
425            0x00, 0x00, 0x00, 0x08, // offset to IFD
426        ];
427        let meta = ImageMetadata {
428            exif: Some(exif_bytes.clone()),
429            icc_profile: None,
430        };
431
432        let temp_path = std::env::temp_dir().join("agx_test_meta_rt.jpg");
433        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
434        let opts = EncodeOptions {
435            jpeg_quality: 92,
436            format: None,
437        };
438        encode_to_file_with_options(&linear, &temp_path, &opts, Some(&meta)).unwrap();
439
440        let meta_out = crate::metadata::extract_metadata(&temp_path);
441        assert!(meta_out.is_some(), "Should have metadata in output");
442        assert!(
443            meta_out.as_ref().unwrap().exif.is_some(),
444            "Should have EXIF in output"
445        );
446
447        let _ = std::fs::remove_file(&temp_path);
448    }
449
450    #[test]
451    fn encode_without_metadata_still_works() {
452        let temp_path = std::env::temp_dir().join("agx_test_no_meta.jpg");
453        let linear: Rgb32FImage = ImageBuffer::from_pixel(4, 4, Rgb([0.5f32, 0.5, 0.5]));
454        let opts = EncodeOptions::default();
455        let result = encode_to_file_with_options(&linear, &temp_path, &opts, None);
456        assert!(result.is_ok());
457        let _ = std::fs::remove_file(result.unwrap());
458    }
459}
agx/encode/mod.rs

agx/encode/
mod.rs
