libjxl

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

enc_jpeg_data.cc (14004B)


      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_bit_writer.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;
     80   bool have_xmp = false;
     81   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
     82     auto& marker = jpeg_data.app_data[i];
     83     if (marker.empty() || marker[0] != kApp1) {
     84       continue;
     85     }
     86     ByteSpan payload;
     87     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
     88       // Something is wrong with this marker; does not care.
     89       continue;
     90     }
     91     if (!have_exif && payload.size() >= sizeof kExifTag &&
     92         !memcmp(payload.data(), kExifTag, sizeof kExifTag)) {
     93       jpeg_data.app_marker_type[i] = AppMarkerType::kExif;
     94       have_exif = true;
     95     }
     96     if (!have_xmp && payload.size() >= sizeof kXMPTag &&
     97         !memcmp(payload.data(), kXMPTag, sizeof kXMPTag)) {
     98       jpeg_data.app_marker_type[i] = AppMarkerType::kXMP;
     99       have_xmp = true;
    100     }
    101   }
    102   return true;
    103 }
    104 
    105 Status ParseChunkedMarker(const jpeg::JPEGData& src, uint8_t marker_type,
    106                           const ByteSpan& tag, IccBytes* output,
    107                           bool allow_permutations = false) {
    108   output->clear();
    109 
    110   std::vector<ByteSpan> chunks;
    111   std::vector<bool> presence;
    112   size_t expected_number_of_parts = 0;
    113   bool is_first_chunk = true;
    114   size_t ordinal = 0;
    115   for (const auto& marker : src.app_data) {
    116     if (marker.empty() || marker[0] != marker_type) {
    117       continue;
    118     }
    119     ByteSpan payload;
    120     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
    121       // Something is wrong with this marker; does not care.
    122       continue;
    123     }
    124     if ((payload.size() < tag.size()) ||
    125         memcmp(payload.data(), tag.data(), tag.size()) != 0) {
    126       continue;
    127     }
    128     payload.remove_prefix(tag.size());
    129     if (payload.size() < 2) {
    130       return JXL_FAILURE("Chunk is too small.");
    131     }
    132     uint8_t index = payload[0];
    133     uint8_t total = payload[1];
    134     ordinal++;
    135     if (!allow_permutations) {
    136       if (index != ordinal) return JXL_FAILURE("Invalid chunk order.");
    137     }
    138 
    139     payload.remove_prefix(2);
    140 
    141     JXL_RETURN_IF_ERROR(total != 0);
    142     if (is_first_chunk) {
    143       is_first_chunk = false;
    144       expected_number_of_parts = total;
    145       // 1-based indices; 0-th element is added for convenience.
    146       chunks.resize(total + 1);
    147       presence.resize(total + 1);
    148     } else {
    149       JXL_RETURN_IF_ERROR(expected_number_of_parts == total);
    150     }
    151 
    152     if (index == 0 || index > total) {
    153       return JXL_FAILURE("Invalid chunk index.");
    154     }
    155 
    156     if (presence[index]) {
    157       return JXL_FAILURE("Duplicate chunk.");
    158     }
    159     presence[index] = true;
    160     chunks[index] = payload;
    161   }
    162 
    163   for (size_t i = 0; i < expected_number_of_parts; ++i) {
    164     // 0-th element is not used.
    165     size_t index = i + 1;
    166     if (!presence[index]) {
    167       return JXL_FAILURE("Missing chunk.");
    168     }
    169     chunks[index].AppendTo(*output);
    170   }
    171 
    172   return true;
    173 }
    174 
    175 Status SetBlobsFromJpegData(const jpeg::JPEGData& jpeg_data, Blobs* blobs) {
    176   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
    177     const auto& marker = jpeg_data.app_data[i];
    178     if (marker.empty() || marker[0] != kApp1) {
    179       continue;
    180     }
    181     ByteSpan payload;
    182     if (!GetMarkerPayload(marker.data(), marker.size(), &payload)) {
    183       // Something is wrong with this marker; does not care.
    184       continue;
    185     }
    186     if (payload.size() >= sizeof kExifTag &&
    187         !memcmp(payload.data(), kExifTag, sizeof kExifTag)) {
    188       if (blobs->exif.empty()) {
    189         blobs->exif.resize(payload.size() - sizeof kExifTag);
    190         memcpy(blobs->exif.data(), payload.data() + sizeof kExifTag,
    191                payload.size() - sizeof kExifTag);
    192       } else {
    193         JXL_WARNING(
    194             "ReJPEG: multiple Exif blobs, storing only first one in the JPEG "
    195             "XL container\n");
    196       }
    197     }
    198     if (payload.size() >= sizeof kXMPTag &&
    199         !memcmp(payload.data(), kXMPTag, sizeof kXMPTag)) {
    200       if (blobs->xmp.empty()) {
    201         blobs->xmp.resize(payload.size() - sizeof kXMPTag);
    202         memcpy(blobs->xmp.data(), payload.data() + sizeof kXMPTag,
    203                payload.size() - sizeof kXMPTag);
    204       } else {
    205         JXL_WARNING(
    206             "ReJPEG: multiple XMP blobs, storing only first one in the JPEG "
    207             "XL container\n");
    208       }
    209     }
    210   }
    211   return true;
    212 }
    213 
    214 inline bool IsJPG(const Span<const uint8_t> bytes) {
    215   return bytes.size() >= 2 && bytes[0] == 0xFF && bytes[1] == 0xD8;
    216 }
    217 
    218 }  // namespace
    219 
    220 void SetColorEncodingFromJpegData(const jpeg::JPEGData& jpg,
    221                                   ColorEncoding* color_encoding) {
    222   IccBytes icc_profile;
    223   if (!ParseChunkedMarker(jpg, kApp2, ByteSpan(kIccProfileTag), &icc_profile)) {
    224     JXL_WARNING("ReJPEG: corrupted ICC profile\n");
    225     icc_profile.clear();
    226   }
    227 
    228   if (icc_profile.empty()) {
    229     bool is_gray = (jpg.components.size() == 1);
    230     *color_encoding = ColorEncoding::SRGB(is_gray);
    231   } else {
    232     color_encoding->SetICCRaw(std::move(icc_profile));
    233   }
    234 }
    235 
    236 Status SetChromaSubsamplingFromJpegData(const JPEGData& jpg,
    237                                         YCbCrChromaSubsampling* cs) {
    238   size_t nbcomp = jpg.components.size();
    239   if (nbcomp != 1 && nbcomp != 3) {
    240     return JXL_FAILURE("Cannot recompress JPEGs with neither 1 nor 3 channels");
    241   }
    242   if (nbcomp == 3) {
    243     uint8_t hsample[3];
    244     uint8_t vsample[3];
    245     for (size_t i = 0; i < nbcomp; i++) {
    246       hsample[i] = jpg.components[i].h_samp_factor;
    247       vsample[i] = jpg.components[i].v_samp_factor;
    248     }
    249     JXL_RETURN_IF_ERROR(cs->Set(hsample, vsample));
    250   } else if (nbcomp == 1) {
    251     uint8_t hsample[3];
    252     uint8_t vsample[3];
    253     for (size_t i = 0; i < 3; i++) {
    254       hsample[i] = jpg.components[0].h_samp_factor;
    255       vsample[i] = jpg.components[0].v_samp_factor;
    256     }
    257     JXL_RETURN_IF_ERROR(cs->Set(hsample, vsample));
    258   }
    259   return true;
    260 }
    261 
    262 Status SetColorTransformFromJpegData(const JPEGData& jpg,
    263                                      ColorTransform* color_transform) {
    264   size_t nbcomp = jpg.components.size();
    265   if (nbcomp != 1 && nbcomp != 3) {
    266     return JXL_FAILURE("Cannot recompress JPEGs with neither 1 nor 3 channels");
    267   }
    268   bool is_rgb = false;
    269   {
    270     const auto& markers = jpg.marker_order;
    271     // If there is a JFIF marker, this is YCbCr. Otherwise...
    272     if (std::find(markers.begin(), markers.end(), 0xE0) == markers.end()) {
    273       // Try to find an 'Adobe' marker.
    274       size_t app_markers = 0;
    275       size_t i = 0;
    276       for (; i < markers.size(); i++) {
    277         // This is an APP marker.
    278         if ((markers[i] & 0xF0) == 0xE0) {
    279           JXL_CHECK(app_markers < jpg.app_data.size());
    280           // APP14 marker
    281           if (markers[i] == 0xEE) {
    282             const auto& data = jpg.app_data[app_markers];
    283             if (data.size() == 15 && data[3] == 'A' && data[4] == 'd' &&
    284                 data[5] == 'o' && data[6] == 'b' && data[7] == 'e') {
    285               // 'Adobe' marker.
    286               is_rgb = data[14] == 0;
    287               break;
    288             }
    289           }
    290           app_markers++;
    291         }
    292       }
    293 
    294       if (i == markers.size()) {
    295         // No 'Adobe' marker, guess from component IDs.
    296         is_rgb = nbcomp == 3 && jpg.components[0].id == 'R' &&
    297                  jpg.components[1].id == 'G' && jpg.components[2].id == 'B';
    298       }
    299     }
    300   }
    301   *color_transform =
    302       (!is_rgb || nbcomp == 1) ? ColorTransform::kYCbCr : ColorTransform::kNone;
    303   return true;
    304 }
    305 
    306 Status EncodeJPEGData(JPEGData& jpeg_data, std::vector<uint8_t>* bytes,
    307                       const CompressParams& cparams) {
    308   bytes->clear();
    309   jpeg_data.app_marker_type.resize(jpeg_data.app_data.size(),
    310                                    AppMarkerType::kUnknown);
    311   JXL_RETURN_IF_ERROR(DetectIccProfile(jpeg_data));
    312   JXL_RETURN_IF_ERROR(DetectBlobs(jpeg_data));
    313 
    314   size_t total_data = 0;
    315   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
    316     if (jpeg_data.app_marker_type[i] != AppMarkerType::kUnknown) {
    317       continue;
    318     }
    319     total_data += jpeg_data.app_data[i].size();
    320   }
    321   for (size_t i = 0; i < jpeg_data.com_data.size(); i++) {
    322     total_data += jpeg_data.com_data[i].size();
    323   }
    324   for (size_t i = 0; i < jpeg_data.inter_marker_data.size(); i++) {
    325     total_data += jpeg_data.inter_marker_data[i].size();
    326   }
    327   total_data += jpeg_data.tail_data.size();
    328   size_t brotli_capacity = BrotliEncoderMaxCompressedSize(total_data);
    329 
    330   BitWriter writer;
    331   JXL_RETURN_IF_ERROR(Bundle::Write(jpeg_data, &writer, 0, nullptr));
    332   writer.ZeroPadToByte();
    333   {
    334     PaddedBytes serialized_jpeg_data = std::move(writer).TakeBytes();
    335     bytes->reserve(serialized_jpeg_data.size() + brotli_capacity);
    336     Bytes(serialized_jpeg_data).AppendTo(*bytes);
    337   }
    338 
    339   BrotliEncoderState* brotli_enc =
    340       BrotliEncoderCreateInstance(nullptr, nullptr, nullptr);
    341   int effort = cparams.brotli_effort;
    342   if (effort < 0) effort = 11 - static_cast<int>(cparams.speed_tier);
    343   BrotliEncoderSetParameter(brotli_enc, BROTLI_PARAM_QUALITY, effort);
    344   size_t initial_size = bytes->size();
    345   BrotliEncoderSetParameter(brotli_enc, BROTLI_PARAM_SIZE_HINT, total_data);
    346   bytes->resize(initial_size + brotli_capacity);
    347   size_t enc_size = 0;
    348   auto br_append = [&](const std::vector<uint8_t>& data, bool last) {
    349     size_t available_in = data.size();
    350     const uint8_t* in = data.data();
    351     uint8_t* out = &(*bytes)[initial_size + enc_size];
    352     do {
    353       uint8_t* out_before = out;
    354       msan::MemoryIsInitialized(in, available_in);
    355       JXL_CHECK(BrotliEncoderCompressStream(
    356           brotli_enc, last ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS,
    357           &available_in, &in, &brotli_capacity, &out, &enc_size));
    358       msan::UnpoisonMemory(out_before, out - out_before);
    359     } while (BrotliEncoderHasMoreOutput(brotli_enc) || available_in > 0);
    360   };
    361 
    362   for (size_t i = 0; i < jpeg_data.app_data.size(); i++) {
    363     if (jpeg_data.app_marker_type[i] != AppMarkerType::kUnknown) {
    364       continue;
    365     }
    366     br_append(jpeg_data.app_data[i], /*last=*/false);
    367   }
    368   for (size_t i = 0; i < jpeg_data.com_data.size(); i++) {
    369     br_append(jpeg_data.com_data[i], /*last=*/false);
    370   }
    371   for (size_t i = 0; i < jpeg_data.inter_marker_data.size(); i++) {
    372     br_append(jpeg_data.inter_marker_data[i], /*last=*/false);
    373   }
    374   br_append(jpeg_data.tail_data, /*last=*/true);
    375   BrotliEncoderDestroyInstance(brotli_enc);
    376   bytes->resize(initial_size + enc_size);
    377   return true;
    378 }
    379 
    380 Status DecodeImageJPG(const Span<const uint8_t> bytes, CodecInOut* io) {
    381   if (!IsJPG(bytes)) return false;
    382   io->frames.clear();
    383   io->frames.reserve(1);
    384   io->frames.emplace_back(&io->metadata.m);
    385   io->Main().jpeg_data = make_unique<jpeg::JPEGData>();
    386   jpeg::JPEGData* jpeg_data = io->Main().jpeg_data.get();
    387   if (!jpeg::ReadJpeg(bytes.data(), bytes.size(), jpeg::JpegReadMode::kReadAll,
    388                       jpeg_data)) {
    389     return JXL_FAILURE("Error reading JPEG");
    390   }
    391   SetColorEncodingFromJpegData(*jpeg_data, &io->metadata.m.color_encoding);
    392   JXL_RETURN_IF_ERROR(SetBlobsFromJpegData(*jpeg_data, &io->blobs));
    393   JXL_RETURN_IF_ERROR(SetChromaSubsamplingFromJpegData(
    394       *jpeg_data, &io->Main().chroma_subsampling));
    395   JXL_RETURN_IF_ERROR(
    396       SetColorTransformFromJpegData(*jpeg_data, &io->Main().color_transform));
    397 
    398   io->metadata.m.SetIntensityTarget(kDefaultIntensityTarget);
    399   io->metadata.m.SetUintSamples(BITS_IN_JSAMPLE);
    400   JXL_ASSIGN_OR_RETURN(Image3F tmp,
    401                        Image3F::Create(jpeg_data->width, jpeg_data->height));
    402   io->SetFromImage(std::move(tmp), io->metadata.m.color_encoding);
    403   SetIntensityTarget(&io->metadata.m);
    404   return true;
    405 }
    406 
    407 }  // namespace jpeg
    408 }  // namespace jxl