support other jxl image modes

Tests/test_file_jxl.py

@@ -8,8 +8,10 @@
 
 from PIL import features, Image, ImageCms
 from Tests.helper import (
+    assert_image_equal,
     assert_image_similar,
     assert_image_similar_tofile,
+    hopper,
     skip_unless_feature,
 )
 
@@ -75,13 +77,21 @@ def test_simple():
         info = im._basic_info
         assert info["size"] == (128, 128)
         assert info["bits_per_sample"] == 8
+        assert info["exponent_bits_per_sample"] == 0
+        assert info["intensity_target"] == 255.0
+        assert info["min_nits"] == 0.0
+        assert info["relative_to_max_display"] is False
+        assert info["linear_below"] == 0.0
         assert info["uses_original_profile"] is False
         assert info["preview_size"] is None
         assert info["animation_info"] is None
         assert info["orientation"] == 1
         assert info["num_color_channels"] == 3
         assert info["num_extra_channels"] == 0
         assert info["alpha_bits"] == 0
+        assert info["alpha_exponent_bits"] == 0
+        assert info["alpha_premultiplied"] == 0
+        assert info["intrinsic_size"] == (128, 128)
         assert info["num_frames"] == 1
         assert info["box_types"] == []
         assert im.size == (128, 128)
@@ -241,3 +251,19 @@ def test_boxes_max_count():
         assert len(im._decoder.get_boxes("jxlp", max_count=1)) == 1
         assert len(im._decoder.get_boxes("jxlp", max_count=2)) == 2
         assert len(im._decoder.get_boxes("jxlp", max_count=3)) == 2
+
+
+@pytest.mark.parametrize("mode", ["L", "LA", "RGB", "RGBA", "I;16", "F"])
+def test_modes(mode):
+    suffix = mode.replace(";", "")
+    prefix = f"Tests/images/hopper_{suffix}"
+
+    tgt = hopper(mode)
+    if "A" in mode:
+        tgt.putalpha(hopper("L").transpose(Image.Transpose.ROTATE_90))
+    if mode == "I;16":
+        tgt = tgt.point(lambda val: val * 64)
+
+    with Image.open(prefix + ".jxl") as im:
+        assert im.mode == mode
+        assert_image_equal(im, tgt)

src/PIL/JpegXLImagePlugin.py

@@ -48,14 +48,29 @@ def _open(self) -> None:
             width, height = height, width
         self._size = (width, height)
 
+        bps = self._basic_info["bits_per_sample"]
+        num_channels = self._basic_info["num_color_channels"]
         have_alpha = self._basic_info["alpha_bits"] != 0
-        if self._basic_info["num_color_channels"] == 1:
-            # TODO check bit depth (image may be "I")
-            self._mode = "LA" if have_alpha else "L"
-        elif self._basic_info["num_color_channels"] == 3:
+        is_float = (
+            self._basic_info["exponent_bits_per_sample"] > 0
+            or self._basic_info["alpha_exponent_bits"] > 0
+        )
+        if num_channels == 1 and not have_alpha:
+            if is_float:
+                self._mode = "F"
+            elif bps > 8:
+                self._mode = "I;16"
+            else:
+                self._mode = "L"
+        elif bps > 8 or is_float:
+            msg = "image mode not supported"
+            raise SyntaxError(msg)
+        elif num_channels == 1:
+            self._mode = "LA"
+        elif num_channels == 3:
             self._mode = "RGBA" if have_alpha else "RGB"
         else:
-            msg = "cannot determine image mode"
+            msg = "image mode not supported"
             raise SyntaxError(msg)
 
         self.n_frames = self._basic_info["num_frames"]

src/PIL/_imagingjxl.pyi

@@ -8,15 +8,24 @@ class AnimationInfo(TypedDict):
     have_timecodes: bool
 
 class BasicInfo(TypedDict):
+    have_container: bool
     size: tuple[int, int]
     bits_per_sample: int
+    exponent_bits_per_sample: int
+    intensity_target: float
+    min_nits: float
+    relative_to_max_display: bool
+    linear_below: float
     uses_original_profile: bool
     preview_size: tuple[int, int] | None
     animation_info: AnimationInfo | None
     orientation: int
     num_color_channels: int
     num_extra_channels: int
     alpha_bits: int
+    alpha_exponent_bits: int
+    alpha_premultiplied: bool
+    intrinsic_size: tuple[int, int]
     num_frames: int
     box_types: list[bytes]
 

src/_imagingjxl.c

