1 files changed, 183 insertions, 15 deletions

diff --git a/lib/extras/packed_image.h b/lib/extras/packed_image.h
index 1296472..d3ba9ce 100644
--- a/lib/extras/packed_image.h
+++ b/lib/extras/packed_image.h
@@ -9,20 +9,24 @@
 // Helper class for storing external (int or float, interleaved) images. This is
 // the common format used by other libraries and in the libjxl API.
 
+#include <jxl/codestream_header.h>
+#include <jxl/encode.h>
+#include <jxl/types.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
 
 #include <algorithm>
+#include <cmath>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
-#include "jxl/codestream_header.h"
-#include "jxl/encode.h"
-#include "jxl/types.h"
-#include "lib/jxl/common.h"
+#include "lib/jxl/base/byte_order.h"
+#include "lib/jxl/base/common.h"
+#include "lib/jxl/base/status.h"
 
 namespace jxl {
 namespace extras {
@@ -33,9 +37,26 @@ class PackedImage {
   PackedImage(size_t xsize, size_t ysize, const JxlPixelFormat& format)
       : PackedImage(xsize, ysize, format, CalcStride(format, xsize)) {}
 
+  PackedImage Copy() const {
+    PackedImage copy(xsize, ysize, format);
+    memcpy(reinterpret_cast<uint8_t*>(copy.pixels()),
+           reinterpret_cast<const uint8_t*>(pixels()), pixels_size);
+    return copy;
+  }
+
   // The interleaved pixels as defined in the storage format.
   void* pixels() const { return pixels_.get(); }
 
+  uint8_t* pixels(size_t y, size_t x, size_t c) const {
+    return (reinterpret_cast<uint8_t*>(pixels_.get()) + y * stride +
+            x * pixel_stride_ + c * bytes_per_channel_);
+  }
+
+  const uint8_t* const_pixels(size_t y, size_t x, size_t c) const {
+    return (reinterpret_cast<const uint8_t*>(pixels_.get()) + y * stride +
+            x * pixel_stride_ + c * bytes_per_channel_);
+  }
+
   // The image size in pixels.
   size_t xsize;
   size_t ysize;
@@ -47,10 +68,7 @@ class PackedImage {
   JxlPixelFormat format;
   size_t pixels_size;
 
-  size_t pixel_stride() const {
-    return (BitsPerChannel(format.data_type) * format.num_channels /
-            jxl::kBitsPerByte);
-  }
+  size_t pixel_stride() const { return pixel_stride_; }
 
   static size_t BitsPerChannel(JxlDataType data_type) {
     switch (data_type) {
@@ -67,6 +85,52 @@ class PackedImage {
     }
   }
 
+  float GetPixelValue(size_t y, size_t x, size_t c) const {
+    const uint8_t* data = const_pixels(y, x, c);
+    switch (format.data_type) {
+      case JXL_TYPE_UINT8:
+        return data[0] * (1.0f / 255);
+      case JXL_TYPE_UINT16: {
+        uint16_t val;
+        memcpy(&val, data, 2);
+        return (swap_endianness_ ? JXL_BSWAP16(val) : val) * (1.0f / 65535);
+      }
+      case JXL_TYPE_FLOAT: {
+        float val;
+        memcpy(&val, data, 4);
+        return swap_endianness_ ? BSwapFloat(val) : val;
+      }
+      default:
+        JXL_ABORT("Unhandled JxlDataType");
+    }
+  }
+
+  void SetPixelValue(size_t y, size_t x, size_t c, float val) {
+    uint8_t* data = pixels(y, x, c);
+    switch (format.data_type) {
+      case JXL_TYPE_UINT8:
+        data[0] = Clamp1(std::round(val * 255), 0.0f, 255.0f);
+        break;
+      case JXL_TYPE_UINT16: {
+        uint16_t val16 = Clamp1(std::round(val * 65535), 0.0f, 65535.0f);
+        if (swap_endianness_) {
+          val16 = JXL_BSWAP16(val16);
+        }
+        memcpy(data, &val16, 2);
+        break;
+      }
+      case JXL_TYPE_FLOAT: {
+        if (swap_endianness_) {
+          val = BSwapFloat(val);
+        }
+        memcpy(data, &val, 4);
+        break;
+      }
+      default:
+        JXL_ABORT("Unhandled JxlDataType");
+    }
+  }
+
  private:
   PackedImage(size_t xsize, size_t ysize, const JxlPixelFormat& format,
               size_t stride)
@@ -75,7 +139,11 @@ class PackedImage {
         stride(stride),
         format(format),
         pixels_size(ysize * stride),
-        pixels_(malloc(std::max<size_t>(1, pixels_size)), free) {}
+        pixels_(malloc(std::max<size_t>(1, pixels_size)), free) {
+    bytes_per_channel_ = BitsPerChannel(format.data_type) / jxl::kBitsPerByte;
+    pixel_stride_ = format.num_channels * bytes_per_channel_;
+    swap_endianness_ = SwapEndianness(format.endianness);
+  }
 
   static size_t CalcStride(const JxlPixelFormat& format, size_t xsize) {
     size_t stride = xsize * (BitsPerChannel(format.data_type) *
@@ -86,6 +154,9 @@ class PackedImage {
     return stride;
   }
 
+  size_t bytes_per_channel_;
+  size_t pixel_stride_;
+  bool swap_endianness_;
   std::unique_ptr<void, decltype(free)*> pixels_;
 };
 
@@ -98,6 +169,18 @@ class PackedFrame {
   template <typename... Args>
   explicit PackedFrame(Args&&... args) : color(std::forward<Args>(args)...) {}
 
+  PackedFrame Copy() const {
+    PackedFrame copy(color.xsize, color.ysize, color.format);
+    copy.frame_info = frame_info;
+    copy.name = name;
+    copy.color = color.Copy();
+    for (size_t i = 0; i < extra_channels.size(); ++i) {
+      PackedImage ec = extra_channels[i].Copy();
+      copy.extra_channels.emplace_back(std::move(ec));
+    }
+    return copy;
+  }
+
   // The Frame metadata.
   JxlFrameHeader frame_info = {};
   std::string name;
@@ -108,6 +191,85 @@ class PackedFrame {
   std::vector<PackedImage> extra_channels;
 };
 
+class ChunkedPackedFrame {
+ public:
+  typedef void (*ReadLine)(void* opaque, size_t xpos, size_t ypos, size_t xsize,
+                           uint8_t* buffer, size_t len);
+  ChunkedPackedFrame(size_t xsize, size_t ysize, const JxlPixelFormat& format,
+                     void* opaque, ReadLine read_line)
+      : xsize(xsize),
+        ysize(ysize),
+        format(format),
+        opaque_(opaque),
+        read_line_(read_line) {}
+
+  JxlChunkedFrameInputSource GetInputSource() {
+    return JxlChunkedFrameInputSource{this,
+                                      GetColorChannelsPixelFormat,
+                                      GetColorChannelDataAt,
+                                      GetExtraChannelPixelFormat,
+                                      GetExtraChannelDataAt,
+                                      ReleaseCurrentData};
+  }
+
+  // The Frame metadata.
+  JxlFrameHeader frame_info = {};
+  std::string name;
+
+  size_t xsize;
+  size_t ysize;
+  JxlPixelFormat format;
+
+ private:
+  static void GetColorChannelsPixelFormat(void* opaque,
+                                          JxlPixelFormat* pixel_format) {
+    ChunkedPackedFrame* self = reinterpret_cast<ChunkedPackedFrame*>(opaque);
+    *pixel_format = self->format;
+  }
+
+  static const void* GetColorChannelDataAt(void* opaque, size_t xpos,
+                                           size_t ypos, size_t xsize,
+                                           size_t ysize, size_t* row_offset) {
+    ChunkedPackedFrame* self = reinterpret_cast<ChunkedPackedFrame*>(opaque);
+    size_t bytes_per_channel =
+        PackedImage::BitsPerChannel(self->format.data_type) / jxl::kBitsPerByte;
+    size_t bytes_per_pixel = bytes_per_channel * self->format.num_channels;
+    *row_offset = xsize * bytes_per_pixel;
+    uint8_t* buffer = reinterpret_cast<uint8_t*>(malloc(ysize * (*row_offset)));
+    for (size_t y = 0; y < ysize; ++y) {
+      self->read_line_(self->opaque_, xpos, ypos + y, xsize,
+                       &buffer[y * (*row_offset)], *row_offset);
+    }
+    self->buffers_.insert(buffer);
+    return buffer;
+  }
+
+  static void GetExtraChannelPixelFormat(void* opaque, size_t ec_index,
+                                         JxlPixelFormat* pixel_format) {
+    JXL_ABORT("Not implemented");
+  }
+
+  static const void* GetExtraChannelDataAt(void* opaque, size_t ec_index,
+                                           size_t xpos, size_t ypos,
+                                           size_t xsize, size_t ysize,
+                                           size_t* row_offset) {
+    JXL_ABORT("Not implemented");
+  }
+
+  static void ReleaseCurrentData(void* opaque, const void* buffer) {
+    ChunkedPackedFrame* self = reinterpret_cast<ChunkedPackedFrame*>(opaque);
+    auto iter = self->buffers_.find(const_cast<void*>(buffer));
+    if (iter != self->buffers_.end()) {
+      free(*iter);
+      self->buffers_.erase(iter);
+    }
+  }
+
+  void* opaque_;
+  ReadLine read_line_;
+  std::set<void*> buffers_;
+};
+
 // Optional metadata associated with a file
 class PackedMetadata {
  public:
@@ -117,17 +279,18 @@ class PackedMetadata {
   std::vector<uint8_t> xmp;
 };
 
+// The extra channel metadata information.
+struct PackedExtraChannel {
+  JxlExtraChannelInfo ec_info;
+  size_t index;
+  std::string name;
+};
+
 // Helper class representing a JXL image file as decoded to pixels from the API.
 class PackedPixelFile {
  public:
   JxlBasicInfo info = {};
 
-  // The extra channel metadata information.
-  struct PackedExtraChannel {
-    JxlExtraChannelInfo ec_info;
-    size_t index;
-    std::string name;
-  };
   std::vector<PackedExtraChannel> extra_channels_info;
 
   // Color information of the decoded pixels.
@@ -139,9 +302,14 @@ class PackedPixelFile {
 
   std::unique_ptr<PackedFrame> preview_frame;
   std::vector<PackedFrame> frames;
+  mutable std::vector<ChunkedPackedFrame> chunked_frames;
 
   PackedMetadata metadata;
   PackedPixelFile() { JxlEncoderInitBasicInfo(&info); };
+
+  size_t num_frames() const {
+    return chunked_frames.empty() ? frames.size() : chunked_frames.size();
+  }
 };
 
 }  // namespace extras