Coverage Report

Created: 2023-08-28 07:24

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