libjxl

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

benchmark_codec_avif.cc (15322B)


      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 #include "tools/benchmark/benchmark_codec_avif.h"
      6 
      7 #include <avif/avif.h>
      8 #include <jxl/cms.h>
      9 
     10 #include "lib/extras/time.h"
     11 #include "lib/jxl/base/span.h"
     12 #include "lib/jxl/codec_in_out.h"
     13 #include "lib/jxl/dec_external_image.h"
     14 #include "lib/jxl/enc_external_image.h"
     15 #include "tools/cmdline.h"
     16 #include "tools/thread_pool_internal.h"
     17 
     18 #define JXL_RETURN_IF_AVIF_ERROR(result)                                       \
     19   do {                                                                         \
     20     avifResult jxl_return_if_avif_error_result = (result);                     \
     21     if (jxl_return_if_avif_error_result != AVIF_RESULT_OK) {                   \
     22       return JXL_FAILURE("libavif error: %s",                                  \
     23                          avifResultToString(jxl_return_if_avif_error_result)); \
     24     }                                                                          \
     25   } while (false)
     26 
     27 namespace jpegxl {
     28 namespace tools {
     29 
     30 using ::jxl::Bytes;
     31 using ::jxl::CodecInOut;
     32 using ::jxl::IccBytes;
     33 using ::jxl::ImageBundle;
     34 using ::jxl::Primaries;
     35 using ::jxl::Span;
     36 using ::jxl::ThreadPool;
     37 using ::jxl::TransferFunction;
     38 using ::jxl::WhitePoint;
     39 
     40 namespace {
     41 
     42 size_t GetNumThreads(ThreadPool* pool) {
     43   size_t result = 0;
     44   const auto count_threads = [&](const size_t num_threads) {
     45     result = num_threads;
     46     return true;
     47   };
     48   const auto no_op = [&](const uint32_t /*task*/, size_t /*thread*/) {};
     49   (void)jxl::RunOnPool(pool, 0, 1, count_threads, no_op, "Compress");
     50   return result;
     51 }
     52 
     53 struct AvifArgs {
     54   avifPixelFormat chroma_subsampling = AVIF_PIXEL_FORMAT_YUV444;
     55 };
     56 
     57 AvifArgs* const avifargs = new AvifArgs;
     58 
     59 bool ParseChromaSubsampling(const char* arg, avifPixelFormat* subsampling) {
     60   if (strcmp(arg, "444") == 0) {
     61     *subsampling = AVIF_PIXEL_FORMAT_YUV444;
     62     return true;
     63   }
     64   if (strcmp(arg, "422") == 0) {
     65     *subsampling = AVIF_PIXEL_FORMAT_YUV422;
     66     return true;
     67   }
     68   if (strcmp(arg, "420") == 0) {
     69     *subsampling = AVIF_PIXEL_FORMAT_YUV420;
     70     return true;
     71   }
     72   if (strcmp(arg, "400") == 0) {
     73     *subsampling = AVIF_PIXEL_FORMAT_YUV400;
     74     return true;
     75   }
     76   return false;
     77 }
     78 
     79 void SetUpAvifColor(const ColorEncoding& color, bool rgb,
     80                     avifImage* const image) {
     81   bool need_icc = (color.GetWhitePointType() != WhitePoint::kD65);
     82 
     83   image->matrixCoefficients =
     84       rgb ? AVIF_MATRIX_COEFFICIENTS_IDENTITY : AVIF_MATRIX_COEFFICIENTS_BT709;
     85   if (!color.HasPrimaries()) {
     86     need_icc = true;
     87   } else {
     88     switch (color.GetPrimariesType()) {
     89       case Primaries::kSRGB:
     90         image->colorPrimaries = AVIF_COLOR_PRIMARIES_BT709;
     91         break;
     92       case Primaries::k2100:
     93         image->colorPrimaries = AVIF_COLOR_PRIMARIES_BT2020;
     94         image->matrixCoefficients = AVIF_MATRIX_COEFFICIENTS_BT2020_NCL;
     95         break;
     96       default:
     97         need_icc = true;
     98         image->colorPrimaries = AVIF_COLOR_PRIMARIES_UNKNOWN;
     99         break;
    100     }
    101   }
    102 
    103   switch (color.Tf().GetTransferFunction()) {
    104     case TransferFunction::kSRGB:
    105       image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_SRGB;
    106       break;
    107     case TransferFunction::kLinear:
    108       image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_LINEAR;
    109       break;
    110     case TransferFunction::kPQ:
    111       image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_SMPTE2084;
    112       break;
    113     case TransferFunction::kHLG:
    114       image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_HLG;
    115       break;
    116     default:
    117       need_icc = true;
    118       image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_UNKNOWN;
    119       break;
    120   }
    121 
    122   if (need_icc) {
    123     avifImageSetProfileICC(image, color.ICC().data(), color.ICC().size());
    124   }
    125 }
    126 
    127 Status ReadAvifColor(const avifImage* const image, ColorEncoding* const color) {
    128   if (image->icc.size != 0) {
    129     IccBytes icc;
    130     icc.assign(image->icc.data, image->icc.data + image->icc.size);
    131     return color->SetICC(std::move(icc), JxlGetDefaultCms());
    132   }
    133 
    134   JXL_RETURN_IF_ERROR(color->SetWhitePointType(WhitePoint::kD65));
    135   switch (image->colorPrimaries) {
    136     case AVIF_COLOR_PRIMARIES_BT709:
    137       JXL_RETURN_IF_ERROR(color->SetPrimariesType(Primaries::kSRGB));
    138       break;
    139     case AVIF_COLOR_PRIMARIES_BT2020:
    140       JXL_RETURN_IF_ERROR(color->SetPrimariesType(Primaries::k2100));
    141       break;
    142     default:
    143       return JXL_FAILURE("unsupported avif primaries");
    144   }
    145   jxl::cms::CustomTransferFunction& tf = color->Tf();
    146   switch (image->transferCharacteristics) {
    147     case AVIF_TRANSFER_CHARACTERISTICS_BT470M:
    148       JXL_RETURN_IF_ERROR(tf.SetGamma(2.2));
    149       break;
    150     case AVIF_TRANSFER_CHARACTERISTICS_BT470BG:
    151       JXL_RETURN_IF_ERROR(tf.SetGamma(2.8));
    152       break;
    153     case AVIF_TRANSFER_CHARACTERISTICS_LINEAR:
    154       tf.SetTransferFunction(TransferFunction::kLinear);
    155       break;
    156     case AVIF_TRANSFER_CHARACTERISTICS_SRGB:
    157       tf.SetTransferFunction(TransferFunction::kSRGB);
    158       break;
    159     case AVIF_TRANSFER_CHARACTERISTICS_SMPTE2084:
    160       tf.SetTransferFunction(TransferFunction::kPQ);
    161       break;
    162     case AVIF_TRANSFER_CHARACTERISTICS_HLG:
    163       tf.SetTransferFunction(TransferFunction::kHLG);
    164       break;
    165     default:
    166       return JXL_FAILURE("unsupported avif TRC");
    167   }
    168   return color->CreateICC();
    169 }
    170 
    171 }  // namespace
    172 
    173 Status AddCommandLineOptionsAvifCodec(BenchmarkArgs* args) {
    174   args->cmdline.AddOptionValue(
    175       '\0', "avif_chroma_subsampling", "444/422/420/400",
    176       "default AVIF chroma subsampling (default: 444).",
    177       &avifargs->chroma_subsampling, &ParseChromaSubsampling);
    178   return true;
    179 }
    180 
    181 class AvifCodec : public ImageCodec {
    182  public:
    183   explicit AvifCodec(const BenchmarkArgs& args) : ImageCodec(args) {
    184     chroma_subsampling_ = avifargs->chroma_subsampling;
    185   }
    186 
    187   Status ParseParam(const std::string& param) override {
    188     if (param.compare(0, 3, "yuv") == 0) {
    189       if (param.size() != 6) return false;
    190       return ParseChromaSubsampling(param.c_str() + 3, &chroma_subsampling_);
    191     }
    192     if (param == "rgb") {
    193       rgb_ = true;
    194       return true;
    195     }
    196     if (param.compare(0, 10, "log2_cols=") == 0) {
    197       log2_cols = strtol(param.c_str() + 10, nullptr, 10);
    198       return true;
    199     }
    200     if (param.compare(0, 10, "log2_rows=") == 0) {
    201       log2_rows = strtol(param.c_str() + 10, nullptr, 10);
    202       return true;
    203     }
    204     if (param[0] == 's') {
    205       speed_ = strtol(param.c_str() + 1, nullptr, 10);
    206       return true;
    207     }
    208     if (param == "aomenc") {
    209       encoder_ = AVIF_CODEC_CHOICE_AOM;
    210       return true;
    211     }
    212     if (param == "aomdec") {
    213       decoder_ = AVIF_CODEC_CHOICE_AOM;
    214       return true;
    215     }
    216     if (param == "aom") {
    217       encoder_ = AVIF_CODEC_CHOICE_AOM;
    218       decoder_ = AVIF_CODEC_CHOICE_AOM;
    219       return true;
    220     }
    221     if (param == "rav1e") {
    222       encoder_ = AVIF_CODEC_CHOICE_RAV1E;
    223       return true;
    224     }
    225     if (param == "dav1d") {
    226       decoder_ = AVIF_CODEC_CHOICE_DAV1D;
    227       return true;
    228     }
    229     if (param.compare(0, 2, "a=") == 0) {
    230       std::string subparam = param.substr(2);
    231       size_t pos = subparam.find('=');
    232       if (pos == std::string::npos) {
    233         codec_specific_options_.emplace_back(subparam, "");
    234       } else {
    235         std::string key = subparam.substr(0, pos);
    236         std::string value = subparam.substr(pos + 1);
    237         codec_specific_options_.emplace_back(key, value);
    238       }
    239       return true;
    240     }
    241     return ImageCodec::ParseParam(param);
    242   }
    243 
    244   Status Compress(const std::string& filename, const PackedPixelFile& ppf,
    245                   ThreadPool* pool, std::vector<uint8_t>* compressed,
    246                   jpegxl::tools::SpeedStats* speed_stats) override {
    247     CodecInOut io;
    248     JXL_RETURN_IF_ERROR(
    249         jxl::extras::ConvertPackedPixelFileToCodecInOut(ppf, pool, &io));
    250     return Compress(filename, &io, pool, compressed, speed_stats);
    251   }
    252 
    253   Status Compress(const std::string& filename, const CodecInOut* io,
    254                   ThreadPool* pool, std::vector<uint8_t>* compressed,
    255                   SpeedStats* speed_stats) {
    256     double elapsed_convert_image = 0;
    257     size_t max_threads = GetNumThreads(pool);
    258     const double start = jxl::Now();
    259     {
    260       const auto depth =
    261           std::min<int>(16, io->metadata.m.bit_depth.bits_per_sample);
    262       std::unique_ptr<avifEncoder, void (*)(avifEncoder*)> encoder(
    263           avifEncoderCreate(), &avifEncoderDestroy);
    264       encoder->codecChoice = encoder_;
    265       // TODO(sboukortt): configure this separately.
    266       encoder->minQuantizer = 0;
    267       encoder->maxQuantizer = 63;
    268 #if AVIF_VERSION >= 1000300
    269       encoder->quality = q_target_;
    270       encoder->qualityAlpha = q_target_;
    271 #endif
    272       encoder->tileColsLog2 = log2_cols;
    273       encoder->tileRowsLog2 = log2_rows;
    274       encoder->speed = speed_;
    275       encoder->maxThreads = max_threads;
    276       for (const auto& opts : codec_specific_options_) {
    277 #if AVIF_VERSION_MAJOR >= 1
    278         JXL_RETURN_IF_AVIF_ERROR(avifEncoderSetCodecSpecificOption(
    279             encoder.get(), opts.first.c_str(), opts.second.c_str()));
    280 #else
    281         (void)avifEncoderSetCodecSpecificOption(
    282             encoder.get(), opts.first.c_str(), opts.second.c_str());
    283 #endif
    284       }
    285       avifAddImageFlags add_image_flags = AVIF_ADD_IMAGE_FLAG_SINGLE;
    286       if (io->metadata.m.have_animation) {
    287         encoder->timescale = std::lround(
    288             static_cast<float>(io->metadata.m.animation.tps_numerator) /
    289             io->metadata.m.animation.tps_denominator);
    290         add_image_flags = AVIF_ADD_IMAGE_FLAG_NONE;
    291       }
    292       for (const ImageBundle& ib : io->frames) {
    293         std::unique_ptr<avifImage, void (*)(avifImage*)> image(
    294             avifImageCreate(ib.xsize(), ib.ysize(), depth, chroma_subsampling_),
    295             &avifImageDestroy);
    296         image->width = ib.xsize();
    297         image->height = ib.ysize();
    298         image->depth = depth;
    299         SetUpAvifColor(ib.c_current(), rgb_, image.get());
    300         std::unique_ptr<avifRWData, void (*)(avifRWData*)> icc_freer(
    301             &image->icc, &avifRWDataFree);
    302         avifRGBImage rgb_image;
    303         avifRGBImageSetDefaults(&rgb_image, image.get());
    304         rgb_image.format =
    305             ib.HasAlpha() ? AVIF_RGB_FORMAT_RGBA : AVIF_RGB_FORMAT_RGB;
    306         avifRGBImageAllocatePixels(&rgb_image);
    307         std::unique_ptr<avifRGBImage, void (*)(avifRGBImage*)> pixels_freer(
    308             &rgb_image, &avifRGBImageFreePixels);
    309         const double start_convert_image = jxl::Now();
    310         JXL_RETURN_IF_ERROR(ConvertToExternal(
    311             ib, depth, /*float_out=*/false,
    312             /*num_channels=*/ib.HasAlpha() ? 4 : 3, JXL_NATIVE_ENDIAN,
    313             /*stride=*/rgb_image.rowBytes, pool, rgb_image.pixels,
    314             rgb_image.rowBytes * rgb_image.height,
    315             /*out_callback=*/{}, jxl::Orientation::kIdentity));
    316         const double end_convert_image = jxl::Now();
    317         elapsed_convert_image += end_convert_image - start_convert_image;
    318         JXL_RETURN_IF_AVIF_ERROR(avifImageRGBToYUV(image.get(), &rgb_image));
    319         JXL_RETURN_IF_AVIF_ERROR(avifEncoderAddImage(
    320             encoder.get(), image.get(), ib.duration, add_image_flags));
    321       }
    322       avifRWData buffer = AVIF_DATA_EMPTY;
    323       JXL_RETURN_IF_AVIF_ERROR(avifEncoderFinish(encoder.get(), &buffer));
    324       compressed->assign(buffer.data, buffer.data + buffer.size);
    325       avifRWDataFree(&buffer);
    326     }
    327     const double end = jxl::Now();
    328     speed_stats->NotifyElapsed(end - start - elapsed_convert_image);
    329     return true;
    330   }
    331 
    332   Status Decompress(const std::string& filename,
    333                     const Span<const uint8_t> compressed, ThreadPool* pool,
    334                     PackedPixelFile* ppf,
    335                     jpegxl::tools::SpeedStats* speed_stats) override {
    336     CodecInOut io;
    337     JXL_RETURN_IF_ERROR(
    338         Decompress(filename, compressed, pool, &io, speed_stats));
    339     JxlPixelFormat format{0, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
    340     return jxl::extras::ConvertCodecInOutToPackedPixelFile(
    341         io, format, io.Main().c_current(), pool, ppf);
    342   };
    343 
    344   Status Decompress(const std::string& filename,
    345                     const Span<const uint8_t> compressed, ThreadPool* pool,
    346                     CodecInOut* io, SpeedStats* speed_stats) {
    347     io->frames.clear();
    348     size_t max_threads = GetNumThreads(pool);
    349     double elapsed_convert_image = 0;
    350     const double start = jxl::Now();
    351     {
    352       std::unique_ptr<avifDecoder, void (*)(avifDecoder*)> decoder(
    353           avifDecoderCreate(), &avifDecoderDestroy);
    354       decoder->codecChoice = decoder_;
    355       decoder->maxThreads = max_threads;
    356       JXL_RETURN_IF_AVIF_ERROR(avifDecoderSetIOMemory(
    357           decoder.get(), compressed.data(), compressed.size()));
    358       JXL_RETURN_IF_AVIF_ERROR(avifDecoderParse(decoder.get()));
    359       const bool has_alpha = decoder->alphaPresent;
    360       io->metadata.m.have_animation = decoder->imageCount > 1;
    361       io->metadata.m.animation.tps_numerator = decoder->timescale;
    362       io->metadata.m.animation.tps_denominator = 1;
    363       io->metadata.m.SetUintSamples(decoder->image->depth);
    364       io->SetSize(decoder->image->width, decoder->image->height);
    365       avifResult next_image;
    366       while ((next_image = avifDecoderNextImage(decoder.get())) ==
    367              AVIF_RESULT_OK) {
    368         ColorEncoding color;
    369         JXL_RETURN_IF_ERROR(ReadAvifColor(decoder->image, &color));
    370         avifRGBImage rgb_image;
    371         avifRGBImageSetDefaults(&rgb_image, decoder->image);
    372         rgb_image.format =
    373             has_alpha ? AVIF_RGB_FORMAT_RGBA : AVIF_RGB_FORMAT_RGB;
    374         avifRGBImageAllocatePixels(&rgb_image);
    375         std::unique_ptr<avifRGBImage, void (*)(avifRGBImage*)> pixels_freer(
    376             &rgb_image, &avifRGBImageFreePixels);
    377         JXL_RETURN_IF_AVIF_ERROR(avifImageYUVToRGB(decoder->image, &rgb_image));
    378         const double start_convert_image = jxl::Now();
    379         {
    380           JxlPixelFormat format = {
    381               (has_alpha ? 4u : 3u),
    382               (rgb_image.depth <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16),
    383               JXL_NATIVE_ENDIAN, 0};
    384           ImageBundle ib(&io->metadata.m);
    385           JXL_RETURN_IF_ERROR(ConvertFromExternal(
    386               Bytes(rgb_image.pixels, rgb_image.height * rgb_image.rowBytes),
    387               rgb_image.width, rgb_image.height, color, rgb_image.depth, format,
    388               pool, &ib));
    389           io->frames.push_back(std::move(ib));
    390         }
    391         const double end_convert_image = jxl::Now();
    392         elapsed_convert_image += end_convert_image - start_convert_image;
    393       }
    394       if (next_image != AVIF_RESULT_NO_IMAGES_REMAINING) {
    395         JXL_RETURN_IF_AVIF_ERROR(next_image);
    396       }
    397     }
    398     const double end = jxl::Now();
    399     speed_stats->NotifyElapsed(end - start - elapsed_convert_image);
    400     return true;
    401   }
    402 
    403  protected:
    404   avifPixelFormat chroma_subsampling_;
    405   avifCodecChoice encoder_ = AVIF_CODEC_CHOICE_AUTO;
    406   avifCodecChoice decoder_ = AVIF_CODEC_CHOICE_AUTO;
    407   bool rgb_ = false;
    408   int speed_ = AVIF_SPEED_DEFAULT;
    409   int log2_cols = 0;
    410   int log2_rows = 0;
    411   std::vector<std::pair<std::string, std::string>> codec_specific_options_;
    412 };
    413 
    414 ImageCodec* CreateNewAvifCodec(const BenchmarkArgs& args) {
    415   return new AvifCodec(args);
    416 }
    417 
    418 }  // namespace tools
    419 }  // namespace jpegxl