libjxl

FORK: libjxl patches used on blog
git clone https://git.neptards.moe/blog/libjxl.git
Log | Files | Refs | Submodules | README | LICENSE

npy.cc (9641B)

      1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
      2 //
      3 // Use of this source code is governed by a BSD-style
      4 // license that can be found in the LICENSE file.
      5 
      6 #include "lib/extras/enc/npy.h"
      7 
      8 #include <jxl/types.h>
      9 
     10 #include <memory>
     11 #include <sstream>
     12 #include <string>
     13 #include <vector>
     14 
     15 #include "lib/extras/packed_image.h"
     16 #include "lib/jxl/base/common.h"
     17 
     18 namespace jxl {
     19 namespace extras {
     20 namespace {
     21 
     22 // JSON value writing
     23 
     24 class JSONField {
     25  public:
     26   virtual ~JSONField() = default;
     27   virtual void Write(std::ostream& o, uint32_t indent) const = 0;
     28 
     29  protected:
     30   JSONField() = default;
     31 };
     32 
     33 class JSONValue : public JSONField {
     34  public:
     35   template <typename T>
     36   explicit JSONValue(const T& value) : value_(std::to_string(value)) {}
     37 
     38   explicit JSONValue(const std::string& value) : value_("\"" + value + "\"") {}
     39 
     40   explicit JSONValue(bool value) : value_(value ? "true" : "false") {}
     41 
     42   void Write(std::ostream& o, uint32_t indent) const override { o << value_; }
     43 
     44  private:
     45   std::string value_;
     46 };
     47 
     48 class JSONDict : public JSONField {
     49  public:
     50   JSONDict() = default;
     51 
     52   template <typename T>
     53   T* AddEmpty(const std::string& key) {
     54     static_assert(std::is_convertible<T*, JSONField*>::value,
     55                   "T must be a JSONField");
     56     T* ret = new T();
     57     JSONField* field = static_cast<JSONField*>(ret);
     58     auto handle = std::unique_ptr<JSONField>(field);
     59     values_.emplace_back(key, std::move(handle));
     60     return ret;
     61   }
     62 
     63   template <typename T>
     64   void Add(const std::string& key, const T& value) {
     65     JSONField* field = static_cast<JSONField*>(new JSONValue(value));
     66     auto handle = std::unique_ptr<JSONField>(field);
     67     values_.emplace_back(key, std::move(handle));
     68   }
     69 
     70   void Write(std::ostream& o, uint32_t indent) const override {
     71     std::string indent_str(indent, ' ');
     72     o << "{";
     73     bool is_first = true;
     74     for (const auto& key_value : values_) {
     75       if (!is_first) {
     76         o << ",";
     77       }
     78       is_first = false;
     79       o << "\n" << indent_str << "  \"" << key_value.first << "\": ";
     80       key_value.second->Write(o, indent + 2);
     81     }
     82     if (!values_.empty()) {
     83       o << "\n" << indent_str;
     84     }
     85     o << "}";
     86   }
     87 
     88  private:
     89   // Dictionary with order.
     90   std::vector<std::pair<std::string, std::unique_ptr<JSONField>>> values_;
     91 };
     92 
     93 class JSONArray : public JSONField {
     94  public:
     95   JSONArray() = default;
     96 
     97   template <typename T>
     98   T* AddEmpty() {
     99     static_assert(std::is_convertible<T*, JSONField*>::value,
    100                   "T must be a JSONField");
    101     T* ret = new T();
    102     values_.emplace_back(ret);
    103     return ret;
    104   }
    105 
    106   template <typename T>
    107   void Add(const T& value) {
    108     values_.emplace_back(new JSONValue(value));
    109   }
    110 
    111   void Write(std::ostream& o, uint32_t indent) const override {
    112     std::string indent_str(indent, ' ');
    113     o << "[";
    114     bool is_first = true;
    115     for (const auto& value : values_) {
    116       if (!is_first) {
    117         o << ",";
    118       }
    119       is_first = false;
    120       o << "\n" << indent_str << "  ";
    121       value->Write(o, indent + 2);
    122     }
    123     if (!values_.empty()) {
    124       o << "\n" << indent_str;
    125     }
    126     o << "]";
    127   }
    128 
    129  private:
    130   std::vector<std::unique_ptr<JSONField>> values_;
    131 };
    132 
    133 void GenerateMetadata(const PackedPixelFile& ppf, std::vector<uint8_t>* out) {
    134   JSONDict meta;
    135   // Same order as in 18181-3 CD.
    136 
    137   // Frames.
    138   auto* meta_frames = meta.AddEmpty<JSONArray>("frames");
    139   for (size_t i = 0; i < ppf.frames.size(); i++) {
    140     auto* frame_i = meta_frames->AddEmpty<JSONDict>();
    141     if (ppf.info.have_animation) {
    142       frame_i->Add("duration",
    143                    JSONValue(ppf.frames[i].frame_info.duration * 1.0f *
    144                              ppf.info.animation.tps_denominator /
    145                              ppf.info.animation.tps_numerator));
    146     }
    147 
    148     frame_i->Add("name", JSONValue(ppf.frames[i].name));
    149 
    150     if (ppf.info.animation.have_timecodes) {
    151       frame_i->Add("timecode", JSONValue(ppf.frames[i].frame_info.timecode));
    152     }
    153   }
    154 
    155 #define METADATA(FIELD) meta.Add(#FIELD, ppf.info.FIELD)
    156 
    157   METADATA(intensity_target);
    158   METADATA(min_nits);
    159   METADATA(relative_to_max_display);
    160   METADATA(linear_below);
    161 
    162   if (ppf.info.have_preview) {
    163     meta.AddEmpty<JSONDict>("preview");
    164     // TODO(veluca): can we have duration/name/timecode here?
    165   }
    166 
    167   {
    168     auto* ectype = meta.AddEmpty<JSONArray>("extra_channel_type");
    169     auto* bps = meta.AddEmpty<JSONArray>("bits_per_sample");
    170     auto* ebps = meta.AddEmpty<JSONArray>("exp_bits_per_sample");
    171     bps->Add(ppf.info.bits_per_sample);
    172     ebps->Add(ppf.info.exponent_bits_per_sample);
    173     for (const auto& eci : ppf.extra_channels_info) {
    174       switch (eci.ec_info.type) {
    175         case JXL_CHANNEL_ALPHA: {
    176           ectype->Add(std::string("Alpha"));
    177           break;
    178         }
    179         case JXL_CHANNEL_DEPTH: {
    180           ectype->Add(std::string("Depth"));
    181           break;
    182         }
    183         case JXL_CHANNEL_SPOT_COLOR: {
    184           ectype->Add(std::string("SpotColor"));
    185           break;
    186         }
    187         case JXL_CHANNEL_SELECTION_MASK: {
    188           ectype->Add(std::string("SelectionMask"));
    189           break;
    190         }
    191         case JXL_CHANNEL_BLACK: {
    192           ectype->Add(std::string("Black"));
    193           break;
    194         }
    195         case JXL_CHANNEL_CFA: {
    196           ectype->Add(std::string("CFA"));
    197           break;
    198         }
    199         case JXL_CHANNEL_THERMAL: {
    200           ectype->Add(std::string("Thermal"));
    201           break;
    202         }
    203         default: {
    204           ectype->Add(std::string("UNKNOWN"));
    205           break;
    206         }
    207       }
    208       bps->Add(eci.ec_info.bits_per_sample);
    209       ebps->Add(eci.ec_info.exponent_bits_per_sample);
    210     }
    211   }
    212 
    213   std::ostringstream os;
    214   meta.Write(os, 0);
    215   out->resize(os.str().size());
    216   memcpy(out->data(), os.str().data(), os.str().size());
    217 }
    218 
    219 void Append(std::vector<uint8_t>* out, const void* data, size_t size) {
    220   size_t pos = out->size();
    221   out->resize(pos + size);
    222   memcpy(out->data() + pos, data, size);
    223 }
    224 
    225 void WriteNPYHeader(size_t xsize, size_t ysize, uint32_t num_channels,
    226                     size_t num_frames, std::vector<uint8_t>* out) {
    227   const uint8_t header[] = "\x93NUMPY\x01\x00";
    228   Append(out, header, 8);
    229   std::stringstream ss;
    230   ss << "{'descr': '<f4', 'fortran_order': False, 'shape': (" << num_frames
    231      << ", " << ysize << ", " << xsize << ", " << num_channels << "), }\n";
    232   // 16-bit little endian header length.
    233   uint8_t header_len[2] = {static_cast<uint8_t>(ss.str().size() % 256),
    234                            static_cast<uint8_t>(ss.str().size() / 256)};
    235   Append(out, header_len, 2);
    236   Append(out, ss.str().data(), ss.str().size());
    237 }
    238 
    239 bool WriteFrameToNPYArray(size_t xsize, size_t ysize, const PackedFrame& frame,
    240                           std::vector<uint8_t>* out) {
    241   const auto& color = frame.color;
    242   if (color.xsize != xsize || color.ysize != ysize) {
    243     return false;
    244   }
    245   for (const auto& ec : frame.extra_channels) {
    246     if (ec.xsize != xsize || ec.ysize != ysize) {
    247       return false;
    248     }
    249   }
    250   // interleave the samples from color and extra channels
    251   for (size_t y = 0; y < ysize; ++y) {
    252     for (size_t x = 0; x < xsize; ++x) {
    253       {
    254         size_t sample_size = color.pixel_stride();
    255         size_t offset = y * color.stride + x * sample_size;
    256         uint8_t* pixels = reinterpret_cast<uint8_t*>(color.pixels());
    257         JXL_ASSERT(offset + sample_size <= color.pixels_size);
    258         Append(out, pixels + offset, sample_size);
    259       }
    260       for (const auto& ec : frame.extra_channels) {
    261         size_t sample_size = ec.pixel_stride();
    262         size_t offset = y * ec.stride + x * sample_size;
    263         uint8_t* pixels = reinterpret_cast<uint8_t*>(ec.pixels());
    264         JXL_ASSERT(offset + sample_size <= ec.pixels_size);
    265         Append(out, pixels + offset, sample_size);
    266       }
    267     }
    268   }
    269   return true;
    270 }
    271 
    272 // Writes a PackedPixelFile as a numpy 4D ndarray in binary format.
    273 bool WriteNPYArray(const PackedPixelFile& ppf, std::vector<uint8_t>* out) {
    274   size_t xsize = ppf.info.xsize;
    275   size_t ysize = ppf.info.ysize;
    276   WriteNPYHeader(xsize, ysize,
    277                  ppf.info.num_color_channels + ppf.extra_channels_info.size(),
    278                  ppf.frames.size(), out);
    279   for (const auto& frame : ppf.frames) {
    280     if (!WriteFrameToNPYArray(xsize, ysize, frame, out)) {
    281       return false;
    282     }
    283   }
    284   return true;
    285 }
    286 
    287 class NumPyEncoder : public Encoder {
    288  public:
    289   Status Encode(const PackedPixelFile& ppf, EncodedImage* encoded_image,
    290                 ThreadPool* pool) const override {
    291     JXL_RETURN_IF_ERROR(VerifyBasicInfo(ppf.info));
    292     GenerateMetadata(ppf, &encoded_image->metadata);
    293     encoded_image->bitstreams.emplace_back();
    294     if (!WriteNPYArray(ppf, &encoded_image->bitstreams.back())) {
    295       return false;
    296     }
    297     if (ppf.preview_frame) {
    298       size_t xsize = ppf.info.preview.xsize;
    299       size_t ysize = ppf.info.preview.ysize;
    300       WriteNPYHeader(xsize, ysize, ppf.info.num_color_channels, 1,
    301                      &encoded_image->preview_bitstream);
    302       if (!WriteFrameToNPYArray(xsize, ysize, *ppf.preview_frame,
    303                                 &encoded_image->preview_bitstream)) {
    304         return false;
    305       }
    306     }
    307     return true;
    308   }
    309   std::vector<JxlPixelFormat> AcceptedFormats() const override {
    310     std::vector<JxlPixelFormat> formats;
    311     for (const uint32_t num_channels : {1, 3}) {
    312       formats.push_back(JxlPixelFormat{num_channels, JXL_TYPE_FLOAT,
    313                                        JXL_LITTLE_ENDIAN, /*align=*/0});
    314     }
    315     return formats;
    316   }
    317 };
    318 
    319 }  // namespace
    320 
    321 std::unique_ptr<Encoder> GetNumPyEncoder() {
    322   return jxl::make_unique<NumPyEncoder>();
    323 }
    324 
    325 }  // namespace extras
    326 }  // namespace jxl