jxl-oxide: Accept integer output buffer (#366)

* Accept integer output buffer

* Update CHANGELOG.md

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `jxl-oxide`: Accept `u8` and `u16` output buffers in `ImageStream::write_to_buffer` (#366).
+
 ### Changed
 - `jxl-color`: Use better PQ to HLG method (#348).
 

crates/jxl-oxide-cli/src/output.rs

@@ -1,6 +1,6 @@
 use std::io::prelude::*;
 
-use jxl_oxide::{FrameBuffer, JxlImage, PixelFormat, Render};
+use jxl_oxide::{JxlImage, PixelFormat, Render};
 
 #[cfg(feature = "__ffmpeg")]
 mod video;
@@ -90,27 +90,27 @@ pub(crate) fn write_png<W: Write>(
         }
 
         let mut stream = keyframe.stream();
-        let mut fb = FrameBuffer::new(
-            stream.width() as usize,
-            stream.height() as usize,
-            stream.channels() as usize,
-        );
-        stream.write_to_buffer(fb.buf_mut());
 
         if sixteen_bits {
-            let mut buf = vec![0u8; fb.width() * fb.height() * fb.channels() * 2];
-            for (b, s) in buf.chunks_exact_mut(2).zip(fb.buf()) {
-                let w = (*s * 65535.0 + 0.5).clamp(0.0, 65535.0) as u16;
-                let [b0, b1] = w.to_be_bytes();
-                b[0] = b0;
-                b[1] = b1;
+            let mut fb_row = vec![0u16; (stream.width() * stream.channels()) as usize];
+            let mut buf =
+                vec![0u8; (stream.width() * stream.height() * stream.channels() * 2) as usize];
+
+            let buf_it = buf.chunks_exact_mut((stream.width() * stream.channels() * 2) as usize);
+            for buf_row in buf_it {
+                stream.write_to_buffer(&mut fb_row);
+                for (b, s) in buf_row.chunks_exact_mut(2).zip(&fb_row) {
+                    let [b0, b1] = s.to_be_bytes();
+                    b[0] = b0;
+                    b[1] = b1;
+                }
             }
+
             writer.write_image_data(&buf)?;
         } else {
-            let mut buf = vec![0u8; fb.width() * fb.height() * fb.channels()];
-            for (b, s) in buf.iter_mut().zip(fb.buf()) {
-                *b = (*s * 255.0 + 0.5).clamp(0.0, 255.0) as u8;
-            }
+            let mut buf =
+                vec![0u8; (stream.width() * stream.height() * stream.channels()) as usize];
+            stream.write_to_buffer(&mut buf);
             writer.write_image_data(&buf)?;
         }
     }

crates/jxl-oxide-cli/src/output/video/context.rs

@@ -543,7 +543,6 @@ impl<W> VideoContext<W> {
 
         let data_ptr = data[0];
         let stride = linesize[0];
-        let mut tmp = vec![0f32; width * channels];
         for y in 0..height {
             let output_row = unsafe {
                 let base_ptr = if stride < 0 {
@@ -555,10 +554,7 @@ impl<W> VideoContext<W> {
                 std::slice::from_raw_parts_mut(base_ptr as *mut u16, video_width * channels)
             };
             let (output_row, output_trailing) = output_row.split_at_mut(width * channels);
-            stream.write_to_buffer(&mut tmp);
-            for (o, i) in output_row.iter_mut().zip(&tmp) {
-                *o = (*i * 65535.0 + 0.5).max(0.0) as u16;
-            }
+            stream.write_to_buffer(output_row);
             output_trailing.fill(0);
         }
 

crates/jxl-oxide-wasm/src/lib.rs

@@ -260,12 +260,6 @@ impl RenderResult {
     pub fn into_png(self) -> Result<Vec<u8>, String> {
         let image = self.image;
         let mut stream = image.stream();
-        let mut fb = jxl_oxide::FrameBuffer::new(
-            stream.width() as usize,
-            stream.height() as usize,
-            stream.channels() as usize,
-        );
-        stream.write_to_buffer(fb.buf_mut());
 
         let mut out = Vec::new();
         let mut encoder = png::Encoder::new(&mut out, stream.width(), stream.height());
@@ -302,19 +296,25 @@ impl RenderResult {
         }
 
         if self.need_high_precision {
-            let mut buf = vec![0u8; fb.width() * fb.height() * fb.channels() * 2];
-            for (b, s) in buf.chunks_exact_mut(2).zip(fb.buf()) {
-                let w = (*s * 65535.0 + 0.5).clamp(0.0, 65535.0) as u16;
-                let [b0, b1] = w.to_be_bytes();
-                b[0] = b0;
-                b[1] = b1;
+            let mut fb_row = vec![0u16; (stream.width() * stream.channels()) as usize];
+            let mut buf =
+                vec![0u8; (stream.width() * stream.height() * stream.channels() * 2) as usize];
+
+            let buf_it = buf.chunks_exact_mut((stream.width() * stream.channels() * 2) as usize);
+            for buf_row in buf_it {
+                stream.write_to_buffer(&mut fb_row);
+                for (b, s) in buf_row.chunks_exact_mut(2).zip(&fb_row) {
+                    let [b0, b1] = s.to_be_bytes();
+                    b[0] = b0;
+                    b[1] = b1;
+                }
             }
+
             writer.write_image_data(&buf).map_err(|e| e.to_string())?;
         } else {
-            let mut buf = vec![0u8; fb.width() * fb.height() * fb.channels()];
-            for (b, s) in buf.iter_mut().zip(fb.buf()) {
-                *b = (*s * 255.0 + 0.5).clamp(0.0, 255.0) as u8;
-            }
+            let mut buf =
+                vec![0u8; (stream.width() * stream.height() * stream.channels()) as usize];
+            stream.write_to_buffer(&mut buf);
             writer.write_image_data(&buf).map_err(|e| e.to_string())?;
         }
 

crates/jxl-oxide/src/fb.rs

@@ -1,5 +1,6 @@
 use jxl_image::BitDepth;
 use jxl_render::{ImageBuffer, Region};
+use private::Sealed;
 
 /// Frame buffer representing a decoded image.
 #[derive(Debug, Clone)]
@@ -301,7 +302,7 @@ impl ImageStream<'_> {
     }
 
     /// Writes next samples to the buffer, returning how many samples are written.
-    pub fn write_to_buffer(&mut self, buf: &mut [f32]) -> usize {
+    pub fn write_to_buffer<Sample: FrameBufferSample>(&mut self, buf: &mut [Sample]) -> usize {
         let channels = self.grids.len() as u32;
         let mut buf_it = buf.iter_mut();
         let mut count = 0usize;
@@ -317,7 +318,7 @@ impl ImageStream<'_> {
                         orig_x.checked_add_signed(start_x),
                         orig_y.checked_add_signed(start_y),
                     ) else {
-                        *v = 0.0;
+                        *v = Sample::default();
                         count += 1;
                         self.c += 1;
                         continue;
@@ -326,17 +327,13 @@ impl ImageStream<'_> {
                     let y = y as usize;
                     let grid = &self.grids[self.c as usize];
                     let bit_depth = self.bit_depth[self.c as usize];
-                    *v = match grid {
-                        ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
-                        ImageBuffer::I32(g) => {
-                            bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0))
-                        }
-                        ImageBuffer::I16(g) => {
-                            bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0) as i32)
-                        }
-                    };
 
-                    if self.c < 3 {
+                    if self.c >= 3 || self.spot_colors.is_empty() {
+                        v.copy_from_grid(grid, x, y, bit_depth);
+                    } else {
+                        let mut tmp_sample = 0f32;
+                        tmp_sample.copy_from_grid(grid, x, y, bit_depth);
+
                         for spot in &self.spot_colors {
                             let ImageStreamSpotColor {
                                 grid,
@@ -353,21 +350,17 @@ impl ImageStream<'_> {
                             let mix = if let (Some(x), Some(y)) = xy {
                                 let x = x as usize;
                                 let y = y as usize;
-                                let val = match grid {
-                                    ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
-                                    ImageBuffer::I32(g) => bit_depth
-                                        .parse_integer_sample(g.get(x, y).copied().unwrap_or(0)),
-                                    ImageBuffer::I16(g) => bit_depth.parse_integer_sample(
-                                        g.get(x, y).copied().unwrap_or(0) as i32,
-                                    ),
-                                };
+                                let mut val = 0f32;
+                                val.copy_from_grid(grid, x, y, bit_depth);
                                 val * solidity
                             } else {
                                 0.0
                             };
 
-                            *v = color * mix + *v * (1.0 - mix);
+                            tmp_sample = color * mix + tmp_sample * (1.0 - mix);
                         }
+
+                        v.copy_from_f32(tmp_sample);
                     }
 
                     count += 1;
@@ -407,3 +400,111 @@ struct ImageStreamSpotColor<'r> {
     rgb: (f32, f32, f32),
     solidity: f32,
 }
+
+/// Marker trait for supported output sample types.
+pub trait FrameBufferSample: private::Sealed {}
+
+/// Output as 32-bit float samples, with nominal range of `[0, 1]`.
+impl FrameBufferSample for f32 {}
+
+/// Output as 16-bit unsigned integer samples.
+impl FrameBufferSample for u16 {}
+
+/// Output as 8-bit unsigned integer samples.
+impl FrameBufferSample for u8 {}
+
+mod private {
+    use jxl_image::BitDepth;
+    use jxl_render::ImageBuffer;
+
+    pub trait Sealed: Sized + Default {
+        fn copy_from_grid(&mut self, grid: &ImageBuffer, x: usize, y: usize, bit_depth: BitDepth);
+        fn copy_from_f32(&mut self, val: f32);
+    }
+
+    impl Sealed for f32 {
+        #[inline]
+        fn copy_from_grid(&mut self, grid: &ImageBuffer, x: usize, y: usize, bit_depth: BitDepth) {
+            *self = match grid {
+                ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
+                ImageBuffer::I32(g) => {
+                    bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0))
+                }
+                ImageBuffer::I16(g) => {
+                    bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0) as i32)
+                }
+            };
+        }
+
+        #[inline]
+        fn copy_from_f32(&mut self, val: f32) {
+            *self = val;
+        }
+    }
+
+    impl Sealed for u16 {
+        #[inline]
+        fn copy_from_grid(&mut self, grid: &ImageBuffer, x: usize, y: usize, bit_depth: BitDepth) {
+            if matches!(
+                bit_depth,
+                BitDepth::IntegerSample {
+                    bits_per_sample: 16
+                }
+            ) {
+                *self = match grid {
+                    ImageBuffer::F32(g) => (g.get(x, y).copied().unwrap_or(0.0) * 65535.0 + 0.5)
+                        .clamp(0.0, 65535.0) as u16,
+                    ImageBuffer::I32(g) => g.get(x, y).copied().unwrap_or(0).clamp(0, 65535) as u16,
+                    ImageBuffer::I16(g) => g.get(x, y).copied().unwrap_or(0).max(0) as u16,
+                };
+            } else {
+                let flt = match grid {
+                    ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
+                    ImageBuffer::I32(g) => {
+                        bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0))
+                    }
+                    ImageBuffer::I16(g) => {
+                        bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0) as i32)
+                    }
+                };
+                self.copy_from_f32(flt);
+            }
+        }
+
+        #[inline]
+        fn copy_from_f32(&mut self, val: f32) {
+            *self = (val * 65535.0 + 0.5).clamp(0.0, 65535.0) as u16;
+        }
+    }
+
+    impl Sealed for u8 {
+        #[inline]
+        fn copy_from_grid(&mut self, grid: &ImageBuffer, x: usize, y: usize, bit_depth: BitDepth) {
+            if matches!(bit_depth, BitDepth::IntegerSample { bits_per_sample: 8 }) {
+                *self = match grid {
+                    ImageBuffer::F32(g) => {
+                        (g.get(x, y).copied().unwrap_or(0.0) * 255.0 + 0.5).clamp(0.0, 255.0) as u8
+                    }
+                    ImageBuffer::I32(g) => g.get(x, y).copied().unwrap_or(0).clamp(0, 255) as u8,
+                    ImageBuffer::I16(g) => g.get(x, y).copied().unwrap_or(0).clamp(0, 255) as u8,
+                };
+            } else {
+                let flt = match grid {
+                    ImageBuffer::F32(g) => g.get(x, y).copied().unwrap_or(0.0),
+                    ImageBuffer::I32(g) => {
+                        bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0))
+                    }
+                    ImageBuffer::I16(g) => {
+                        bit_depth.parse_integer_sample(g.get(x, y).copied().unwrap_or(0) as i32)
+                    }
+                };
+                self.copy_from_f32(flt);
+            }
+        }
+
+        #[inline]
+        fn copy_from_f32(&mut self, val: f32) {
+            *self = (val * 255.0 + 0.5).clamp(0.0, 255.0) as u8;
+        }
+    }
+}

crates/jxl-oxide/src/lib.rs

@@ -170,7 +170,7 @@ mod lcms2;
 
 #[cfg(feature = "lcms2")]
 pub use self::lcms2::Lcms2;
-pub use fb::{FrameBuffer, ImageStream};
+pub use fb::{FrameBuffer, FrameBufferSample, ImageStream};
 
 pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync + 'static>>;