lib/jxl/jpeg/enc_jpeg_data.cc

   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/jxl/jpeg/enc_jpeg_data.h"
   7
   8 #include <brotli/encode.h>
   9
  10 #include "lib/jxl/codec_in_out.h"
  11 #include "lib/jxl/enc_fields.h"
  12 #include "lib/jxl/image_bundle.h"
  13 #include "lib/jxl/jpeg/enc_jpeg_data_reader.h"
  14 #include "lib/jxl/luminance.h"
  15 #include "lib/jxl/sanitizers.h"
  16
  17 namespace jxl {
  18 namespace jpeg {
  19
  20 namespace {
  21
  22 constexpr int BITS_IN_JSAMPLE = 8;
  23 using ByteSpan = Span<const uint8_t>;
  24
  25 // TODO(eustas): move to jpeg_data, to use from codec_jpg as well.
  26 // See if there is a canonically chunked ICC profile and mark corresponding
  27 // app-tags with AppMarkerType::kICC.
  28 Status DetectIccProfile(JPEGData& jpeg_data) {
  29   JXL_DASSERT(jpeg_data.app_data.size() == jpeg_data.app_marker_type.size());
  30   size_t num_icc = 0;
  31   size_t num_icc_jpeg = 0;
  32   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
  33     const auto& app = jpeg_data.app_data[i];
  34     size_t pos = 0;
  35     if (app[pos++] != 0xE2) continue;
  36     // At least APPn + size; otherwise it should be intermarker-data.
  37     JXL_DASSERT(app.size() >= 3);
  38     size_t tag_length = (app[pos] << 8) + app[pos + 1];
  39     pos += 2;
  40     JXL_DASSERT(app.size() == tag_length + 1);
  41     // Empty payload is 2 bytes for tag length itself + signature
  42     if (tag_length < 2 + sizeof kIccProfileTag) continue;
  43
  44     if (memcmp(&app[pos], kIccProfileTag, sizeof kIccProfileTag) != 0) continue;
  45     pos += sizeof kIccProfileTag;
  46     uint8_t chunk_id = app[pos++];
  47     uint8_t num_chunks = app[pos++];
  48     if (chunk_id != num_icc + 1) continue;
  49     if (num_icc_jpeg == 0) num_icc_jpeg = num_chunks;
  50     if (num_icc_jpeg != num_chunks) continue;
  51     num_icc++;
  52     jpeg_data.app_marker_type[i] = AppMarkerType::kICC;
  53   }
  54   if (num_icc != num_icc_jpeg) {
  55     return JXL_FAILURE("Invalid ICC chunks");
  56   }
  57   return true;
  58 }
  59
  60 bool GetMarkerPayload(const uint8_t* data, size_t size, ByteSpan* payload) {
  61   if (size < 3) {
  62     return false;
  63   }
  64   size_t hi = data[1];
  65   size_t lo = data[2];
  66   size_t internal_size = (hi << 8u) | lo;
  67   // Second byte of marker is not counted towards size.
  68   if (internal_size != size - 1) {
  69     return false;
  70   }
  71   // cut second marker byte and "length" from payload.
  72   *payload = ByteSpan(data, size);
  73   payload->remove_prefix(3);
  74   return true;
  75 }
  76
  77 Status DetectBlobs(jpeg::JPEGData& jpeg_data) {
  78   JXL_DASSERT(jpeg_data.app_data.size() == jpeg_data.app_marker_type.size());
  79   bool have_exif = false, have_xmp = false;
  80   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
  81     auto& marker = jpeg_data.app_data[i];
  82     if (marker.empty() || marker[0] != kApp1) {
  83       continue;
  84     }
  85     ByteSpan payload;
  86     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
  87       // Something is wrong with this marker; does not care.
  88       continue;
  89     }
  90     if (!have_exif && payload.size() >= sizeof kExifTag &&
  91         !memcmp(payload.data(), kExifTag, sizeof kExifTag)) {
  92       jpeg_data.app_marker_type[i] = AppMarkerType::kExif;
  93       have_exif = true;
  94     }
  95     if (!have_xmp && payload.size() >= sizeof kXMPTag &&
  96         !memcmp(payload.data(), kXMPTag, sizeof kXMPTag)) {
  97       jpeg_data.app_marker_type[i] = AppMarkerType::kXMP;
  98       have_xmp = true;
  99     }
 100   }
 101   return true;
 102 }
 103
 104 Status ParseChunkedMarker(const jpeg::JPEGData& src, uint8_t marker_type,
 105                           const ByteSpan& tag, IccBytes* output,
 106                           bool allow_permutations = false) {
 107   output->clear();
 108
 109   std::vector<ByteSpan> chunks;
 110   std::vector<bool> presence;
 111   size_t expected_number_of_parts = 0;
 112   bool is_first_chunk = true;
 113   size_t ordinal = 0;
 114   for (const auto& marker : src.app_data) {
 115     if (marker.empty() || marker[0] != marker_type) {
 116       continue;
 117     }
 118     ByteSpan payload;
 119     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
 120       // Something is wrong with this marker; does not care.
 121       continue;
 122     }
 123     if ((payload.size() < tag.size()) ||
 124         memcmp(payload.data(), tag.data(), tag.size()) != 0) {
 125       continue;
 126     }
 127     payload.remove_prefix(tag.size());
 128     if (payload.size() < 2) {
 129       return JXL_FAILURE("Chunk is too small.");
 130     }
 131     uint8_t index = payload[0];
 132     uint8_t total = payload[1];
 133     ordinal++;
 134     if (!allow_permutations) {
 135       if (index != ordinal) return JXL_FAILURE("Invalid chunk order.");
 136     }
 137
 138     payload.remove_prefix(2);
 139
 140     JXL_RETURN_IF_ERROR(total != 0);
 141     if (is_first_chunk) {
 142       is_first_chunk = false;
 143       expected_number_of_parts = total;
 144       // 1-based indices; 0-th element is added for convenience.
 145       chunks.resize(total + 1);
 146       presence.resize(total + 1);
 147     } else {
 148       JXL_RETURN_IF_ERROR(expected_number_of_parts == total);
 149     }
 150
 151     if (index == 0 || index > total) {
 152       return JXL_FAILURE("Invalid chunk index.");
 153     }
 154
 155     if (presence[index]) {
 156       return JXL_FAILURE("Duplicate chunk.");
 157     }
 158     presence[index] = true;
 159     chunks[index] = payload;
 160   }
 161
 162   for (size_t i = 0; i < expected_number_of_parts; ++i) {
 163     // 0-th element is not used.
 164     size_t index = i + 1;
 165     if (!presence[index]) {
 166       return JXL_FAILURE("Missing chunk.");
 167     }
 168     chunks[index].AppendTo(output);
 169   }
 170
 171   return true;
 172 }
 173
 174 Status SetBlobsFromJpegData(const jpeg::JPEGData& jpeg_data, Blobs* blobs) {
 175   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
 176     auto& marker = jpeg_data.app_data[i];
 177     if (marker.empty() || marker[0] != kApp1) {
 178       continue;
 179     }
 180     ByteSpan payload;
 181     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
 182       // Something is wrong with this marker; does not care.
 183       continue;
 184     }
 185     if (payload.size() >= sizeof kExifTag &&
 186         !memcmp(payload.data(), kExifTag, sizeof kExifTag)) {
 187       if (blobs->exif.empty()) {
 188         blobs->exif.resize(payload.size() - sizeof kExifTag);
 189         memcpy(blobs->exif.data(), payload.data() + sizeof kExifTag,
 190                payload.size() - sizeof kExifTag);
 191       } else {
 192         JXL_WARNING(
 193             "ReJPEG: multiple Exif blobs, storing only first one in the JPEG "
 194             "XL container\n");
 195       }
 196     }
 197     if (payload.size() >= sizeof kXMPTag &&
 198         !memcmp(payload.data(), kXMPTag, sizeof kXMPTag)) {
 199       if (blobs->xmp.empty()) {
 200         blobs->xmp.resize(payload.size() - sizeof kXMPTag);
 201         memcpy(blobs->xmp.data(), payload.data() + sizeof kXMPTag,
 202                payload.size() - sizeof kXMPTag);
 203       } else {
 204         JXL_WARNING(
 205             "ReJPEG: multiple XMP blobs, storing only first one in the JPEG "
 206             "XL container\n");
 207       }
 208     }
 209   }
 210   return true;
 211 }
 212
 213 static inline bool IsJPG(const Span<const uint8_t> bytes) {
 214   return bytes.size() >= 2 && bytes[0] == 0xFF && bytes[1] == 0xD8;
 215 }
 216
 217 }  // namespace
 218
 219 void SetColorEncodingFromJpegData(const jpeg::JPEGData& jpg,
 220                                   ColorEncoding* color_encoding) {
 221   IccBytes icc_profile;
 222   if (!ParseChunkedMarker(jpg, kApp2, ByteSpan(kIccProfileTag), &icc_profile)) {
 223     JXL_WARNING("ReJPEG: corrupted ICC profile\n");
 224     icc_profile.clear();
 225   }
 226
 227   if (icc_profile.empty()) {
 228     bool is_gray = (jpg.components.size() == 1);
 229     *color_encoding = ColorEncoding::SRGB(is_gray);
 230   } else {
 231     color_encoding->SetICCRaw(std::move(icc_profile));
 232   }
 233 }
 234
 235 Status EncodeJPEGData(JPEGData& jpeg_data, std::vector<uint8_t>* bytes,
 236                       const CompressParams& cparams) {
 237   bytes->clear();
 238   jpeg_data.app_marker_type.resize(jpeg_data.app_data.size(),
 239                                    AppMarkerType::kUnknown);
 240   JXL_RETURN_IF_ERROR(DetectIccProfile(jpeg_data));
 241   JXL_RETURN_IF_ERROR(DetectBlobs(jpeg_data));
 242
 243   size_t total_data = 0;
 244   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
 245     if (jpeg_data.app_marker_type[i] != AppMarkerType::kUnknown) {
 246       continue;
 247     }
 248     total_data += jpeg_data.app_data[i].size();
 249   }
 250   for (size_t i = 0; i < jpeg_data.com_data.size(); i++) {
 251     total_data += jpeg_data.com_data[i].size();
 252   }
 253   for (size_t i = 0; i < jpeg_data.inter_marker_data.size(); i++) {
 254     total_data += jpeg_data.inter_marker_data[i].size();
 255   }
 256   total_data += jpeg_data.tail_data.size();
 257   size_t brotli_capacity = BrotliEncoderMaxCompressedSize(total_data);
 258
 259   BitWriter writer;
 260   JXL_RETURN_IF_ERROR(Bundle::Write(jpeg_data, &writer, 0, nullptr));
 261   writer.ZeroPadToByte();
 262   {
 263     PaddedBytes serialized_jpeg_data = std::move(writer).TakeBytes();
 264     bytes->reserve(serialized_jpeg_data.size() + brotli_capacity);
 265     Bytes(serialized_jpeg_data).AppendTo(bytes);
 266   }
 267
 268   BrotliEncoderState* brotli_enc =
 269       BrotliEncoderCreateInstance(nullptr, nullptr, nullptr);
 270   int effort = cparams.brotli_effort;
 271   if (effort < 0) effort = 11 - static_cast<int>(cparams.speed_tier);
 272   BrotliEncoderSetParameter(brotli_enc, BROTLI_PARAM_QUALITY, effort);
 273   size_t initial_size = bytes->size();
 274   BrotliEncoderSetParameter(brotli_enc, BROTLI_PARAM_SIZE_HINT, total_data);
 275   bytes->resize(initial_size + brotli_capacity);
 276   size_t enc_size = 0;
 277   auto br_append = [&](const std::vector<uint8_t>& data, bool last) {
 278     size_t available_in = data.size();
 279     const uint8_t* in = data.data();
 280     uint8_t* out = &(*bytes)[initial_size + enc_size];
 281     do {
 282       uint8_t* out_before = out;
 283       msan::MemoryIsInitialized(in, available_in);
 284       JXL_CHECK(BrotliEncoderCompressStream(
 285           brotli_enc, last ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS,
 286           &available_in, &in, &brotli_capacity, &out, &enc_size));
 287       msan::UnpoisonMemory(out_before, out - out_before);
 288     } while (BrotliEncoderHasMoreOutput(brotli_enc) || available_in > 0);
 289   };
 290
 291   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
 292     if (jpeg_data.app_marker_type[i] != AppMarkerType::kUnknown) {
 293       continue;
 294     }
 295     br_append(jpeg_data.app_data[i], /*last=*/false);
 296   }
 297   for (size_t i = 0; i < jpeg_data.com_data.size(); i++) {
 298     br_append(jpeg_data.com_data[i], /*last=*/false);
 299   }
 300   for (size_t i = 0; i < jpeg_data.inter_marker_data.size(); i++) {
 301     br_append(jpeg_data.inter_marker_data[i], /*last=*/false);
 302   }
 303   br_append(jpeg_data.tail_data, /*last=*/true);
 304   BrotliEncoderDestroyInstance(brotli_enc);
 305   bytes->resize(initial_size + enc_size);
 306   return true;
 307 }
 308
 309 Status DecodeImageJPG(const Span<const uint8_t> bytes, CodecInOut* io) {
 310   if (!IsJPG(bytes)) return false;
 311   io->frames.clear();
 312   io->frames.reserve(1);
 313   io->frames.emplace_back(&io->metadata.m);
 314   io->Main().jpeg_data = make_unique<jpeg::JPEGData>();
 315   jpeg::JPEGData* jpeg_data = io->Main().jpeg_data.get();
 316   if (!jpeg::ReadJpeg(bytes.data(), bytes.size(), jpeg::JpegReadMode::kReadAll,
 317                       jpeg_data)) {
 318     return JXL_FAILURE("Error reading JPEG");
 319   }
 320   SetColorEncodingFromJpegData(*jpeg_data, &io->metadata.m.color_encoding);
 321   JXL_RETURN_IF_ERROR(SetBlobsFromJpegData(*jpeg_data, &io->blobs));
 322   size_t nbcomp = jpeg_data->components.size();
 323   if (nbcomp != 1 && nbcomp != 3) {
 324     return JXL_FAILURE("Cannot recompress JPEGs with neither 1 nor 3 channels");
 325   }
 326   YCbCrChromaSubsampling cs;
 327   if (nbcomp == 3) {
 328     uint8_t hsample[3], vsample[3];
 329     for (size_t i = 0; i < nbcomp; i++) {
 330       hsample[i] = jpeg_data->components[i].h_samp_factor;
 331       vsample[i] = jpeg_data->components[i].v_samp_factor;
 332     }
 333     JXL_RETURN_IF_ERROR(cs.Set(hsample, vsample));
 334   } else if (nbcomp == 1) {
 335     uint8_t hsample[3], vsample[3];
 336     for (size_t i = 0; i < 3; i++) {
 337       hsample[i] = jpeg_data->components[0].h_samp_factor;
 338       vsample[i] = jpeg_data->components[0].v_samp_factor;
 339     }
 340     JXL_RETURN_IF_ERROR(cs.Set(hsample, vsample));
 341   }
 342   bool is_rgb = false;
 343   {
 344     const auto& markers = jpeg_data->marker_order;
 345     // If there is a JFIF marker, this is YCbCr. Otherwise...
 346     if (std::find(markers.begin(), markers.end(), 0xE0) == markers.end()) {
 347       // Try to find an 'Adobe' marker.
 348       size_t app_markers = 0;
 349       size_t i = 0;
 350       for (; i < markers.size(); i++) {
 351         // This is an APP marker.
 352         if ((markers[i] & 0xF0) == 0xE0) {
 353           JXL_CHECK(app_markers < jpeg_data->app_data.size());
 354           // APP14 marker
 355           if (markers[i] == 0xEE) {
 356             const auto& data = jpeg_data->app_data[app_markers];
 357             if (data.size() == 15 && data[3] == 'A' && data[4] == 'd' &&
 358                 data[5] == 'o' && data[6] == 'b' && data[7] == 'e') {
 359               // 'Adobe' marker.
 360               is_rgb = data[14] == 0;
 361               break;
 362             }
 363           }
 364           app_markers++;
 365         }
 366       }
 367
 368       if (i == markers.size()) {
 369         // No 'Adobe' marker, guess from component IDs.
 370         is_rgb = nbcomp == 3 && jpeg_data->components[0].id == 'R' &&
 371                  jpeg_data->components[1].id == 'G' &&
 372                  jpeg_data->components[2].id == 'B';
 373       }
 374     }
 375   }
 376
 377   io->Main().chroma_subsampling = cs;
 378   io->Main().color_transform =
 379       (!is_rgb || nbcomp == 1) ? ColorTransform::kYCbCr : ColorTransform::kNone;
 380
 381   io->metadata.m.SetIntensityTarget(kDefaultIntensityTarget);
 382   io->metadata.m.SetUintSamples(BITS_IN_JSAMPLE);
 383   io->SetFromImage(Image3F(jpeg_data->width, jpeg_data->height),
 384                    io->metadata.m.color_encoding);
 385   SetIntensityTarget(&io->metadata.m);
 386   return true;
 387 }
 388
 389 }  // namespace jpeg
 390 }  // namespace jxl