/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 |