Populate duration before load

Author: radarhere

Tests/test_file_jxl_animated.py

@@ -23,8 +23,12 @@ def test_n_frames() -> None:
 
 def test_float_duration() -> None:
     with Image.open("Tests/images/iss634.jxl") as im:
-        im.load()
         assert im.info["duration"] == 70
+        assert im.info["timestamp"] == 0
+
+        im.seek(2)
+        assert im.info["duration"] == 60
+        assert im.info["timestamp"] == 140
 
 
 def test_seek() -> None:

Tests/test_file_jxl_metadata.py

@@ -97,8 +97,8 @@ class JpegXlDecoder:
             def __init__(self, b: bytes) -> None:
                 pass
 
-            def get_info(self) -> tuple[tuple[int, int], str, int, int, int, int]:
-                return ((1, 1), "L", 0, 0, 0, 0)
+            def get_info(self) -> tuple[tuple[int, int], str, int, int, int, int, int]:
+                return ((1, 1), "L", 0, 0, 0, 0, 0)
 
             def get_icc(self) -> None:
                 pass

src/PIL/JpegXlImagePlugin.py

@@ -25,9 +25,7 @@ def _accept(prefix: bytes) -> bool | str:
 class JpegXlImageFile(ImageFile.ImageFile):
     format = "JPEG XL"
     format_description = "JPEG XL image"
-    __loaded = -1
-    __logical_frame = 0
-    __physical_frame = 0
+    __frame = 0
 
     def _open(self) -> None:
         self._decoder = _jpegxl.JpegXlDecoder(self.fp.read())
@@ -39,13 +37,15 @@ def _open(self) -> None:
             tps_num,
             tps_denom,
             self.info["loop"],
+            tps_duration,
         ) = self._decoder.get_info()
 
         self._n_frames = None if self.is_animated else 1
         self._tps_dur_secs = tps_num / tps_denom if tps_denom != 0 else 1
+        self.info["duration"] = 1000 * tps_duration * (1 / self._tps_dur_secs)
 
         # TODO: handle libjxl time codes
-        self.__timestamp = 0
+        self.info["timestamp"] = 0
 
         if icc := self._decoder.get_icc():
             self.info["icc_profile"] = icc
@@ -57,6 +57,7 @@ def _open(self) -> None:
                 self.info["exif"] = exif[exif_start_offset + 4 :]
         if xmp := self._decoder.get_xmp():
             self.info["xmp"] = xmp
+        self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)]
 
     @property
     def n_frames(self) -> int:
@@ -67,72 +68,51 @@ def n_frames(self) -> int:
 
         return self._n_frames
 
-    def _get_next(self) -> tuple[bytes, float, float]:
-        # Get next frame
-        next_frame = self._decoder.get_next()
-        self.__physical_frame += 1
+    def _get_next(self) -> bytes:
+        data, tps_duration, is_last = self._decoder.get_next()
 
-        # this actually means EOF, errors are raised in _jxl
-        if next_frame is None:
-            msg = "failed to decode next frame in JXL file"
-            raise EOFError(msg)
-
-        data, tps_duration, is_last = next_frame
         if is_last and self._n_frames is None:
-            # libjxl said this frame is the last one
-            self._n_frames = self.__physical_frame
+            self._n_frames = self.__frame
 
         # duration in milliseconds
-        duration = 1000 * tps_duration * (1 / self._tps_dur_secs)
-        timestamp = self.__timestamp
-        self.__timestamp += duration
-
-        return data, timestamp, duration
+        self.info["timestamp"] += self.info["duration"]
+        self.info["duration"] = 1000 * tps_duration * (1 / self._tps_dur_secs)
 
-    def _seek(self, frame: int) -> None:
-        if frame == self.__physical_frame:
-            return  # Nothing to do
-        if frame < self.__physical_frame:
-            # also rewind libjxl decoder instance
-            self._decoder.rewind()
-            self.__physical_frame = 0
-            self.__loaded = -1
-            self.__timestamp = 0
-
-        while self.__physical_frame < frame:
-            self._get_next()  # Advance to the requested frame
+        return data
 
     def seek(self, frame: int) -> None:
-        if self._n_frames is None:
-            self.n_frames
         if not self._seek_check(frame):
             return
 
-        # Set logical frame to requested position
-        self.__logical_frame = frame
+        if frame < self.__frame:
+            self.__frame = 0
+            self._decoder.rewind()
+            self.info["timestamp"] = 0
 
-    def load(self) -> Image.core.PixelAccess | None:
-        if self.__loaded != self.__logical_frame:
-            self._seek(self.__logical_frame)
+        while self.__frame < frame:
+            self.__frame += 1
+            self._get_next()
+            if self._n_frames is not None and self._n_frames < frame:
+                msg = "no more images in JPEG XL file"
+                raise EOFError(msg)
 