@@ -155,16 +155,25 @@ PyObject *jxl_decoder_get_info(JxlDecoderObject *self, PyObject *Py_UNUSED(ignor
         goto err;
     }
 
-    self->info = Py_BuildValue("{s(II)sIsNsNsNsisIsIsIsnsN}",
+    self->info = Py_BuildValue("{sNs(II)sIsIsfsfsNsfsNsNsNsisIsIsIsIsNs(II)snsN}",
+            "have_container", PyBool_FromLong(info.have_container),
             "size", info.xsize, info.ysize,
             "bits_per_sample", info.bits_per_sample,
+            "exponent_bits_per_sample", info.exponent_bits_per_sample,
+            "intensity_target", info.intensity_target,
+            "min_nits", info.min_nits,
+            "relative_to_max_display", PyBool_FromLong(info.relative_to_max_display),
+            "linear_below", info.linear_below,
             "uses_original_profile", PyBool_FromLong(info.uses_original_profile),
             "preview_size", preview_size,
             "animation_info", animation_info,
             "orientation", (int) info.orientation,
             "num_color_channels", info.num_color_channels,
             "num_extra_channels", info.num_extra_channels,
             "alpha_bits", info.alpha_bits,
+            "alpha_exponent_bits", info.alpha_exponent_bits,
+            "alpha_premultiplied", PyBool_FromLong(info.alpha_premultiplied),
+            "intrinsic_size", info.intrinsic_xsize, info.intrinsic_ysize,
             "num_frames", num_frames,
             "box_types", box_types);
     if (!self->info) {
@@ -318,22 +327,55 @@ PyObject *jxl_decoder_get_icc_profile(JxlDecoderObject *self, PyObject *Py_UNUSE
     return icc_data;
 }
 
-void _jxl_decoder_image_out_callback_RGB(Imaging im, size_t x, size_t y, size_t num_pixels, const UINT8 *pixels) {
-    if (x >= im->xsize || y >= im->ysize) {
-        return;  /* TODO set error? */
+void _jxl_decoder_image_out_callback_L(Imaging im, size_t x, size_t y, size_t num_pixels, const UINT8 *pixels) {
+    if (x < im->xsize && y < im->ysize) {
+        if (num_pixels > im->xsize - x) {
+            num_pixels = im->xsize - x;
+        }
+        memcpy((UINT8 *)im->image8[y] + x, pixels, num_pixels);
+    }
+}
+
+void _jxl_decoder_image_out_callback_LA(Imaging im, size_t x, size_t y, size_t num_pixels, const UINT8 *pixels) {
+    if (x < im->xsize && y < im->ysize) {
+        if (num_pixels > im->xsize - x) {
+            num_pixels = im->xsize - x;
+        }
+        size_t w = x + num_pixels;
+        UINT8 *target = (UINT8 *)im->image[y] + x * 4;
+        for (; x < w; x++) {
+            target[0] = target[1] = target[2] = pixels[0];
+            target[3] = pixels[1];
+            target += 4;
+            pixels += 2;
+        }
+    }
+}
+
+void _jxl_decoder_image_out_callback_I16(Imaging im, size_t x, size_t y, size_t num_pixels, const UINT16 *pixels) {
+    if (x < im->xsize && y < im->ysize) {
+        if (num_pixels > im->xsize - x) {
+            num_pixels = im->xsize - x;
+        }
+        memcpy((UINT16 *)im->image[y] + x, pixels, num_pixels * 2);
     }
-    if (num_pixels >= im->xsize - x) {
-        num_pixels = im->xsize - x;
+}
+
+void _jxl_decoder_image_out_callback_F(Imaging im, size_t x, size_t y, size_t num_pixels, const FLOAT32 *pixels) {
+    if (x < im->xsize && y < im->ysize) {
+        if (num_pixels > im->xsize - x) {
+            num_pixels = im->xsize - x;
+        }
+        memcpy((FLOAT32 *)im->image32[y] + x, pixels, num_pixels * 4);
     }
-    size_t w = x + num_pixels;
-    UINT8 *row = (UINT8 *)im->image[y] + x * 4;
-    for (; x < w; x++) {
-        row[0] = pixels[0];
-        row[1] = pixels[1];
-        row[2] = pixels[2];
-        row[3] = 0xff;
-        row += 4;
-        pixels += 3;
+}
+
+void _jxl_decoder_image_out_callback_RGBA(Imaging im, size_t x, size_t y, size_t num_pixels, const UINT8 *pixels) {
+    if (x < im->xsize && y < im->ysize) {
+        if (num_pixels > im->xsize - x) {
+            num_pixels = im->xsize - x;
+        }
+        memcpy((UINT8 *)im->image[y] + x * 4, pixels, num_pixels * 4);
     }
 }
 
@@ -342,10 +384,29 @@ PyObject *jxl_decoder_next(JxlDecoderObject *self, PyObject *arg) {
     if (PyErr_Occurred()) {
         return NULL;
     }
-
-    /* TODO figure out format and callback */
-    JxlPixelFormat format = {3, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
-    JxlImageOutCallback callback = (JxlImageOutCallback)&_jxl_decoder_image_out_callback_RGB;
+    Imaging im = (Imaging)id;
+
+    /* supported modes: L, LA, I;16, F, RGB, RGBA */
+    JxlPixelFormat format = {1, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
+    JxlImageOutCallback callback;
+    if (!strcmp(im->mode, "L")) {
+        callback = (JxlImageOutCallback)_jxl_decoder_image_out_callback_L;
+    } else if (!strcmp(im->mode, "LA")) {
+        format.num_channels = 2;
+        callback = (JxlImageOutCallback)_jxl_decoder_image_out_callback_LA;
+    } else if (!strcmp(im->mode, "I;16")) {
+        format.data_type = JXL_TYPE_UINT16;
+        callback = (JxlImageOutCallback)_jxl_decoder_image_out_callback_I16;
+    } else if (!strcmp(im->mode, "F")) {
+        format.data_type = JXL_TYPE_FLOAT;
+        callback = (JxlImageOutCallback)_jxl_decoder_image_out_callback_F;
+    } else if (!strcmp(im->mode, "RGB") || !strcmp(im->mode, "RGBA")) {
+        format.num_channels = 4;  /* pad RGB with 0xff */
+        callback = (JxlImageOutCallback)_jxl_decoder_image_out_callback_RGBA;
+    } else {
+        PyErr_SetString(PyExc_ValueError, "image buffer mode not supported");
+        return NULL;
+    }
 
     JxlFrameHeader header;
     PyObject *frame_name = NULL;