blob: ef5498c27429619191d733270251f95ec1b788e6 [file] [log] [blame]
Dichen Zhang85b37562022-10-11 11:08:28 -07001/*
2 * Copyright 2022 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include <jpegrecoverymap/recoverymap.h>
Nick Deakinf6bca5a2022-11-04 10:43:43 -040018#include <jpegrecoverymap/jpegencoder.h>
19#include <jpegrecoverymap/jpegdecoder.h>
20#include <jpegrecoverymap/recoverymapmath.h>
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +000021#include <jpegrecoverymap/recoverymaputils.h>
Nick Deakinf6bca5a2022-11-04 10:43:43 -040022
Dichen Zhanga8766262022-11-07 23:48:24 +000023#include <image_io/jpeg/jpeg_marker.h>
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +000024#include <image_io/jpeg/jpeg_info.h>
25#include <image_io/jpeg/jpeg_scanner.h>
26#include <image_io/jpeg/jpeg_info_builder.h>
27#include <image_io/base/data_segment_data_source.h>
Nick Deakinf6bca5a2022-11-04 10:43:43 -040028
29#include <memory>
Dichen Zhang72fd2b12022-11-01 06:11:50 +000030#include <sstream>
31#include <string>
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +000032#include <cmath>
Harish Mahendrakar72b6f302022-12-16 10:39:15 -080033#include <condition_variable>
34#include <deque>
35#include <mutex>
36#include <thread>
37#include <unistd.h>
Dichen Zhang72fd2b12022-11-01 06:11:50 +000038
39using namespace std;
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +000040using namespace photos_editing_formats::image_io;
Dichen Zhang85b37562022-10-11 11:08:28 -070041
42namespace android::recoverymap {
43
Harish Mahendrakar555a06b2022-12-14 09:37:27 -080044#define USE_SRGB_INVOETF_LUT 1
45#define USE_HLG_OETF_LUT 1
46#define USE_PQ_OETF_LUT 1
47#define USE_HLG_INVOETF_LUT 1
48#define USE_PQ_INVOETF_LUT 1
Harish Mahendrakarf25991f2022-12-16 11:57:44 -080049#define USE_APPLY_RECOVERY_LUT 1
Harish Mahendrakar555a06b2022-12-14 09:37:27 -080050
Nick Deakinf6bca5a2022-11-04 10:43:43 -040051#define JPEGR_CHECK(x) \
52 { \
53 status_t status = (x); \
54 if ((status) != NO_ERROR) { \
55 return status; \
56 } \
57 }
58
Nick Deakin6bd90432022-11-20 16:26:37 -050059// The current JPEGR version that we encode to
60static const uint32_t kJpegrVersion = 1;
61
Nick Deakinf6bca5a2022-11-04 10:43:43 -040062// Map is quarter res / sixteenth size
63static const size_t kMapDimensionScaleFactor = 4;
Dichen Zhang0b9f7de2022-11-18 06:52:46 +000064// JPEG compress quality (0 ~ 100) for recovery map
65static const int kMapCompressQuality = 85;
Nick Deakinf6bca5a2022-11-04 10:43:43 -040066
Nick Deakin6bd90432022-11-20 16:26:37 -050067// TODO: fill in st2086 metadata
68static const st2086_metadata kSt2086Metadata = {
69 {0.0f, 0.0f},
70 {0.0f, 0.0f},
71 {0.0f, 0.0f},
72 {0.0f, 0.0f},
73 0,
74 1.0f,
75};
Nick Deakinf6bca5a2022-11-04 10:43:43 -040076
Harish Mahendrakar72b6f302022-12-16 10:39:15 -080077#define CONFIG_MULTITHREAD 1
78int GetCPUCoreCount() {
79 int cpuCoreCount = 1;
80#if CONFIG_MULTITHREAD
81#if defined(_SC_NPROCESSORS_ONLN)
82 cpuCoreCount = sysconf(_SC_NPROCESSORS_ONLN);
83#else
84 // _SC_NPROC_ONLN must be defined...
85 cpuCoreCount = sysconf(_SC_NPROC_ONLN);
86#endif
87#endif
88 return cpuCoreCount;
89}
90
Dichen Zhang72fd2b12022-11-01 06:11:50 +000091/*
Dichen Zhanga8766262022-11-07 23:48:24 +000092 * Helper function used for writing data to destination.
93 *
94 * @param destination destination of the data to be written.
95 * @param source source of data being written.
96 * @param length length of the data to be written.
97 * @param position cursor in desitination where the data is to be written.
98 * @return status of succeed or error code.
99 */
100status_t Write(jr_compressed_ptr destination, const void* source, size_t length, int &position) {
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000101 if (position + length > destination->maxLength) {
Dichen Zhanga8766262022-11-07 23:48:24 +0000102 return ERROR_JPEGR_BUFFER_TOO_SMALL;
103 }
104
105 memcpy((uint8_t*)destination->data + sizeof(uint8_t) * position, source, length);
106 position += length;
107 return NO_ERROR;
108}
109
Dichen Zhangd18bc302022-12-16 20:55:24 +0000110status_t Write(jr_exif_ptr destination, const void* source, size_t length, int &position) {
111 memcpy((uint8_t*)destination->data + sizeof(uint8_t) * position, source, length);
112 position += length;
113 return NO_ERROR;
114}
115
116// If the EXIF package doesn't exist in the input JPEG, we'll create one with one entry
117// where the length is represented by this value.
118const size_t PSEUDO_EXIF_PACKAGE_LENGTH = 28;
119// If the EXIF package exists in the input JPEG, we'll add an "JR" entry where the length is
120// represented by this value.
121const size_t EXIF_J_R_ENTRY_LENGTH = 12;
122
123/*
124 * Helper function
125 * Add J R entry to existing exif, or create a new one with J R entry if it's null.
126 * EXIF syntax / change:
127 * ori:
128 * FF E1 - APP1
129 * 01 FC - size of APP1 (to be calculated)
130 * -----------------------------------------------------
131 * 45 78 69 66 00 00 - Exif\0\0 "Exif header"
132 * 49 49 2A 00 - TIFF Header
133 * 08 00 00 00 - offset to the IFD (image file directory)
134 * 06 00 - 6 entries
135 * 00 01 - Width Tag
136 * 03 00 - 'Short' type
137 * 01 00 00 00 - one entry
138 * 00 05 00 00 - image with 0x500
139 *--------------------------------------------------------------------------
140 * new:
141 * FF E1 - APP1
142 * 02 08 - new size, equals to old size + EXIF_J_R_ENTRY_LENGTH (12)
143 *-----------------------------------------------------
144 * 45 78 69 66 00 00 - Exif\0\0 "Exif header"
145 * 49 49 2A 00 - TIFF Header
146 * 08 00 00 00 - offset to the IFD (image file directory)
147 * 07 00 - +1 entry
148 * 4A 52 Custom ('J''R') Tag
149 * 07 00 - Unknown type
150 * 01 00 00 00 - one element
151 * 00 00 00 00 - empty data
152 * 00 01 - Width Tag
153 * 03 00 - 'Short' type
154 * 01 00 00 00 - one entry
155 * 00 05 00 00 - image with 0x500
156 */
157status_t updateExif(jr_exif_ptr exif, jr_exif_ptr dest) {
158 if (exif == nullptr || exif->data == nullptr) {
159 uint8_t data[PSEUDO_EXIF_PACKAGE_LENGTH] = {
160 0x45, 0x78, 0x69, 0x66, 0x00, 0x00,
161 0x49, 0x49, 0x2A, 0x00,
162 0x08, 0x00, 0x00, 0x00,
163 0x01, 0x00,
164 0x4A, 0x52,
165 0x07, 0x00,
166 0x01, 0x00, 0x00, 0x00,
167 0x00, 0x00, 0x00, 0x00};
168 int pos = 0;
169 Write(dest, data, PSEUDO_EXIF_PACKAGE_LENGTH, pos);
170 return NO_ERROR;
171 }
172
173 int num_entry = 0;
174 uint8_t num_entry_low = 0;
175 uint8_t num_entry_high = 0;
176 bool use_big_endian = false;
177 if (reinterpret_cast<uint16_t*>(exif->data)[3] == 0x4949) {
178 num_entry_low = reinterpret_cast<uint8_t*>(exif->data)[14];
179 num_entry_high = reinterpret_cast<uint8_t*>(exif->data)[15];
180 } else if (reinterpret_cast<uint16_t*>(exif->data)[3] == 0x4d4d) {
181 use_big_endian = true;
182 num_entry_high = reinterpret_cast<uint8_t*>(exif->data)[14];
183 num_entry_low = reinterpret_cast<uint8_t*>(exif->data)[15];
184 } else {
185 return ERROR_JPEGR_METADATA_ERROR;
186 }
187 num_entry = (num_entry_high << 8) | num_entry_low;
188 num_entry += 1;
189 num_entry_low = num_entry & 0xff;
190 num_entry_high = (num_entry << 8) & 0xff;
191
192 int pos = 0;
193 Write(dest, (uint8_t*)exif->data, 14, pos);
194
195 if (use_big_endian) {
196 Write(dest, &num_entry_high, 1, pos);
197 Write(dest, &num_entry_low, 1, pos);
198 uint8_t data[EXIF_J_R_ENTRY_LENGTH] = {
199 0x4A, 0x52,
200 0x07, 0x00,
201 0x01, 0x00, 0x00, 0x00,
202 0x00, 0x00, 0x00, 0x00};
203 Write(dest, data, EXIF_J_R_ENTRY_LENGTH, pos);
204 } else {
205 Write(dest, &num_entry_low, 1, pos);
206 Write(dest, &num_entry_high, 1, pos);
207 uint8_t data[EXIF_J_R_ENTRY_LENGTH] = {
208 0x4A, 0x52,
209 0x00, 0x07,
210 0x00, 0x00, 0x00, 0x01,
211 0x00, 0x00, 0x00, 0x00};
212 Write(dest, data, EXIF_J_R_ENTRY_LENGTH, pos);
213 }
214
215 Write(dest, (uint8_t*)exif->data + 16, exif->length - 16, pos);
216
217 return NO_ERROR;
218}
219
Dichen Zhang63d92512023-01-04 12:01:16 -0800220/*
221 * Helper function copies the JPEG image from without EXIF.
222 *
223 * @param dest destination of the data to be written.
224 * @param source source of data being written.
225 * @param exif_pos position of the EXIF package, which is aligned with jpegdecoder.getEXIFPos().
226 * (4 bypes offset to FF sign, the byte after FF E1 XX XX <this byte>).
227 * @param exif_size exif size without the initial 4 bytes, aligned with jpegdecoder.getEXIFSize().
228 */
229void copyJpegWithoutExif(jr_compressed_ptr dest,
230 jr_compressed_ptr source,
231 size_t exif_pos,
232 size_t exif_size) {
233 memcpy(dest, source, sizeof(jpegr_compressed_struct));
234
235 const size_t exif_offset = 4; //exif_pos has 4 bypes offset to the FF sign
236 dest->length = source->length - exif_size - exif_offset;
237 dest->data = malloc(dest->length);
238
239 memcpy(dest->data, source->data, exif_pos - exif_offset);
240 memcpy((uint8_t*)dest->data + exif_pos - exif_offset,
241 (uint8_t*)source->data + exif_pos + exif_size,
242 source->length - exif_pos - exif_size);
243}
244
Dichen Zhang636f5242022-12-07 20:25:44 +0000245/* Encode API-0 */
246status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
247 jpegr_transfer_function hdr_tf,
248 jr_compressed_ptr dest,
249 int quality,
Dichen Zhangd18bc302022-12-16 20:55:24 +0000250 jr_exif_ptr exif) {
Dichen Zhang636f5242022-12-07 20:25:44 +0000251 if (uncompressed_p010_image == nullptr || dest == nullptr) {
252 return ERROR_JPEGR_INVALID_NULL_PTR;
253 }
254
255 if (quality < 0 || quality > 100) {
256 return ERROR_JPEGR_INVALID_INPUT_TYPE;
257 }
258
259 jpegr_metadata metadata;
260 metadata.version = kJpegrVersion;
261 metadata.transferFunction = hdr_tf;
262 if (hdr_tf == JPEGR_TF_PQ) {
263 metadata.hdr10Metadata.st2086Metadata = kSt2086Metadata;
264 }
265
266 jpegr_uncompressed_struct uncompressed_yuv_420_image;
Dichen Zhangc3437ca2023-01-04 14:00:08 -0800267 unique_ptr<uint8_t[]> uncompressed_yuv_420_image_data = make_unique<uint8_t[]>(
268 uncompressed_p010_image->width * uncompressed_p010_image->height * 3 / 2);
269 uncompressed_yuv_420_image.data = uncompressed_yuv_420_image_data.get();
Dichen Zhang636f5242022-12-07 20:25:44 +0000270 JPEGR_CHECK(toneMap(uncompressed_p010_image, &uncompressed_yuv_420_image));
271
272 jpegr_uncompressed_struct map;
273 JPEGR_CHECK(generateRecoveryMap(
274 &uncompressed_yuv_420_image, uncompressed_p010_image, &metadata, &map));
275 std::unique_ptr<uint8_t[]> map_data;
276 map_data.reset(reinterpret_cast<uint8_t*>(map.data));
277
278 jpegr_compressed_struct compressed_map;
279 compressed_map.maxLength = map.width * map.height;
280 unique_ptr<uint8_t[]> compressed_map_data = make_unique<uint8_t[]>(compressed_map.maxLength);
281 compressed_map.data = compressed_map_data.get();
282 JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
283
284 JpegEncoder jpeg_encoder;
285 // TODO: determine ICC data based on color gamut information
286 if (!jpeg_encoder.compressImage(uncompressed_yuv_420_image.data,
287 uncompressed_yuv_420_image.width,
288 uncompressed_yuv_420_image.height, quality, nullptr, 0)) {
289 return ERROR_JPEGR_ENCODE_ERROR;
290 }
291 jpegr_compressed_struct jpeg;
292 jpeg.data = jpeg_encoder.getCompressedImagePtr();
293 jpeg.length = jpeg_encoder.getCompressedImageSize();
294
Dichen Zhangd18bc302022-12-16 20:55:24 +0000295 jpegr_exif_struct new_exif;
Dichen Zhangc3437ca2023-01-04 14:00:08 -0800296 if (exif == nullptr || exif->data == nullptr) {
Dichen Zhangd18bc302022-12-16 20:55:24 +0000297 new_exif.length = PSEUDO_EXIF_PACKAGE_LENGTH;
298 } else {
299 new_exif.length = exif->length + EXIF_J_R_ENTRY_LENGTH;
300 }
301 new_exif.data = new uint8_t[new_exif.length];
302 std::unique_ptr<uint8_t[]> new_exif_data;
303 new_exif_data.reset(reinterpret_cast<uint8_t*>(new_exif.data));
304 JPEGR_CHECK(updateExif(exif, &new_exif));
305
306 JPEGR_CHECK(appendRecoveryMap(&jpeg, &compressed_map, &new_exif, &metadata, dest));
Dichen Zhang636f5242022-12-07 20:25:44 +0000307
308 return NO_ERROR;
309}
310
311/* Encode API-1 */
Dichen Zhang6947d532022-10-22 02:16:21 +0000312status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
313 jr_uncompressed_ptr uncompressed_yuv_420_image,
Nick Deakin6bd90432022-11-20 16:26:37 -0500314 jpegr_transfer_function hdr_tf,
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400315 jr_compressed_ptr dest,
Dichen Zhangffa34012022-11-03 23:21:13 +0000316 int quality,
Dichen Zhangd18bc302022-12-16 20:55:24 +0000317 jr_exif_ptr exif) {
Dichen Zhang6947d532022-10-22 02:16:21 +0000318 if (uncompressed_p010_image == nullptr
319 || uncompressed_yuv_420_image == nullptr
320 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000321 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang6947d532022-10-22 02:16:21 +0000322 }
323
Dichen Zhangffa34012022-11-03 23:21:13 +0000324 if (quality < 0 || quality > 100) {
325 return ERROR_JPEGR_INVALID_INPUT_TYPE;
326 }
327
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400328 if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
329 || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
330 return ERROR_JPEGR_RESOLUTION_MISMATCH;
331 }
332
Nick Deakin6bd90432022-11-20 16:26:37 -0500333 jpegr_metadata metadata;
334 metadata.version = kJpegrVersion;
335 metadata.transferFunction = hdr_tf;
336 if (hdr_tf == JPEGR_TF_PQ) {
337 metadata.hdr10Metadata.st2086Metadata = kSt2086Metadata;
338 }
339
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400340 jpegr_uncompressed_struct map;
Dichen Zhanga8766262022-11-07 23:48:24 +0000341 JPEGR_CHECK(generateRecoveryMap(
Nick Deakin6bd90432022-11-20 16:26:37 -0500342 uncompressed_yuv_420_image, uncompressed_p010_image, &metadata, &map));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400343 std::unique_ptr<uint8_t[]> map_data;
344 map_data.reset(reinterpret_cast<uint8_t*>(map.data));
345
346 jpegr_compressed_struct compressed_map;
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000347 compressed_map.maxLength = map.width * map.height;
348 unique_ptr<uint8_t[]> compressed_map_data = make_unique<uint8_t[]>(compressed_map.maxLength);
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400349 compressed_map.data = compressed_map_data.get();
350 JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
351
352 JpegEncoder jpeg_encoder;
Nick Deakin6bd90432022-11-20 16:26:37 -0500353 // TODO: determine ICC data based on color gamut information
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400354 if (!jpeg_encoder.compressImage(uncompressed_yuv_420_image->data,
355 uncompressed_yuv_420_image->width,
Dichen Zhang95cbb9f2022-11-07 18:32:05 +0000356 uncompressed_yuv_420_image->height, quality, nullptr, 0)) {
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400357 return ERROR_JPEGR_ENCODE_ERROR;
358 }
359 jpegr_compressed_struct jpeg;
360 jpeg.data = jpeg_encoder.getCompressedImagePtr();
361 jpeg.length = jpeg_encoder.getCompressedImageSize();
362
Dichen Zhangd18bc302022-12-16 20:55:24 +0000363 jpegr_exif_struct new_exif;
364 if (exif == nullptr || exif->data == nullptr) {
365 new_exif.length = PSEUDO_EXIF_PACKAGE_LENGTH;
366 } else {
367 new_exif.length = exif->length + EXIF_J_R_ENTRY_LENGTH;
368 }
369
370 new_exif.data = new uint8_t[new_exif.length];
371 std::unique_ptr<uint8_t[]> new_exif_data;
372 new_exif_data.reset(reinterpret_cast<uint8_t*>(new_exif.data));
373 JPEGR_CHECK(updateExif(exif, &new_exif));
374
375 JPEGR_CHECK(appendRecoveryMap(&jpeg, &compressed_map, &new_exif, &metadata, dest));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400376
Dichen Zhang6947d532022-10-22 02:16:21 +0000377 return NO_ERROR;
378}
379
Dichen Zhang636f5242022-12-07 20:25:44 +0000380/* Encode API-2 */
Dichen Zhang6947d532022-10-22 02:16:21 +0000381status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
382 jr_uncompressed_ptr uncompressed_yuv_420_image,
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400383 jr_compressed_ptr compressed_jpeg_image,
Nick Deakin6bd90432022-11-20 16:26:37 -0500384 jpegr_transfer_function hdr_tf,
Dichen Zhang95cbb9f2022-11-07 18:32:05 +0000385 jr_compressed_ptr dest) {
Dichen Zhang6947d532022-10-22 02:16:21 +0000386 if (uncompressed_p010_image == nullptr
387 || uncompressed_yuv_420_image == nullptr
388 || compressed_jpeg_image == nullptr
389 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000390 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang6947d532022-10-22 02:16:21 +0000391 }
392
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400393 if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
394 || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
395 return ERROR_JPEGR_RESOLUTION_MISMATCH;
396 }
397
Nick Deakin6bd90432022-11-20 16:26:37 -0500398 jpegr_metadata metadata;
399 metadata.version = kJpegrVersion;
400 metadata.transferFunction = hdr_tf;
401 if (hdr_tf == JPEGR_TF_PQ) {
402 metadata.hdr10Metadata.st2086Metadata = kSt2086Metadata;
403 }
404
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400405 jpegr_uncompressed_struct map;
Dichen Zhanga8766262022-11-07 23:48:24 +0000406 JPEGR_CHECK(generateRecoveryMap(
Nick Deakin6bd90432022-11-20 16:26:37 -0500407 uncompressed_yuv_420_image, uncompressed_p010_image, &metadata, &map));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400408 std::unique_ptr<uint8_t[]> map_data;
409 map_data.reset(reinterpret_cast<uint8_t*>(map.data));
410
411 jpegr_compressed_struct compressed_map;
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000412 compressed_map.maxLength = map.width * map.height;
413 unique_ptr<uint8_t[]> compressed_map_data = make_unique<uint8_t[]>(compressed_map.maxLength);
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400414 compressed_map.data = compressed_map_data.get();
415 JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
416
Dichen Zhangd18bc302022-12-16 20:55:24 +0000417 // Extract EXIF from JPEG without decoding.
418 JpegDecoder jpeg_decoder;
419 if (!jpeg_decoder.extractEXIF(compressed_jpeg_image->data, compressed_jpeg_image->length)) {
420 return ERROR_JPEGR_DECODE_ERROR;
421 }
422
Dichen Zhang63d92512023-01-04 12:01:16 -0800423 // Update exif.
Dichen Zhangd18bc302022-12-16 20:55:24 +0000424 jpegr_exif_struct exif;
425 exif.data = nullptr;
426 exif.length = 0;
Dichen Zhang63d92512023-01-04 12:01:16 -0800427 jpegr_compressed_struct new_jpeg_image;
428 new_jpeg_image.data = nullptr;
429 new_jpeg_image.length = 0;
Dichen Zhangd18bc302022-12-16 20:55:24 +0000430 if (jpeg_decoder.getEXIFPos() != 0) {
Dichen Zhang63d92512023-01-04 12:01:16 -0800431 copyJpegWithoutExif(&new_jpeg_image,
432 compressed_jpeg_image,
433 jpeg_decoder.getEXIFPos(),
434 jpeg_decoder.getEXIFSize());
Dichen Zhangd18bc302022-12-16 20:55:24 +0000435 exif.data = jpeg_decoder.getEXIFPtr();
436 exif.length = jpeg_decoder.getEXIFSize();
437 }
438
439 jpegr_exif_struct new_exif;
440 if (exif.data == nullptr) {
441 new_exif.length = PSEUDO_EXIF_PACKAGE_LENGTH;
442 } else {
443 new_exif.length = exif.length + EXIF_J_R_ENTRY_LENGTH;
444 }
445
446 new_exif.data = new uint8_t[new_exif.length];
447 std::unique_ptr<uint8_t[]> new_exif_data;
448 new_exif_data.reset(reinterpret_cast<uint8_t*>(new_exif.data));
449 JPEGR_CHECK(updateExif(&exif, &new_exif));
450
451 JPEGR_CHECK(appendRecoveryMap(
Dichen Zhang63d92512023-01-04 12:01:16 -0800452 new_jpeg_image.data == nullptr ? compressed_jpeg_image : &new_jpeg_image,
453 &compressed_map, &new_exif, &metadata, dest));
454
455 if (new_jpeg_image.data != nullptr) {
456 free(new_jpeg_image.data);
457 }
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400458
Dichen Zhang6947d532022-10-22 02:16:21 +0000459 return NO_ERROR;
460}
461
Dichen Zhang636f5242022-12-07 20:25:44 +0000462/* Encode API-3 */
Dichen Zhang6947d532022-10-22 02:16:21 +0000463status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400464 jr_compressed_ptr compressed_jpeg_image,
Nick Deakin6bd90432022-11-20 16:26:37 -0500465 jpegr_transfer_function hdr_tf,
Dichen Zhang95cbb9f2022-11-07 18:32:05 +0000466 jr_compressed_ptr dest) {
Dichen Zhang6947d532022-10-22 02:16:21 +0000467 if (uncompressed_p010_image == nullptr
468 || compressed_jpeg_image == nullptr
469 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000470 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang6947d532022-10-22 02:16:21 +0000471 }
472
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400473 JpegDecoder jpeg_decoder;
474 if (!jpeg_decoder.decompressImage(compressed_jpeg_image->data, compressed_jpeg_image->length)) {
475 return ERROR_JPEGR_DECODE_ERROR;
476 }
477 jpegr_uncompressed_struct uncompressed_yuv_420_image;
478 uncompressed_yuv_420_image.data = jpeg_decoder.getDecompressedImagePtr();
479 uncompressed_yuv_420_image.width = jpeg_decoder.getDecompressedImageWidth();
480 uncompressed_yuv_420_image.height = jpeg_decoder.getDecompressedImageHeight();
Nick Deakin6bd90432022-11-20 16:26:37 -0500481 uncompressed_yuv_420_image.colorGamut = compressed_jpeg_image->colorGamut;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400482
Dichen Zhang63d92512023-01-04 12:01:16 -0800483 // Update exif.
Dichen Zhangd18bc302022-12-16 20:55:24 +0000484 jpegr_exif_struct exif;
485 exif.data = nullptr;
486 exif.length = 0;
Dichen Zhang63d92512023-01-04 12:01:16 -0800487 jpegr_compressed_struct new_jpeg_image;
488 new_jpeg_image.data = nullptr;
489 new_jpeg_image.length = 0;
Dichen Zhangd18bc302022-12-16 20:55:24 +0000490 if (jpeg_decoder.getEXIFPos() != 0) {
Dichen Zhang63d92512023-01-04 12:01:16 -0800491 copyJpegWithoutExif(&new_jpeg_image,
492 compressed_jpeg_image,
493 jpeg_decoder.getEXIFPos(),
494 jpeg_decoder.getEXIFSize());
Dichen Zhangd18bc302022-12-16 20:55:24 +0000495 exif.data = jpeg_decoder.getEXIFPtr();
496 exif.length = jpeg_decoder.getEXIFSize();
497 }
498
499 jpegr_exif_struct new_exif;
500 if (exif.data == nullptr) {
501 new_exif.length = PSEUDO_EXIF_PACKAGE_LENGTH;
502 } else {
503 new_exif.length = exif.length + EXIF_J_R_ENTRY_LENGTH;
504 }
505 new_exif.data = new uint8_t[new_exif.length];
506 std::unique_ptr<uint8_t[]> new_exif_data;
507 new_exif_data.reset(reinterpret_cast<uint8_t*>(new_exif.data));
508 JPEGR_CHECK(updateExif(&exif, &new_exif));
509
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400510 if (uncompressed_p010_image->width != uncompressed_yuv_420_image.width
511 || uncompressed_p010_image->height != uncompressed_yuv_420_image.height) {
512 return ERROR_JPEGR_RESOLUTION_MISMATCH;
513 }
514
Nick Deakin6bd90432022-11-20 16:26:37 -0500515 jpegr_metadata metadata;
516 metadata.version = kJpegrVersion;
517 metadata.transferFunction = hdr_tf;
518 if (hdr_tf == JPEGR_TF_PQ) {
519 metadata.hdr10Metadata.st2086Metadata = kSt2086Metadata;
520 }
521
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400522 jpegr_uncompressed_struct map;
Dichen Zhanga8766262022-11-07 23:48:24 +0000523 JPEGR_CHECK(generateRecoveryMap(
Nick Deakin6bd90432022-11-20 16:26:37 -0500524 &uncompressed_yuv_420_image, uncompressed_p010_image, &metadata, &map));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400525 std::unique_ptr<uint8_t[]> map_data;
526 map_data.reset(reinterpret_cast<uint8_t*>(map.data));
527
528 jpegr_compressed_struct compressed_map;
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000529 compressed_map.maxLength = map.width * map.height;
530 unique_ptr<uint8_t[]> compressed_map_data = make_unique<uint8_t[]>(compressed_map.maxLength);
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400531 compressed_map.data = compressed_map_data.get();
532 JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
533
Dichen Zhangd18bc302022-12-16 20:55:24 +0000534 JPEGR_CHECK(appendRecoveryMap(
Dichen Zhang63d92512023-01-04 12:01:16 -0800535 new_jpeg_image.data == nullptr ? compressed_jpeg_image : &new_jpeg_image,
536 &compressed_map, &new_exif, &metadata, dest));
537
538 if (new_jpeg_image.data != nullptr) {
539 free(new_jpeg_image.data);
540 }
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400541
Dichen Zhang6947d532022-10-22 02:16:21 +0000542 return NO_ERROR;
543}
544
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000545status_t RecoveryMap::getJPEGRInfo(jr_compressed_ptr compressed_jpegr_image,
546 jr_info_ptr jpegr_info) {
547 if (compressed_jpegr_image == nullptr || jpegr_info == nullptr) {
548 return ERROR_JPEGR_INVALID_NULL_PTR;
549 }
550
551 jpegr_compressed_struct primary_image, recovery_map;
552 JPEGR_CHECK(extractPrimaryImageAndRecoveryMap(compressed_jpegr_image,
553 &primary_image, &recovery_map));
554
555 JpegDecoder jpeg_decoder;
556 if (!jpeg_decoder.getCompressedImageParameters(primary_image.data, primary_image.length,
557 &jpegr_info->width, &jpegr_info->height,
558 jpegr_info->iccData, jpegr_info->exifData)) {
559 return ERROR_JPEGR_DECODE_ERROR;
560 }
561
562 return NO_ERROR;
563}
564
Dichen Zhang636f5242022-12-07 20:25:44 +0000565/* Decode API */
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400566status_t RecoveryMap::decodeJPEGR(jr_compressed_ptr compressed_jpegr_image,
Dichen Zhangffa34012022-11-03 23:21:13 +0000567 jr_uncompressed_ptr dest,
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000568 jr_exif_ptr exif,
569 bool request_sdr) {
Dichen Zhang6947d532022-10-22 02:16:21 +0000570 if (compressed_jpegr_image == nullptr || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000571 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang6947d532022-10-22 02:16:21 +0000572 }
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000573 // TODO: fill EXIF data
574 (void) exif;
575
Fyodor Kyslovea9180f2023-01-06 01:11:43 +0000576 if (request_sdr) {
577 JpegDecoder jpeg_decoder;
578 if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length,
579 true)) {
580 return ERROR_JPEGR_DECODE_ERROR;
581 }
582 jpegr_uncompressed_struct uncompressed_rgba_image;
583 uncompressed_rgba_image.data = jpeg_decoder.getDecompressedImagePtr();
584 uncompressed_rgba_image.width = jpeg_decoder.getDecompressedImageWidth();
585 uncompressed_rgba_image.height = jpeg_decoder.getDecompressedImageHeight();
586 memcpy(dest->data, uncompressed_rgba_image.data,
587 uncompressed_rgba_image.width * uncompressed_rgba_image.height * 4);
588 dest->width = uncompressed_rgba_image.width;
589 dest->height = uncompressed_rgba_image.height;
590 return NO_ERROR;
591 }
592
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400593 jpegr_compressed_struct compressed_map;
Nick Deakin6bd90432022-11-20 16:26:37 -0500594 jpegr_metadata metadata;
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000595 JPEGR_CHECK(extractRecoveryMap(compressed_jpegr_image, &compressed_map));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400596
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400597 JpegDecoder jpeg_decoder;
598 if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length)) {
599 return ERROR_JPEGR_DECODE_ERROR;
600 }
601
Fyodor Kyslovbf241572022-12-13 22:38:07 +0000602 JpegDecoder recovery_map_decoder;
Fyodor Kyslovea9180f2023-01-06 01:11:43 +0000603 if (!recovery_map_decoder.decompressImage(compressed_map.data, compressed_map.length)) {
Fyodor Kyslovbf241572022-12-13 22:38:07 +0000604 return ERROR_JPEGR_DECODE_ERROR;
605 }
606
607 jpegr_uncompressed_struct map;
608 map.data = recovery_map_decoder.getDecompressedImagePtr();
609 map.width = recovery_map_decoder.getDecompressedImageWidth();
610 map.height = recovery_map_decoder.getDecompressedImageHeight();
611
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400612 jpegr_uncompressed_struct uncompressed_yuv_420_image;
613 uncompressed_yuv_420_image.data = jpeg_decoder.getDecompressedImagePtr();
614 uncompressed_yuv_420_image.width = jpeg_decoder.getDecompressedImageWidth();
615 uncompressed_yuv_420_image.height = jpeg_decoder.getDecompressedImageHeight();
616
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000617 if (!getMetadataFromXMP(static_cast<uint8_t*>(jpeg_decoder.getXMPPtr()),
Fyodor Kyslovea9180f2023-01-06 01:11:43 +0000618 jpeg_decoder.getXMPSize(), &metadata)) {
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000619 return ERROR_JPEGR_DECODE_ERROR;
620 }
621
Fyodor Kyslovea9180f2023-01-06 01:11:43 +0000622 JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, &metadata, dest));
Dichen Zhang6947d532022-10-22 02:16:21 +0000623 return NO_ERROR;
624}
625
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400626status_t RecoveryMap::compressRecoveryMap(jr_uncompressed_ptr uncompressed_recovery_map,
627 jr_compressed_ptr dest) {
Dichen Zhang596a7562022-10-12 14:57:05 -0700628 if (uncompressed_recovery_map == nullptr || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000629 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang596a7562022-10-12 14:57:05 -0700630 }
631
Nick Deakin6bd90432022-11-20 16:26:37 -0500632 // TODO: should we have ICC data for the map?
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400633 JpegEncoder jpeg_encoder;
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000634 if (!jpeg_encoder.compressImage(uncompressed_recovery_map->data,
635 uncompressed_recovery_map->width,
636 uncompressed_recovery_map->height,
637 kMapCompressQuality,
638 nullptr,
639 0,
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400640 true /* isSingleChannel */)) {
641 return ERROR_JPEGR_ENCODE_ERROR;
642 }
643
Dichen Zhang0b9f7de2022-11-18 06:52:46 +0000644 if (dest->maxLength < jpeg_encoder.getCompressedImageSize()) {
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400645 return ERROR_JPEGR_BUFFER_TOO_SMALL;
646 }
647
648 memcpy(dest->data, jpeg_encoder.getCompressedImagePtr(), jpeg_encoder.getCompressedImageSize());
649 dest->length = jpeg_encoder.getCompressedImageSize();
Nick Deakin6bd90432022-11-20 16:26:37 -0500650 dest->colorGamut = JPEGR_COLORGAMUT_UNSPECIFIED;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400651
Dichen Zhang6947d532022-10-22 02:16:21 +0000652 return NO_ERROR;
Dichen Zhang596a7562022-10-12 14:57:05 -0700653}
654
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800655const int kJobSzInRows = 16;
656static_assert(kJobSzInRows > 0 && kJobSzInRows % kMapDimensionScaleFactor == 0,
657 "align job size to kMapDimensionScaleFactor");
658
659class JobQueue {
660 public:
661 bool dequeueJob(size_t& rowStart, size_t& rowEnd);
662 void enqueueJob(size_t rowStart, size_t rowEnd);
663 void markQueueForEnd();
664 void reset();
665
666 private:
667 bool mQueuedAllJobs = false;
668 std::deque<std::tuple<size_t, size_t>> mJobs;
669 std::mutex mMutex;
670 std::condition_variable mCv;
671};
672
673bool JobQueue::dequeueJob(size_t& rowStart, size_t& rowEnd) {
674 std::unique_lock<std::mutex> lock{mMutex};
675 while (true) {
676 if (mJobs.empty()) {
677 if (mQueuedAllJobs) {
678 return false;
679 } else {
680 mCv.wait(lock);
681 }
682 } else {
683 auto it = mJobs.begin();
684 rowStart = std::get<0>(*it);
685 rowEnd = std::get<1>(*it);
686 mJobs.erase(it);
687 return true;
688 }
689 }
690 return false;
691}
692
693void JobQueue::enqueueJob(size_t rowStart, size_t rowEnd) {
694 std::unique_lock<std::mutex> lock{mMutex};
695 mJobs.push_back(std::make_tuple(rowStart, rowEnd));
696 lock.unlock();
697 mCv.notify_one();
698}
699
700void JobQueue::markQueueForEnd() {
701 std::unique_lock<std::mutex> lock{mMutex};
702 mQueuedAllJobs = true;
703}
704
705void JobQueue::reset() {
706 std::unique_lock<std::mutex> lock{mMutex};
707 mJobs.clear();
708 mQueuedAllJobs = false;
709}
710
Dichen Zhang6947d532022-10-22 02:16:21 +0000711status_t RecoveryMap::generateRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_image,
712 jr_uncompressed_ptr uncompressed_p010_image,
Nick Deakin6bd90432022-11-20 16:26:37 -0500713 jr_metadata_ptr metadata,
714 jr_uncompressed_ptr dest) {
Dichen Zhang596a7562022-10-12 14:57:05 -0700715 if (uncompressed_yuv_420_image == nullptr
716 || uncompressed_p010_image == nullptr
Nick Deakin6bd90432022-11-20 16:26:37 -0500717 || metadata == nullptr
Dichen Zhang596a7562022-10-12 14:57:05 -0700718 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000719 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang596a7562022-10-12 14:57:05 -0700720 }
721
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400722 if (uncompressed_yuv_420_image->width != uncompressed_p010_image->width
723 || uncompressed_yuv_420_image->height != uncompressed_p010_image->height) {
724 return ERROR_JPEGR_RESOLUTION_MISMATCH;
725 }
726
Nick Deakin6bd90432022-11-20 16:26:37 -0500727 if (uncompressed_yuv_420_image->colorGamut == JPEGR_COLORGAMUT_UNSPECIFIED
728 || uncompressed_p010_image->colorGamut == JPEGR_COLORGAMUT_UNSPECIFIED) {
729 return ERROR_JPEGR_INVALID_COLORGAMUT;
730 }
731
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400732 size_t image_width = uncompressed_yuv_420_image->width;
733 size_t image_height = uncompressed_yuv_420_image->height;
734 size_t map_width = image_width / kMapDimensionScaleFactor;
735 size_t map_height = image_height / kMapDimensionScaleFactor;
736
737 dest->width = map_width;
738 dest->height = map_height;
Nick Deakin6bd90432022-11-20 16:26:37 -0500739 dest->colorGamut = JPEGR_COLORGAMUT_UNSPECIFIED;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400740 dest->data = new uint8_t[map_width * map_height];
741 std::unique_ptr<uint8_t[]> map_data;
742 map_data.reset(reinterpret_cast<uint8_t*>(dest->data));
743
Nick Deakin6bd90432022-11-20 16:26:37 -0500744 ColorTransformFn hdrInvOetf = nullptr;
Nick Deakin65f492a2022-11-29 22:47:40 -0500745 float hdr_white_nits = 0.0f;
Nick Deakin6bd90432022-11-20 16:26:37 -0500746 switch (metadata->transferFunction) {
Dichen Zhangdc8452b2022-11-23 17:17:56 +0000747 case JPEGR_TF_LINEAR:
748 hdrInvOetf = identityConversion;
749 break;
Nick Deakin6bd90432022-11-20 16:26:37 -0500750 case JPEGR_TF_HLG:
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800751#if USE_HLG_INVOETF_LUT
752 hdrInvOetf = hlgInvOetfLUT;
753#else
Nick Deakin6bd90432022-11-20 16:26:37 -0500754 hdrInvOetf = hlgInvOetf;
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800755#endif
Nick Deakin65f492a2022-11-29 22:47:40 -0500756 hdr_white_nits = kHlgMaxNits;
Nick Deakin6bd90432022-11-20 16:26:37 -0500757 break;
758 case JPEGR_TF_PQ:
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800759#if USE_PQ_INVOETF_LUT
760 hdrInvOetf = pqInvOetfLUT;
761#else
Nick Deakin6bd90432022-11-20 16:26:37 -0500762 hdrInvOetf = pqInvOetf;
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800763#endif
Nick Deakin65f492a2022-11-29 22:47:40 -0500764 hdr_white_nits = kPqMaxNits;
Nick Deakin6bd90432022-11-20 16:26:37 -0500765 break;
Dichen Zhangb27d06d2022-12-14 19:57:50 +0000766 case JPEGR_TF_UNSPECIFIED:
767 // Should be impossible to hit after input validation.
768 return ERROR_JPEGR_INVALID_TRANS_FUNC;
Nick Deakin6bd90432022-11-20 16:26:37 -0500769 }
770
771 ColorTransformFn hdrGamutConversionFn = getHdrConversionFn(
772 uncompressed_yuv_420_image->colorGamut, uncompressed_p010_image->colorGamut);
773
774 ColorCalculationFn luminanceFn = nullptr;
775 switch (uncompressed_yuv_420_image->colorGamut) {
776 case JPEGR_COLORGAMUT_BT709:
777 luminanceFn = srgbLuminance;
778 break;
779 case JPEGR_COLORGAMUT_P3:
780 luminanceFn = p3Luminance;
781 break;
782 case JPEGR_COLORGAMUT_BT2100:
783 luminanceFn = bt2100Luminance;
784 break;
785 case JPEGR_COLORGAMUT_UNSPECIFIED:
786 // Should be impossible to hit after input validation.
787 return ERROR_JPEGR_INVALID_COLORGAMUT;
788 }
789
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800790 std::mutex mutex;
Nick Deakin594a4ca2022-11-16 20:57:42 -0500791 float hdr_y_nits_max = 0.0f;
Nick Deakin6bd90432022-11-20 16:26:37 -0500792 double hdr_y_nits_avg = 0.0f;
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800793 const int threads = std::clamp(GetCPUCoreCount(), 1, 4);
794 size_t rowStep = threads == 1 ? image_height : kJobSzInRows;
795 JobQueue jobQueue;
Nick Deakin594a4ca2022-11-16 20:57:42 -0500796
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800797 std::function<void()> computeMetadata = [uncompressed_p010_image, hdrInvOetf,
798 hdrGamutConversionFn, luminanceFn, hdr_white_nits,
799 threads, &mutex, &hdr_y_nits_avg,
800 &hdr_y_nits_max, &jobQueue]() -> void {
801 size_t rowStart, rowEnd;
802 float hdr_y_nits_max_th = 0.0f;
803 double hdr_y_nits_avg_th = 0.0f;
804 while (jobQueue.dequeueJob(rowStart, rowEnd)) {
805 for (size_t y = rowStart; y < rowEnd; ++y) {
806 for (size_t x = 0; x < uncompressed_p010_image->width; ++x) {
807 Color hdr_yuv_gamma = getP010Pixel(uncompressed_p010_image, x, y);
808 Color hdr_rgb_gamma = bt2100YuvToRgb(hdr_yuv_gamma);
809 Color hdr_rgb = hdrInvOetf(hdr_rgb_gamma);
810 hdr_rgb = hdrGamutConversionFn(hdr_rgb);
811 float hdr_y_nits = luminanceFn(hdr_rgb) * hdr_white_nits;
812
813 hdr_y_nits_avg_th += hdr_y_nits;
814 if (hdr_y_nits > hdr_y_nits_max_th) {
815 hdr_y_nits_max_th = hdr_y_nits;
816 }
817 }
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400818 }
819 }
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800820 std::unique_lock<std::mutex> lock{mutex};
821 hdr_y_nits_avg += hdr_y_nits_avg_th;
822 hdr_y_nits_max = std::max(hdr_y_nits_max, hdr_y_nits_max_th);
823 };
824
825 std::function<void()> generateMap = [uncompressed_yuv_420_image, uncompressed_p010_image,
826 metadata, dest, hdrInvOetf, hdrGamutConversionFn,
827 luminanceFn, hdr_white_nits, &jobQueue]() -> void {
828 size_t rowStart, rowEnd;
829 while (jobQueue.dequeueJob(rowStart, rowEnd)) {
830 for (size_t y = rowStart; y < rowEnd; ++y) {
831 for (size_t x = 0; x < dest->width; ++x) {
832 Color sdr_yuv_gamma =
833 sampleYuv420(uncompressed_yuv_420_image, kMapDimensionScaleFactor, x, y);
834 Color sdr_rgb_gamma = srgbYuvToRgb(sdr_yuv_gamma);
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800835#if USE_SRGB_INVOETF_LUT
836 Color sdr_rgb = srgbInvOetfLUT(sdr_rgb_gamma);
837#else
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800838 Color sdr_rgb = srgbInvOetf(sdr_rgb_gamma);
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800839#endif
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800840 float sdr_y_nits = luminanceFn(sdr_rgb) * kSdrWhiteNits;
841
842 Color hdr_yuv_gamma = sampleP010(uncompressed_p010_image, kMapDimensionScaleFactor, x, y);
843 Color hdr_rgb_gamma = bt2100YuvToRgb(hdr_yuv_gamma);
844 Color hdr_rgb = hdrInvOetf(hdr_rgb_gamma);
845 hdr_rgb = hdrGamutConversionFn(hdr_rgb);
846 float hdr_y_nits = luminanceFn(hdr_rgb) * hdr_white_nits;
847
848 size_t pixel_idx = x + y * dest->width;
849 reinterpret_cast<uint8_t*>(dest->data)[pixel_idx] =
850 encodeRecovery(sdr_y_nits, hdr_y_nits, metadata->rangeScalingFactor);
851 }
852 }
853 }
854 };
855
856 std::vector<std::thread> workers;
857 for (int th = 0; th < threads - 1; th++) {
858 workers.push_back(std::thread(computeMetadata));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400859 }
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800860
861 // compute metadata
862 for (size_t rowStart = 0; rowStart < image_height;) {
863 size_t rowEnd = std::min(rowStart + rowStep, image_height);
864 jobQueue.enqueueJob(rowStart, rowEnd);
865 rowStart = rowEnd;
866 }
867 jobQueue.markQueueForEnd();
868 computeMetadata();
869 std::for_each(workers.begin(), workers.end(), [](std::thread& t) { t.join(); });
870 workers.clear();
Nick Deakin6bd90432022-11-20 16:26:37 -0500871 hdr_y_nits_avg /= image_width * image_height;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400872
Nick Deakin6bd90432022-11-20 16:26:37 -0500873 metadata->rangeScalingFactor = hdr_y_nits_max / kSdrWhiteNits;
874 if (metadata->transferFunction == JPEGR_TF_PQ) {
875 metadata->hdr10Metadata.maxFALL = hdr_y_nits_avg;
876 metadata->hdr10Metadata.maxCLL = hdr_y_nits_max;
877 }
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400878
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800879 // generate map
880 jobQueue.reset();
881 for (int th = 0; th < threads - 1; th++) {
882 workers.push_back(std::thread(generateMap));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400883 }
884
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800885 rowStep = (threads == 1 ? image_height : kJobSzInRows) / kMapDimensionScaleFactor;
886 for (size_t rowStart = 0; rowStart < map_height;) {
887 size_t rowEnd = std::min(rowStart + rowStep, map_height);
888 jobQueue.enqueueJob(rowStart, rowEnd);
889 rowStart = rowEnd;
890 }
891 jobQueue.markQueueForEnd();
892 generateMap();
893 std::for_each(workers.begin(), workers.end(), [](std::thread& t) { t.join(); });
894
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400895 map_data.release();
Dichen Zhang6947d532022-10-22 02:16:21 +0000896 return NO_ERROR;
Dichen Zhang596a7562022-10-12 14:57:05 -0700897}
898
Dichen Zhang6947d532022-10-22 02:16:21 +0000899status_t RecoveryMap::applyRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_image,
900 jr_uncompressed_ptr uncompressed_recovery_map,
Nick Deakin6bd90432022-11-20 16:26:37 -0500901 jr_metadata_ptr metadata,
Dichen Zhang6947d532022-10-22 02:16:21 +0000902 jr_uncompressed_ptr dest) {
Dichen Zhang596a7562022-10-12 14:57:05 -0700903 if (uncompressed_yuv_420_image == nullptr
904 || uncompressed_recovery_map == nullptr
Nick Deakin6bd90432022-11-20 16:26:37 -0500905 || metadata == nullptr
Dichen Zhang596a7562022-10-12 14:57:05 -0700906 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +0000907 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang596a7562022-10-12 14:57:05 -0700908 }
909
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800910 dest->width = uncompressed_yuv_420_image->width;
911 dest->height = uncompressed_yuv_420_image->height;
Ram Mohanfe723d62022-12-15 00:59:11 +0530912 ShepardsIDW idwTable(kMapDimensionScaleFactor);
Harish Mahendrakarf25991f2022-12-16 11:57:44 -0800913 RecoveryLUT recoveryLUT(metadata->rangeScalingFactor);
Ram Mohanfe723d62022-12-15 00:59:11 +0530914
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800915 JobQueue jobQueue;
916 std::function<void()> applyRecMap = [uncompressed_yuv_420_image, uncompressed_recovery_map,
Harish Mahendrakarf25991f2022-12-16 11:57:44 -0800917 metadata, dest, &jobQueue, &idwTable,
918 &recoveryLUT]() -> void {
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800919 const float hdr_ratio = metadata->rangeScalingFactor;
920 size_t width = uncompressed_yuv_420_image->width;
921 size_t height = uncompressed_yuv_420_image->height;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400922
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800923 ColorTransformFn hdrOetf = nullptr;
924 switch (metadata->transferFunction) {
925 case JPEGR_TF_LINEAR:
926 hdrOetf = identityConversion;
927 break;
928 case JPEGR_TF_HLG:
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800929#if USE_HLG_OETF_LUT
930 hdrOetf = hlgOetfLUT;
931#else
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800932 hdrOetf = hlgOetf;
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800933#endif
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800934 break;
935 case JPEGR_TF_PQ:
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800936#if USE_PQ_OETF_LUT
937 hdrOetf = pqOetfLUT;
938#else
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800939 hdrOetf = pqOetf;
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800940#endif
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800941 break;
942 case JPEGR_TF_UNSPECIFIED:
943 // Should be impossible to hit after input validation.
944 hdrOetf = identityConversion;
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400945 }
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800946
947 size_t rowStart, rowEnd;
948 while (jobQueue.dequeueJob(rowStart, rowEnd)) {
949 for (size_t y = rowStart; y < rowEnd; ++y) {
950 for (size_t x = 0; x < width; ++x) {
951 Color yuv_gamma_sdr = getYuv420Pixel(uncompressed_yuv_420_image, x, y);
952 Color rgb_gamma_sdr = srgbYuvToRgb(yuv_gamma_sdr);
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800953#if USE_SRGB_INVOETF_LUT
954 Color rgb_sdr = srgbInvOetfLUT(rgb_gamma_sdr);
955#else
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800956 Color rgb_sdr = srgbInvOetf(rgb_gamma_sdr);
Harish Mahendrakar555a06b2022-12-14 09:37:27 -0800957#endif
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800958 float recovery;
959 // TODO: determine map scaling factor based on actual map dims
960 size_t map_scale_factor = kMapDimensionScaleFactor;
961 // TODO: If map_scale_factor is guaranteed to be an integer, then remove the following.
962 // Currently map_scale_factor is of type size_t, but it could be changed to a float
963 // later.
964 if (map_scale_factor != floorf(map_scale_factor)) {
965 recovery = sampleMap(uncompressed_recovery_map, map_scale_factor, x, y);
966 } else {
967 recovery = sampleMap(uncompressed_recovery_map, map_scale_factor, x, y,
968 idwTable);
969 }
Harish Mahendrakarf25991f2022-12-16 11:57:44 -0800970#if USE_APPLY_RECOVERY_LUT
971 Color rgb_hdr = applyRecoveryLUT(rgb_sdr, recovery, recoveryLUT);
972#else
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800973 Color rgb_hdr = applyRecovery(rgb_sdr, recovery, hdr_ratio);
Harish Mahendrakarf25991f2022-12-16 11:57:44 -0800974#endif
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800975 Color rgb_gamma_hdr = hdrOetf(rgb_hdr / metadata->rangeScalingFactor);
976 uint32_t rgba1010102 = colorToRgba1010102(rgb_gamma_hdr);
977
978 size_t pixel_idx = x + y * width;
979 reinterpret_cast<uint32_t*>(dest->data)[pixel_idx] = rgba1010102;
980 }
981 }
982 }
983 };
984
985 const int threads = std::clamp(GetCPUCoreCount(), 1, 4);
986 std::vector<std::thread> workers;
987 for (int th = 0; th < threads - 1; th++) {
988 workers.push_back(std::thread(applyRecMap));
Nick Deakinf6bca5a2022-11-04 10:43:43 -0400989 }
Harish Mahendrakar72b6f302022-12-16 10:39:15 -0800990 const int rowStep = threads == 1 ? uncompressed_yuv_420_image->height : kJobSzInRows;
991 for (int rowStart = 0; rowStart < uncompressed_yuv_420_image->height;) {
992 int rowEnd = std::min(rowStart + rowStep, uncompressed_yuv_420_image->height);
993 jobQueue.enqueueJob(rowStart, rowEnd);
994 rowStart = rowEnd;
995 }
996 jobQueue.markQueueForEnd();
997 applyRecMap();
998 std::for_each(workers.begin(), workers.end(), [](std::thread& t) { t.join(); });
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +0000999 return NO_ERROR;
1000}
1001
1002status_t RecoveryMap::extractPrimaryImageAndRecoveryMap(jr_compressed_ptr compressed_jpegr_image,
1003 jr_compressed_ptr primary_image,
1004 jr_compressed_ptr recovery_map) {
1005 if (compressed_jpegr_image == nullptr) {
1006 return ERROR_JPEGR_INVALID_NULL_PTR;
1007 }
1008
1009 MessageHandler msg_handler;
1010 std::shared_ptr<DataSegment> seg =
1011 DataSegment::Create(DataRange(0, compressed_jpegr_image->length),
1012 static_cast<const uint8_t*>(compressed_jpegr_image->data),
1013 DataSegment::BufferDispositionPolicy::kDontDelete);
1014 DataSegmentDataSource data_source(seg);
1015 JpegInfoBuilder jpeg_info_builder;
1016 jpeg_info_builder.SetImageLimit(2);
1017 JpegScanner jpeg_scanner(&msg_handler);
1018 jpeg_scanner.Run(&data_source, &jpeg_info_builder);
1019 data_source.Reset();
1020
1021 if (jpeg_scanner.HasError()) {
1022 return ERROR_JPEGR_INVALID_INPUT_TYPE;
1023 }
1024
1025 const auto& jpeg_info = jpeg_info_builder.GetInfo();
1026 const auto& image_ranges = jpeg_info.GetImageRanges();
1027 if (image_ranges.empty()) {
1028 return ERROR_JPEGR_INVALID_INPUT_TYPE;
1029 }
1030
1031 if (image_ranges.size() != 2) {
1032 // Must be 2 JPEG Images
1033 return ERROR_JPEGR_INVALID_INPUT_TYPE;
1034 }
1035
1036 if (primary_image != nullptr) {
1037 primary_image->data = static_cast<uint8_t*>(compressed_jpegr_image->data) +
1038 image_ranges[0].GetBegin();
1039 primary_image->length = image_ranges[0].GetLength();
1040 }
1041
1042 if (recovery_map != nullptr) {
1043 recovery_map->data = static_cast<uint8_t*>(compressed_jpegr_image->data) +
1044 image_ranges[1].GetBegin();
1045 recovery_map->length = image_ranges[1].GetLength();
1046 }
Nick Deakinf6bca5a2022-11-04 10:43:43 -04001047
Dichen Zhang6947d532022-10-22 02:16:21 +00001048 return NO_ERROR;
Dichen Zhang596a7562022-10-12 14:57:05 -07001049}
1050
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +00001051
Nick Deakinf6bca5a2022-11-04 10:43:43 -04001052status_t RecoveryMap::extractRecoveryMap(jr_compressed_ptr compressed_jpegr_image,
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +00001053 jr_compressed_ptr dest) {
1054 if (compressed_jpegr_image == nullptr || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +00001055 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang85b37562022-10-11 11:08:28 -07001056 }
1057
Fyodor Kyslov1dcc4422022-11-16 01:40:53 +00001058 return extractPrimaryImageAndRecoveryMap(compressed_jpegr_image, nullptr, dest);
Dichen Zhang85b37562022-10-11 11:08:28 -07001059}
1060
Dichen Zhangd18bc302022-12-16 20:55:24 +00001061// JPEG/R structure:
1062// SOI (ff d8)
1063// APP1 (ff e1)
1064// 2 bytes of length (2 + length of exif package)
1065// EXIF package (this includes the first two bytes representing the package length)
1066// APP1 (ff e1)
1067// 2 bytes of length (2 + 29 + length of xmp package)
1068// name space ("http://ns.adobe.com/xap/1.0/\0")
1069// xmp
1070// primary image (without the first two bytes (SOI) and without EXIF, may have other packages)
1071// secondary image (the recovery map)
1072//
1073// Metadata versions we are using:
1074// ECMA TR-98 for JFIF marker
1075// Exif 2.2 spec for EXIF marker
1076// Adobe XMP spec part 3 for XMP marker
1077// ICC v4.3 spec for ICC
Nick Deakinf6bca5a2022-11-04 10:43:43 -04001078status_t RecoveryMap::appendRecoveryMap(jr_compressed_ptr compressed_jpeg_image,
1079 jr_compressed_ptr compressed_recovery_map,
Dichen Zhangd18bc302022-12-16 20:55:24 +00001080 jr_exif_ptr exif,
Nick Deakin6bd90432022-11-20 16:26:37 -05001081 jr_metadata_ptr metadata,
Nick Deakinf6bca5a2022-11-04 10:43:43 -04001082 jr_compressed_ptr dest) {
Dichen Zhang6947d532022-10-22 02:16:21 +00001083 if (compressed_jpeg_image == nullptr
1084 || compressed_recovery_map == nullptr
Dichen Zhangd18bc302022-12-16 20:55:24 +00001085 || exif == nullptr
Nick Deakin6bd90432022-11-20 16:26:37 -05001086 || metadata == nullptr
Dichen Zhang6947d532022-10-22 02:16:21 +00001087 || dest == nullptr) {
Dichen Zhang80b72482022-11-02 01:55:35 +00001088 return ERROR_JPEGR_INVALID_NULL_PTR;
Dichen Zhang85b37562022-10-11 11:08:28 -07001089 }
1090
Dichen Zhanga8766262022-11-07 23:48:24 +00001091 int pos = 0;
1092
Dichen Zhangd18bc302022-12-16 20:55:24 +00001093 // Write SOI
Dichen Zhanga8766262022-11-07 23:48:24 +00001094 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kStart, 1, pos));
1095 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kSOI, 1, pos));
Dichen Zhangd18bc302022-12-16 20:55:24 +00001096
1097 // Write EXIF
1098 {
1099 const int length = 2 + exif->length;
1100 const uint8_t lengthH = ((length >> 8) & 0xff);
1101 const uint8_t lengthL = (length & 0xff);
1102 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kStart, 1, pos));
1103 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kAPP1, 1, pos));
1104 JPEGR_CHECK(Write(dest, &lengthH, 1, pos));
1105 JPEGR_CHECK(Write(dest, &lengthL, 1, pos));
1106 JPEGR_CHECK(Write(dest, exif->data, exif->length, pos));
1107 }
1108
1109 // Prepare and write XMP
1110 {
1111 const string xmp = generateXmp(compressed_recovery_map->length, *metadata);
1112 const string nameSpace = "http://ns.adobe.com/xap/1.0/\0";
1113 const int nameSpaceLength = nameSpace.size() + 1; // need to count the null terminator
1114 // 2 bytes: representing the length of the package
1115 // 29 bytes: length of name space "http://ns.adobe.com/xap/1.0/\0",
1116 // x bytes: length of xmp packet
Dichen Zhang25df9c82023-01-03 17:04:10 -08001117 const int length = 2 + nameSpaceLength + xmp.size();
Dichen Zhangd18bc302022-12-16 20:55:24 +00001118 const uint8_t lengthH = ((length >> 8) & 0xff);
1119 const uint8_t lengthL = (length & 0xff);
1120 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kStart, 1, pos));
1121 JPEGR_CHECK(Write(dest, &photos_editing_formats::image_io::JpegMarker::kAPP1, 1, pos));
1122 JPEGR_CHECK(Write(dest, &lengthH, 1, pos));
1123 JPEGR_CHECK(Write(dest, &lengthL, 1, pos));
1124 JPEGR_CHECK(Write(dest, (void*)nameSpace.c_str(), nameSpaceLength, pos));
1125 JPEGR_CHECK(Write(dest, (void*)xmp.c_str(), xmp.size(), pos));
1126 }
1127
1128 // Write primary image
Dichen Zhanga8766262022-11-07 23:48:24 +00001129 JPEGR_CHECK(Write(dest,
1130 (uint8_t*)compressed_jpeg_image->data + 2, compressed_jpeg_image->length - 2, pos));
Dichen Zhangd18bc302022-12-16 20:55:24 +00001131
1132 // Write secondary image
Dichen Zhanga8766262022-11-07 23:48:24 +00001133 JPEGR_CHECK(Write(dest, compressed_recovery_map->data, compressed_recovery_map->length, pos));
Dichen Zhangd18bc302022-12-16 20:55:24 +00001134
1135 // Set back length
Dichen Zhanga8766262022-11-07 23:48:24 +00001136 dest->length = pos;
1137
Dichen Zhangd18bc302022-12-16 20:55:24 +00001138 // Done!
Dichen Zhang6947d532022-10-22 02:16:21 +00001139 return NO_ERROR;
Dichen Zhang85b37562022-10-11 11:08:28 -07001140}
1141
Dichen Zhangc3437ca2023-01-04 14:00:08 -08001142status_t RecoveryMap::toneMap(jr_uncompressed_ptr src,
Dichen Zhang636f5242022-12-07 20:25:44 +00001143 jr_uncompressed_ptr dest) {
Dichen Zhangc3437ca2023-01-04 14:00:08 -08001144 if (src == nullptr || dest == nullptr) {
Dichen Zhang636f5242022-12-07 20:25:44 +00001145 return ERROR_JPEGR_INVALID_NULL_PTR;
1146 }
1147
Dichen Zhangc3437ca2023-01-04 14:00:08 -08001148 dest->width = src->width;
1149 dest->height = src->height;
Dichen Zhang636f5242022-12-07 20:25:44 +00001150
Dichen Zhangc3437ca2023-01-04 14:00:08 -08001151 size_t pixel_count = src->width * src->height;
1152 for (size_t y = 0; y < src->height; ++y) {
1153 for (size_t x = 0; x < src->width; ++x) {
1154 size_t pixel_y_idx = x + y * src->width;
1155 size_t pixel_uv_idx = x / 2 + (y / 2) * (src->width / 2);
1156
1157 uint16_t y_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_y_idx]
1158 >> 6;
1159 uint16_t u_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_count + pixel_uv_idx * 2]
1160 >> 6;
1161 uint16_t v_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_count + pixel_uv_idx * 2 + 1]
1162 >> 6;
1163
1164 uint8_t* y = &reinterpret_cast<uint8_t*>(dest->data)[pixel_y_idx];
1165 uint8_t* u = &reinterpret_cast<uint8_t*>(dest->data)[pixel_count + pixel_uv_idx];
1166 uint8_t* v = &reinterpret_cast<uint8_t*>(dest->data)[pixel_count * 5 / 4 + pixel_uv_idx];
1167
1168 *y = static_cast<uint8_t>((y_uint >> 2) & 0xff);
1169 *u = static_cast<uint8_t>((u_uint >> 2) & 0xff);
1170 *v = static_cast<uint8_t>((v_uint >> 2) & 0xff);
1171 }
1172 }
1173
1174 dest->colorGamut = src->colorGamut;
Dichen Zhang636f5242022-12-07 20:25:44 +00001175
1176 return NO_ERROR;
1177}
1178
Dichen Zhang85b37562022-10-11 11:08:28 -07001179} // namespace android::recoverymap