-            data, self.info["timestamp"], self.info["duration"] = self._get_next()
-            self.__loaded = self.__logical_frame
+        self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)]
+
+    def load(self) -> Image.core.PixelAccess | None:
+        if self.tile:
+            data = self._get_next()
 
-            # Set tile
             if self.fp and self._exclusive_fp:
                 self.fp.close()
-            # this is horribly memory inefficient
-            # you need probably 2*(raw image plane) bytes of memory
             self.fp = BytesIO(data)
-            self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)]
 
         return super().load()
 
     def load_seek(self, pos: int) -> None:
         pass
 
     def tell(self) -> int:
-        return self.__logical_frame
+        return self.__frame
 
 
 Image.register_open(JpegXlImageFile.format, JpegXlImageFile, _accept)

src/_jpegxl.c

@@ -85,10 +85,10 @@ typedef struct {
     Py_ssize_t jxl_data_len;  // length of input jxl bitstream
 
     uint8_t *outbuf;
-    Py_ssize_t outbuf_len;
+    size_t outbuf_len;
 
     uint8_t *jxl_icc;
-    Py_ssize_t jxl_icc_len;
+    size_t jxl_icc_len;
     uint8_t *jxl_exif;
     Py_ssize_t jxl_exif_len;
     uint8_t *jxl_xmp;
@@ -317,7 +317,7 @@ _jxl_decoder_new(PyObject *self, PyObject *args) {
                 continue;
             }
 
-            size_t cur_compr_box_size;
+            uint64_t cur_compr_box_size;
             decp->status = JxlDecoderGetBoxSizeRaw(decp->decoder, &cur_compr_box_size);
             _JXL_CHECK("JxlDecoderGetBoxSizeRaw");
 
@@ -396,16 +396,21 @@ _jxl_decoder_new(PyObject *self, PyObject *args) {
 PyObject *
 _jxl_decoder_get_info(PyObject *self) {
     JpegXlDecoderObject *decp = (JpegXlDecoderObject *)self;
-
+    JxlFrameHeader fhdr = {};
+    if (JxlDecoderGetFrameHeader(decp->decoder, &fhdr) != JXL_DEC_SUCCESS) {
+        PyErr_SetString(PyExc_OSError, "Error determining duration");
+        return NULL;
+    }
     return Py_BuildValue(
-        "(II)sOIII",
+        "(II)sOIIII",
         decp->basic_info.xsize,
         decp->basic_info.ysize,
         decp->mode,
         decp->basic_info.have_animation ? Py_True : Py_False,
         decp->basic_info.animation.tps_numerator,
         decp->basic_info.animation.tps_denominator,
-        decp->basic_info.animation.num_loops
+        decp->basic_info.animation.num_loops,
+        fhdr.duration
     );
 }
 
@@ -426,11 +431,6 @@ _jxl_decoder_get_next(PyObject *self) {
     while (decp->status != JXL_DEC_NEED_IMAGE_OUT_BUFFER) {
         decp->status = JxlDecoderProcessInput(decp->decoder);
 
-        // every frame was decoded successfully
-        if (decp->status == JXL_DEC_SUCCESS) {
-            Py_RETURN_NONE;
-        }
-
         // this should only occur after rewind
         if (decp->status == JXL_DEC_NEED_MORE_INPUT) {
             _jxl_decoder_set_input((PyObject *)decp);
@@ -454,7 +454,7 @@ _jxl_decoder_get_next(PyObject *self) {
         uint8_t *_new_outbuf = realloc(decp->outbuf, decp->outbuf_len);
         if (!_new_outbuf) {
             PyErr_SetString(PyExc_OSError, "failed to allocate outbuf");
-            goto end_with_custom_error;
+            return NULL;
         }
         decp->outbuf = _new_outbuf;
     }
@@ -469,7 +469,7 @@ _jxl_decoder_get_next(PyObject *self) {
 
     if (decp->status != JXL_DEC_FULL_IMAGE) {
         PyErr_SetString(PyExc_OSError, "failed to read next frame");
-        goto end_with_custom_error;
+        return NULL;
     }
 
     bytes = PyBytes_FromStringAndSize((char *)(decp->outbuf), decp->outbuf_len);
@@ -493,13 +493,6 @@ _jxl_decoder_get_next(PyObject *self) {
         decp->status
     );
     PyErr_SetString(PyExc_OSError, err_msg);
-
-end_with_custom_error:
-
-    // no need to deallocate anything here
-    // user can just ignore error
-
-    return NULL;
 }
 
 PyObject *