Preserve JPEG XL metadata and associated alpha

crates/jxl/src/decode.rs

@@ -1,7 +1,7 @@
 use std::panic::{AssertUnwindSafe, catch_unwind};
 
-use jxl_oxide::color::{ColourSpace, EnumColourEncoding, Primaries, RenderingIntent, TransferFunction, WhitePoint};
-use jxl_oxide::{AllocTracker, AuxBoxData, InitializeResult, JxlImage};
+use jxl_oxide::color::{ColourEncoding, ColourSpace, EnumColourEncoding, Primaries, RenderingIntent, TransferFunction, WhitePoint};
+use jxl_oxide::{AllocTracker, AuxBoxData, ExtraChannelType, InitializeResult, JxlImage};
 
 use crate::{DecodeLimits, Decoded, Error, Layout, Metadata, OutputColor, Result};
 
@@ -70,7 +70,7 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
     }
     image.finalize().map_err(decode_error)?;
 
-    let (bit_depth, float_samples, grayscale, source_orientation, xyb_encoded) = {
+    let (bit_depth, float_samples, grayscale, source_orientation, xyb_encoded, extra_channels, colour_encoding) = {
         let metadata = &image.image_header().metadata;
         (
             metadata.bit_depth.bits_per_sample() as u8,
@@ -78,11 +78,29 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
             metadata.grayscale(),
             metadata.orientation,
             metadata.xyb_encoded,
+            metadata.ec_info.iter().map(|info| info.ty).collect::<Vec<_>>(),
+            metadata.colour_encoding.clone(),
         )
     };
     let source_icc = image.original_icc().map(ToOwned::to_owned);
+    let has_source_icc = source_icc.is_some();
+    if has_black_channel(&extra_channels) || image.pixel_format().has_black() {
+        return Err(Error::Unsupported("CMYK JPEG XL rendering requires an external color-management transform".into()));
+    }
+    let alpha_channel = first_alpha(&extra_channels);
+    let alpha_associated = alpha_channel.is_some_and(|(_, associated)| associated);
     let via_srgb = output_color == OutputColor::LinearRec2020 && source_icc.is_none() && !grayscale && xyb_encoded;
-    if output_color == OutputColor::LinearRec2020 && source_icc.is_none() {
+    // Associated samples must be divided while still in their coded nonlinear space. Leave these
+    // frames in their own encoding so the caller can apply its color transform after unpremultiplying.
+    // XYB and unsupported transfer curves stay on the renderer's linear target path; those cases
+    // are premultiplied after decoding instead.
+    let enum_source_has_supported_curve = matches!(
+        &colour_encoding,
+        ColourEncoding::Enum(encoding) if enum_encoding_has_supported_curve(encoding)
+    );
+    let preserve_source_encoding =
+        output_color == OutputColor::LinearRec2020 && alpha_associated && !xyb_encoded && (source_icc.is_some() || enum_source_has_supported_curve);
+    if output_color == OutputColor::LinearRec2020 && source_icc.is_none() && !preserve_source_encoding {
         image.request_color_encoding(EnumColourEncoding {
             colour_space: if grayscale { ColourSpace::Grey } else { ColourSpace::Rgb },
             white_point: WhitePoint::D65,
@@ -96,7 +114,7 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
     if pixel_format.has_black() {
         return Err(Error::Unsupported("CMYK JPEG XL rendering requires an external color-management transform".into()));
     }
-    let has_alpha = pixel_format.has_alpha();
+    let has_alpha = alpha_channel.is_some();
     let layout = match (grayscale, has_alpha) {
         (true, false) => Layout::Gray,
         (true, true) => Layout::GrayAlpha,
@@ -110,6 +128,11 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
     if channels < colors || (has_alpha && channels < keep) {
         return Err(Error::Malformed("decoded frame has too few channels".into()));
     }
+    let color_channels = channels
+        .checked_sub(extra_channels.len())
+        .filter(|&count| count >= colors)
+        .ok_or_else(|| Error::Malformed("decoded frame has too few color channels".into()))?;
+    let alpha_index = alpha_channel.map(|(index, _)| color_channels + index);
     let (rendered_width, rendered_height) = (
         u32::try_from(buffer.width()).map_err(|_| Error::LimitExceeded("rendered width exceeds 32-bit dimensions".into()))?,
         u32::try_from(buffer.height()).map_err(|_| Error::LimitExceeded("rendered height exceeds 32-bit dimensions".into()))?,
@@ -124,7 +147,19 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
     out.try_reserve_exact(sample_count).map_err(|e| Error::LimitExceeded(format!("cannot reserve decoded pixels: {e}")))?;
     for pixel in buffer.buf().chunks_exact(channels) {
         let start = out.len();
-        out.extend_from_slice(&pixel[..keep]);
+        out.extend_from_slice(&pixel[..colors]);
+        if let Some(index) = alpha_index {
+            let alpha_value = pixel.get(index).copied().ok_or_else(|| Error::Malformed("decoded frame is missing its alpha channel".into()))?;
+            if preserve_source_encoding && alpha_associated {
+                let divisor = alpha_value.clamp(0.0, 1.0);
+                if divisor > 0.0 {
+                    for channel in &mut out[start..] {
+                        *channel /= divisor;
+                    }
+                }
+            }
+            out.push(alpha_value);
+        }
         if via_srgb {
             let [r, g, b] = [out[start], out[start + 1], out[start + 2]];
             out[start] = 0.627_403_9 * r + 0.329_283_04 * g + 0.043_313_067 * b;
@@ -136,7 +171,9 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
     let metadata = Metadata {
         icc: source_icc,
         exif: match image.aux_boxes().first_exif() {
-            Ok(AuxBoxData::Data(exif)) => Some(exif.payload().to_vec()),
+            Ok(AuxBoxData::Data(exif)) => {
+                usize::try_from(exif.tiff_header_offset()).ok().and_then(|offset| exif.payload().get(offset..)).map(ToOwned::to_owned)
+            }
             _ => None,
         },
         xmp: match image.aux_boxes().first_xml() {
@@ -144,6 +181,7 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
             _ => None,
         },
     };
+    let pixels_linear_rec2020 = output_color == OutputColor::LinearRec2020 && !has_source_icc && !preserve_source_encoding;
     Ok(Decoded {
         width: rendered_width,
         height: rendered_height,
@@ -156,9 +194,66 @@ fn decode_inner(bytes: &[u8], limits: DecodeLimits, output_color: OutputColor) -
         rendered_cicp: image.rendered_cicp(),
         rendered_icc: image.rendered_icc(),
         xyb_encoded,
+        pixels_linear_rec2020,
+        pixels_premultiplied: alpha_associated && !preserve_source_encoding,
     })
 }
 
+fn first_alpha(channels: &[ExtraChannelType]) -> Option<(usize, bool)> {
+    channels.iter().enumerate().find_map(|(index, channel)| match channel {
+        ExtraChannelType::Alpha { alpha_associated } => Some((index, *alpha_associated)),
+        _ => None,
+    })
+}
+
+fn has_black_channel(channels: &[ExtraChannelType]) -> bool {
+    channels.iter().any(|channel| matches!(channel, ExtraChannelType::Black))
+}
+
+fn enum_encoding_has_supported_curve(encoding: &EnumColourEncoding) -> bool {
+    matches!(encoding.colour_space, ColourSpace::Rgb | ColourSpace::Grey)
+        && match &encoding.tf {
+            TransferFunction::Linear | TransferFunction::Srgb | TransferFunction::Bt709 | TransferFunction::Dci => true,
+            TransferFunction::Gamma { g, .. } => *g > 0,
+            _ => false,
+        }
+}
+
 fn decode_error(error: impl std::fmt::Display) -> Error {
     Error::Malformed(error.to_string())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const STRAIGHT: ExtraChannelType = ExtraChannelType::Alpha { alpha_associated: false };
+    const ASSOCIATED: ExtraChannelType = ExtraChannelType::Alpha { alpha_associated: true };
+    const SPOT: ExtraChannelType = ExtraChannelType::SpotColour { red: 1.0, green: 0.0, blue: 0.0, solidity: 1.0 };
+
+    #[test]
+    fn alpha_is_selected_by_type_and_order() {
+        assert_eq!(first_alpha(&[SPOT, STRAIGHT]), Some((1, false)));
+        assert_eq!(first_alpha(&[STRAIGHT, SPOT]), Some((0, false)));
+        assert_eq!(first_alpha(&[ExtraChannelType::Depth, ASSOCIATED, STRAIGHT]), Some((1, true)));
+        assert_eq!(first_alpha(&[ExtraChannelType::Black, SPOT, ExtraChannelType::Depth]), None);
+    }
+
+    #[test]
+    fn black_extra_channel_is_detected_by_type() {
+        assert!(has_black_channel(&[SPOT, ExtraChannelType::Black]));
+        assert!(!has_black_channel(&[SPOT, STRAIGHT]));
+    }
+
+    #[test]
+    fn only_modeled_nonlinear_enum_encodings_are_preserved_for_unpremultiplication() {
+        let mut encoding = EnumColourEncoding::srgb(RenderingIntent::Relative);
+        assert!(enum_encoding_has_supported_curve(&encoding));
+        encoding.tf = TransferFunction::Pq;
+        assert!(!enum_encoding_has_supported_curve(&encoding));
+        encoding.tf = TransferFunction::Linear;
+        assert!(enum_encoding_has_supported_curve(&encoding));
+        encoding.colour_space = ColourSpace::Xyb;
+        assert!(!enum_encoding_has_supported_curve(&encoding));
+    }
+}

crates/jxl/src/lib.rs

@@ -150,6 +150,12 @@ pub struct Decoded {
     pub rendered_icc: Vec<u8>,
     /// Whether the codestream uses XYB encoding.
     pub xyb_encoded: bool,
+    /// Whether `pixels` are linear Rec. 2020. This is false when associated alpha requires keeping
+    /// a modeled nonlinear source encoding so the caller can unpremultiply before its color transform.
+    pub pixels_linear_rec2020: bool,
+    /// Whether color samples in `pixels` remain associated with alpha. When true, divide them by
+    /// alpha after any necessary transfer conversion; at zero alpha the color values are undefined.
+    pub pixels_premultiplied: bool,
 }
 
 /// Options for bounded decoding.

crates/jxl/tests/roundtrip.rs

@@ -59,6 +59,63 @@ fn lossless_rgba_roundtrip_keeps_alpha_layout() {
     assert_eq!(decoded.metadata.xmp.as_deref(), Some(xmp.as_slice()));
 }
 
+#[test]
+fn exif_payload_starts_at_its_declared_tiff_header_offset() {
+    let image = ImageRef { width: 2, height: 1, layout: Layout::Rgb, samples: Samples::U8(&[10, 20, 30, 40, 50, 60]) };
+    let encoded = encode(&image, EncodeOptions { lossless: true, ..EncodeOptions::default() }, &Metadata::default());
+    assert!(encoded.is_ok(), "valid image should encode: {encoded:?}");
+    let Ok(encoded) = encoded else { return };
+    let raw = raw_codestream(&encoded);
+    let exif = b"II*\0\x08\0\0\0\0\0\0\0";
+
+    for offset in [0_u32, 6, 64] {
+        let mut exif_box = offset.to_be_bytes().to_vec();
+        exif_box.resize(4 + offset as usize, 0xab);
+        exif_box.extend_from_slice(exif);
+        let mut container = b"\0\0\0\x0cJXL \x0d\x0a\x87\x0a".to_vec();
+        box_bytes(&mut container, b"ftyp", b"jxl \0\0\0\0jxl ");
+        box_bytes(&mut container, b"jxlc", &raw);
+        box_bytes(&mut container, b"Exif", &exif_box);
+
+        let decoded = decode(&container, DecodeLimits::default());
+        assert!(decoded.is_ok(), "container with Exif offset {offset} should decode: {decoded:?}");
+        let Ok(decoded) = decoded else { return };
+        assert_eq!(decoded.metadata.exif.as_deref(), Some(exif.as_slice()), "TIFF header offset {offset}");
+    }
+}
+
+#[test]
+#[cfg(not(target_arch = "wasm32"))]
+fn associated_alpha_is_unpremultiplied_before_color_conversion() {
+    let straight = [[200_u8, 100, 50, 128], [30, 180, 240, 192], [100, 80, 60, 255], [220, 110, 40, 160]];
+    let mut premultiplied = Vec::with_capacity(straight.len() * 4);
+    for pixel in &straight {
+        let alpha = f32::from(pixel[3]) / 255.0;
+        premultiplied.extend(pixel[..3].iter().map(|channel| (f32::from(*channel) * alpha).round() as u8));
+        premultiplied.push(pixel[3]);
+    }
+    let metadata = jixel::FlMeta { alpha_premultiplied: true, ..jixel::FlMeta::srgb() };
+    let encoded = jixel::encode_fast_lossless(&premultiplied, 2, 2, jixel::ColorSpace::Rgb, true, &metadata);
+    assert!(encoded.is_ok(), "associated-alpha image should encode: {encoded:?}");
+    let Ok(encoded) = encoded else { return };
+
+    let decoded = decode_with_color(&encoded, DecodeLimits::default(), OutputColor::LinearRec2020);
+    assert!(decoded.is_ok(), "associated-alpha image should decode: {decoded:?}");
+    let Ok(decoded) = decoded else { return };
+    assert_eq!(decoded.layout, Layout::Rgba);
+    assert!(!decoded.pixels_linear_rec2020, "the source encoding is needed for the caller's color transform");
+    assert!(!decoded.pixels_premultiplied, "associated color samples must be divided before nonlinear conversion");
+
+    for (i, source) in straight.iter().enumerate() {
+        let alpha = f32::from(source[3]) / 255.0;
+        assert!((decoded.pixels[i * 4 + 3] - alpha).abs() < 1e-6);
+        for (channel, &value) in source[..3].iter().enumerate() {
+            let expected = f32::from(value) / 255.0;
+            assert!((decoded.pixels[i * 4 + channel] - expected).abs() < 0.01, "pixel {i}, channel {channel}");
+        }
+    }
+}
+
 #[test]
 fn decoder_enforces_input_and_pixel_limits() {
     let samples = [10_u8, 20, 30];