Add ImageEncoder implementation for Hdr (rebase of #2013) (#2246)

Co-authored-by: Johannes-Vollmer <[email protected]>

src/codecs/hdr/encoder.rs

@@ -1,6 +1,7 @@
 use crate::codecs::hdr::{rgbe8, Rgbe8Pixel, SIGNATURE};
 use crate::color::Rgb;
-use crate::error::ImageResult;
+use crate::error::{EncodingError, ImageFormatHint, ImageResult};
+use crate::{ExtendedColorType, ImageEncoder, ImageError, ImageFormat};
 use std::cmp::Ordering;
 use std::io::{Result, Write};
 
@@ -9,16 +10,59 @@ pub struct HdrEncoder<W: Write> {
     w: W,
 }
 
+impl<W: Write> ImageEncoder for HdrEncoder<W> {
+    fn write_image(
+        self,
+        unaligned_bytes: &[u8],
+        width: u32,
+        height: u32,
+        color_type: ExtendedColorType,
+    ) -> ImageResult<()> {
+        match color_type {
+            ExtendedColorType::Rgb32F => {
+                let bytes_per_pixel = color_type.bits_per_pixel() as usize / 8;
+                let rgbe_pixels = unaligned_bytes
+                    .chunks_exact(bytes_per_pixel)
+                    .map(|bytes| to_rgbe8(Rgb::<f32>(bytemuck::pod_read_unaligned(bytes))));
+
+                // the length will be checked inside encode_pixels
+                self.encode_pixels(rgbe_pixels, width as usize, height as usize)
+            }
+
+            _ => Err(ImageError::Encoding(EncodingError::new(
+                ImageFormatHint::Exact(ImageFormat::Hdr),
+                "hdr format currently only supports the `Rgb32F` color type".to_string(),
+            ))),
+        }
+    }
+}
+
 impl<W: Write> HdrEncoder<W> {
     /// Creates encoder
     pub fn new(w: W) -> HdrEncoder<W> {
         HdrEncoder { w }
     }
 
-    /// Encodes the image ```data```
+    /// Encodes the image ```rgb```
     /// that has dimensions ```width``` and ```height```
-    pub fn encode(mut self, data: &[Rgb<f32>], width: usize, height: usize) -> ImageResult<()> {
-        assert!(data.len() >= width * height);
+    pub fn encode(self, rgb: &[Rgb<f32>], width: usize, height: usize) -> ImageResult<()> {
+        self.encode_pixels(rgb.iter().map(|&rgb| to_rgbe8(rgb)), width, height)
+    }
+
+    /// Encodes the image ```flattened_rgbe_pixels```
+    /// that has dimensions ```width``` and ```height```.
+    /// The callback must return the color for the given flattened index of the pixel (row major).
+    fn encode_pixels(
+        mut self,
+        mut flattened_rgbe_pixels: impl ExactSizeIterator<Item = Rgbe8Pixel>,
+        width: usize,
+        height: usize,
+    ) -> ImageResult<()> {
+        assert!(
+            flattened_rgbe_pixels.len() >= width * height,
+            "not enough pixels provided"
+        ); // bonus: this might elide some bounds checks
+
         let w = &mut self.w;
         w.write_all(SIGNATURE)?;
         w.write_all(b"\n")?;
@@ -27,8 +71,8 @@ impl<W: Write> HdrEncoder<W> {
         w.write_all(format!("-Y {} +X {}\n", height, width).as_bytes())?;
 
         if !(8..=32_768).contains(&width) {
-            for &pix in data {
-                write_rgbe8(w, to_rgbe8(pix))?;
+            for pixel in flattened_rgbe_pixels {
+                write_rgbe8(w, pixel)?;
             }
         } else {
             // new RLE marker contains scanline width
@@ -39,20 +83,22 @@ impl<W: Write> HdrEncoder<W> {
             let mut bufb = vec![0; width];
             let mut bufe = vec![0; width];
             let mut rle_buf = vec![0; width];
-            for scanline in data.chunks(width) {
-                for ((((r, g), b), e), &pix) in bufr
+            for _scanline_index in 0..height {
+                assert!(flattened_rgbe_pixels.len() >= width); // may reduce the bound checks
+
+                for ((((r, g), b), e), pixel) in bufr
                     .iter_mut()
                     .zip(bufg.iter_mut())
                     .zip(bufb.iter_mut())
                     .zip(bufe.iter_mut())
-                    .zip(scanline.iter())
+                    .zip(&mut flattened_rgbe_pixels)
                 {
-                    let cp = to_rgbe8(pix);
-                    *r = cp.c[0];
-                    *g = cp.c[1];
-                    *b = cp.c[2];
-                    *e = cp.e;
+                    *r = pixel.c[0];
+                    *g = pixel.c[1];
+                    *b = pixel.c[2];
+                    *e = pixel.e;
                 }
+
                 write_rgbe8(w, marker)?; // New RLE encoding marker
                 rle_buf.clear();
                 rle_compress(&bufr[..], &mut rle_buf);
@@ -77,6 +123,7 @@ enum RunOrNot {
     Run(u8, usize),
     Norun(usize, usize),
 }
+
 use self::RunOrNot::{Norun, Run};
 
 const RUN_MAX_LEN: usize = 127;

src/image.rs

@@ -265,7 +265,7 @@ impl ImageFormat {
             ImageFormat::Farbfeld => true,
             ImageFormat::Avif => true,
             ImageFormat::WebP => true,
-            ImageFormat::Hdr => false,
+            ImageFormat::Hdr => true,
             ImageFormat::OpenExr => true,
             ImageFormat::Dds => false,
             ImageFormat::Qoi => true,
@@ -341,8 +341,8 @@ impl ImageFormat {
             ImageFormat::WebP => cfg!(feature = "webp"),
             ImageFormat::OpenExr => cfg!(feature = "exr"),
             ImageFormat::Qoi => cfg!(feature = "qoi"),
+            ImageFormat::Hdr => cfg!(feature = "hdr"),
             ImageFormat::Dds => false,
-            ImageFormat::Hdr => false,
         }
     }
 
@@ -1818,7 +1818,6 @@ mod tests {
             cfg!(feature = "ff"),
             ImageFormat::Farbfeld.writing_enabled()
         );
-        assert!(!ImageFormat::Hdr.writing_enabled());
         assert!(!ImageFormat::Dds.writing_enabled());
     }
 }

src/io/free_functions.rs

@@ -111,6 +111,10 @@ pub(crate) fn write_buffer_impl<W: std::io::Write + Seek>(
         ImageFormat::WebP => {
             webp::WebPEncoder::new_lossless(buffered_write).write_image(buf, width, height, color)
         }
+        #[cfg(feature = "hdr")]
+        ImageFormat::Hdr => {
+            hdr::HdrEncoder::new(buffered_write).write_image(buf, width, height, color)
+        }
         _ => Err(ImageError::Unsupported(
             UnsupportedError::from_format_and_kind(
                 ImageFormatHint::Unknown,