summaryrefslogtreecommitdiffstats
path: root/third_party/jpeg-xl/lib/extras/enc/npy.cc
diff options
context:
space:
mode:
Diffstat (limited to 'third_party/jpeg-xl/lib/extras/enc/npy.cc')
-rw-r--r--third_party/jpeg-xl/lib/extras/enc/npy.cc322
1 files changed, 322 insertions, 0 deletions
diff --git a/third_party/jpeg-xl/lib/extras/enc/npy.cc b/third_party/jpeg-xl/lib/extras/enc/npy.cc
new file mode 100644
index 0000000000..e7a659184b
--- /dev/null
+++ b/third_party/jpeg-xl/lib/extras/enc/npy.cc
@@ -0,0 +1,322 @@
+// Copyright (c) the JPEG XL Project Authors. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+#include "lib/extras/enc/npy.h"
+
+#include <jxl/types.h>
+#include <stdio.h>
+
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "lib/extras/packed_image.h"
+
+namespace jxl {
+namespace extras {
+namespace {
+
+// JSON value writing
+
+class JSONField {
+ public:
+ virtual ~JSONField() = default;
+ virtual void Write(std::ostream& o, uint32_t indent) const = 0;
+
+ protected:
+ JSONField() = default;
+};
+
+class JSONValue : public JSONField {
+ public:
+ template <typename T>
+ explicit JSONValue(const T& value) : value_(std::to_string(value)) {}
+
+ explicit JSONValue(const std::string& value) : value_("\"" + value + "\"") {}
+
+ explicit JSONValue(bool value) : value_(value ? "true" : "false") {}
+
+ void Write(std::ostream& o, uint32_t indent) const override { o << value_; }
+
+ private:
+ std::string value_;
+};
+
+class JSONDict : public JSONField {
+ public:
+ JSONDict() = default;
+
+ template <typename T>
+ T* AddEmpty(const std::string& key) {
+ static_assert(std::is_convertible<T*, JSONField*>::value,
+ "T must be a JSONField");
+ T* ret = new T();
+ values_.emplace_back(
+ key, std::unique_ptr<JSONField>(static_cast<JSONField*>(ret)));
+ return ret;
+ }
+
+ template <typename T>
+ void Add(const std::string& key, const T& value) {
+ values_.emplace_back(key, std::unique_ptr<JSONField>(new JSONValue(value)));
+ }
+
+ void Write(std::ostream& o, uint32_t indent) const override {
+ std::string indent_str(indent, ' ');
+ o << "{";
+ bool is_first = true;
+ for (const auto& key_value : values_) {
+ if (!is_first) {
+ o << ",";
+ }
+ is_first = false;
+ o << std::endl << indent_str << " \"" << key_value.first << "\": ";
+ key_value.second->Write(o, indent + 2);
+ }
+ if (!values_.empty()) {
+ o << std::endl << indent_str;
+ }
+ o << "}";
+ }
+
+ private:
+ // Dictionary with order.
+ std::vector<std::pair<std::string, std::unique_ptr<JSONField>>> values_;
+};
+
+class JSONArray : public JSONField {
+ public:
+ JSONArray() = default;
+
+ template <typename T>
+ T* AddEmpty() {
+ static_assert(std::is_convertible<T*, JSONField*>::value,
+ "T must be a JSONField");
+ T* ret = new T();
+ values_.emplace_back(ret);
+ return ret;
+ }
+
+ template <typename T>
+ void Add(const T& value) {
+ values_.emplace_back(new JSONValue(value));
+ }
+
+ void Write(std::ostream& o, uint32_t indent) const override {
+ std::string indent_str(indent, ' ');
+ o << "[";
+ bool is_first = true;
+ for (const auto& value : values_) {
+ if (!is_first) {
+ o << ",";
+ }
+ is_first = false;
+ o << std::endl << indent_str << " ";
+ value->Write(o, indent + 2);
+ }
+ if (!values_.empty()) {
+ o << std::endl << indent_str;
+ }
+ o << "]";
+ }
+
+ private:
+ std::vector<std::unique_ptr<JSONField>> values_;
+};
+
+void GenerateMetadata(const PackedPixelFile& ppf, std::vector<uint8_t>* out) {
+ JSONDict meta;
+ // Same order as in 18181-3 CD.
+
+ // Frames.
+ auto* meta_frames = meta.AddEmpty<JSONArray>("frames");
+ for (size_t i = 0; i < ppf.frames.size(); i++) {
+ auto* frame_i = meta_frames->AddEmpty<JSONDict>();
+ if (ppf.info.have_animation) {
+ frame_i->Add("duration",
+ JSONValue(ppf.frames[i].frame_info.duration * 1.0f *
+ ppf.info.animation.tps_denominator /
+ ppf.info.animation.tps_numerator));
+ }
+
+ frame_i->Add("name", JSONValue(ppf.frames[i].name));
+
+ if (ppf.info.animation.have_timecodes) {
+ frame_i->Add("timecode", JSONValue(ppf.frames[i].frame_info.timecode));
+ }
+ }
+
+#define METADATA(FIELD) meta.Add(#FIELD, ppf.info.FIELD)
+
+ METADATA(intensity_target);
+ METADATA(min_nits);
+ METADATA(relative_to_max_display);
+ METADATA(linear_below);
+
+ if (ppf.info.have_preview) {
+ meta.AddEmpty<JSONDict>("preview");
+ // TODO(veluca): can we have duration/name/timecode here?
+ }
+
+ {
+ auto ectype = meta.AddEmpty<JSONArray>("extra_channel_type");
+ auto bps = meta.AddEmpty<JSONArray>("bits_per_sample");
+ auto ebps = meta.AddEmpty<JSONArray>("exp_bits_per_sample");
+ bps->Add(ppf.info.bits_per_sample);
+ ebps->Add(ppf.info.exponent_bits_per_sample);
+ for (size_t i = 0; i < ppf.extra_channels_info.size(); i++) {
+ switch (ppf.extra_channels_info[i].ec_info.type) {
+ case JXL_CHANNEL_ALPHA: {
+ ectype->Add(std::string("Alpha"));
+ break;
+ }
+ case JXL_CHANNEL_DEPTH: {
+ ectype->Add(std::string("Depth"));
+ break;
+ }
+ case JXL_CHANNEL_SPOT_COLOR: {
+ ectype->Add(std::string("SpotColor"));
+ break;
+ }
+ case JXL_CHANNEL_SELECTION_MASK: {
+ ectype->Add(std::string("SelectionMask"));
+ break;
+ }
+ case JXL_CHANNEL_BLACK: {
+ ectype->Add(std::string("Black"));
+ break;
+ }
+ case JXL_CHANNEL_CFA: {
+ ectype->Add(std::string("CFA"));
+ break;
+ }
+ case JXL_CHANNEL_THERMAL: {
+ ectype->Add(std::string("Thermal"));
+ break;
+ }
+ default: {
+ ectype->Add(std::string("UNKNOWN"));
+ break;
+ }
+ }
+ bps->Add(ppf.extra_channels_info[i].ec_info.bits_per_sample);
+ ebps->Add(ppf.extra_channels_info[i].ec_info.exponent_bits_per_sample);
+ }
+ }
+
+ std::ostringstream os;
+ meta.Write(os, 0);
+ out->resize(os.str().size());
+ memcpy(out->data(), os.str().data(), os.str().size());
+}
+
+void Append(std::vector<uint8_t>* out, const void* data, size_t size) {
+ size_t pos = out->size();
+ out->resize(pos + size);
+ memcpy(out->data() + pos, data, size);
+}
+
+void WriteNPYHeader(size_t xsize, size_t ysize, uint32_t num_channels,
+ size_t num_frames, std::vector<uint8_t>* out) {
+ const uint8_t header[] = "\x93NUMPY\x01\x00";
+ Append(out, header, 8);
+ std::stringstream ss;
+ ss << "{'descr': '<f4', 'fortran_order': False, 'shape': (" << num_frames
+ << ", " << ysize << ", " << xsize << ", " << num_channels << "), }\n";
+ // 16-bit little endian header length.
+ uint8_t header_len[2] = {static_cast<uint8_t>(ss.str().size() % 256),
+ static_cast<uint8_t>(ss.str().size() / 256)};
+ Append(out, header_len, 2);
+ Append(out, ss.str().data(), ss.str().size());
+}
+
+bool WriteFrameToNPYArray(size_t xsize, size_t ysize, const PackedFrame& frame,
+ std::vector<uint8_t>* out) {
+ const auto& color = frame.color;
+ if (color.xsize != xsize || color.ysize != ysize) {
+ return false;
+ }
+ for (const auto& ec : frame.extra_channels) {
+ if (ec.xsize != xsize || ec.ysize != ysize) {
+ return false;
+ }
+ }
+ // interleave the samples from color and extra channels
+ for (size_t y = 0; y < ysize; ++y) {
+ for (size_t x = 0; x < xsize; ++x) {
+ {
+ size_t sample_size = color.pixel_stride();
+ size_t offset = y * color.stride + x * sample_size;
+ uint8_t* pixels = reinterpret_cast<uint8_t*>(color.pixels());
+ JXL_ASSERT(offset + sample_size <= color.pixels_size);
+ Append(out, pixels + offset, sample_size);
+ }
+ for (const auto& ec : frame.extra_channels) {
+ size_t sample_size = ec.pixel_stride();
+ size_t offset = y * ec.stride + x * sample_size;
+ uint8_t* pixels = reinterpret_cast<uint8_t*>(ec.pixels());
+ JXL_ASSERT(offset + sample_size <= ec.pixels_size);
+ Append(out, pixels + offset, sample_size);
+ }
+ }
+ }
+ return true;
+}
+
+// Writes a PackedPixelFile as a numpy 4D ndarray in binary format.
+bool WriteNPYArray(const PackedPixelFile& ppf, std::vector<uint8_t>* out) {
+ size_t xsize = ppf.info.xsize;
+ size_t ysize = ppf.info.ysize;
+ WriteNPYHeader(xsize, ysize,
+ ppf.info.num_color_channels + ppf.extra_channels_info.size(),
+ ppf.frames.size(), out);
+ for (const auto& frame : ppf.frames) {
+ if (!WriteFrameToNPYArray(xsize, ysize, frame, out)) {
+ return false;
+ }
+ }
+ return true;
+}
+
+class NumPyEncoder : public Encoder {
+ public:
+ Status Encode(const PackedPixelFile& ppf, EncodedImage* encoded_image,
+ ThreadPool* pool = nullptr) const override {
+ JXL_RETURN_IF_ERROR(VerifyBasicInfo(ppf.info));
+ GenerateMetadata(ppf, &encoded_image->metadata);
+ encoded_image->bitstreams.emplace_back();
+ if (!WriteNPYArray(ppf, &encoded_image->bitstreams.back())) {
+ return false;
+ }
+ if (ppf.preview_frame) {
+ size_t xsize = ppf.info.preview.xsize;
+ size_t ysize = ppf.info.preview.ysize;
+ WriteNPYHeader(xsize, ysize, ppf.info.num_color_channels, 1,
+ &encoded_image->preview_bitstream);
+ if (!WriteFrameToNPYArray(xsize, ysize, *ppf.preview_frame,
+ &encoded_image->preview_bitstream)) {
+ return false;
+ }
+ }
+ return true;
+ }
+ std::vector<JxlPixelFormat> AcceptedFormats() const override {
+ std::vector<JxlPixelFormat> formats;
+ for (const uint32_t num_channels : {1, 3}) {
+ formats.push_back(JxlPixelFormat{num_channels, JXL_TYPE_FLOAT,
+ JXL_LITTLE_ENDIAN, /*align=*/0});
+ }
+ return formats;
+ }
+};
+
+} // namespace
+
+std::unique_ptr<Encoder> GetNumPyEncoder() {
+ return jxl::make_unique<NumPyEncoder>();
+}
+
+} // namespace extras
+} // namespace jxl
generated by cgit v1.2.3 (git 2.39.1) at 2024-07-25 06:47:38 +0000