libavcodec/libjxldec.c

   1 /*
   2  * JPEG XL decoding support via libjxl
   3  * Copyright (c) 2021 Leo Izen <leo.izen@gmail.com>
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 /**
  23  * @file
  24  * JPEG XL decoder using libjxl
  25  */
  26
  27 #include "libavutil/avassert.h"
  28 #include "libavutil/buffer.h"
  29 #include "libavutil/common.h"
  30 #include "libavutil/csp.h"
  31 #include "libavutil/error.h"
  32 #include "libavutil/mem.h"
  33 #include "libavutil/pixdesc.h"
  34 #include "libavutil/pixfmt.h"
  35 #include "libavutil/frame.h"
  36
  37 #include "avcodec.h"
  38 #include "codec_internal.h"
  39 #include "internal.h"
  40
  41 #include <jxl/decode.h>
  42 #include <jxl/thread_parallel_runner.h>
  43 #include "libjxl.h"
  44
  45 typedef struct LibJxlDecodeContext {
  46     void *runner;
  47     JxlDecoder *decoder;
  48     JxlBasicInfo basic_info;
  49     JxlPixelFormat jxl_pixfmt;
  50     JxlDecoderStatus events;
  51     AVBufferRef *iccp;
  52 } LibJxlDecodeContext;
  53
  54 static int libjxl_init_jxl_decoder(AVCodecContext *avctx)
  55 {
  56     LibJxlDecodeContext *ctx = avctx->priv_data;
  57
  58     ctx->events = JXL_DEC_BASIC_INFO | JXL_DEC_FULL_IMAGE | JXL_DEC_COLOR_ENCODING;
  59     if (JxlDecoderSubscribeEvents(ctx->decoder, ctx->events) != JXL_DEC_SUCCESS) {
  60         av_log(avctx, AV_LOG_ERROR, "Error subscribing to JXL events\n");
  61         return AVERROR_EXTERNAL;
  62     }
  63
  64     if (JxlDecoderSetParallelRunner(ctx->decoder, JxlThreadParallelRunner, ctx->runner) != JXL_DEC_SUCCESS) {
  65         av_log(avctx, AV_LOG_ERROR, "Failed to set JxlThreadParallelRunner\n");
  66         return AVERROR_EXTERNAL;
  67     }
  68
  69     memset(&ctx->basic_info, 0, sizeof(JxlBasicInfo));
  70     memset(&ctx->jxl_pixfmt, 0, sizeof(JxlPixelFormat));
  71
  72     return 0;
  73 }
  74
  75 static av_cold int libjxl_decode_init(AVCodecContext *avctx)
  76 {
  77     LibJxlDecodeContext *ctx = avctx->priv_data;
  78     JxlMemoryManager manager;
  79
  80     ff_libjxl_init_memory_manager(&manager);
  81     ctx->decoder = JxlDecoderCreate(&manager);
  82     if (!ctx->decoder) {
  83         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlDecoder\n");
  84         return AVERROR_EXTERNAL;
  85     }
  86
  87     ctx->runner = JxlThreadParallelRunnerCreate(&manager, ff_libjxl_get_threadcount(avctx->thread_count));
  88     if (!ctx->runner) {
  89         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlThreadParallelRunner\n");
  90         return AVERROR_EXTERNAL;
  91     }
  92
  93     return libjxl_init_jxl_decoder(avctx);
  94 }
  95
  96 static enum AVPixelFormat libjxl_get_pix_fmt(void *avctx, const JxlBasicInfo *basic_info, JxlPixelFormat *format)
  97 {
  98     format->endianness = JXL_NATIVE_ENDIAN;
  99     format->num_channels = basic_info->num_color_channels + (basic_info->alpha_bits > 0);
 100     /* Gray */
 101     if (basic_info->num_color_channels == 1) {
 102         if (basic_info->bits_per_sample <= 8) {
 103             format->data_type = JXL_TYPE_UINT8;
 104             return basic_info->alpha_bits ? AV_PIX_FMT_YA8 : AV_PIX_FMT_GRAY8;
 105         }
 106         if (basic_info->exponent_bits_per_sample || basic_info->bits_per_sample > 16) {
 107             if (basic_info->alpha_bits)
 108                 return AV_PIX_FMT_NONE;
 109             format->data_type = JXL_TYPE_FLOAT;
 110             return AV_PIX_FMT_GRAYF32;
 111         }
 112         format->data_type = JXL_TYPE_UINT16;
 113         return basic_info->alpha_bits ? AV_PIX_FMT_YA16 : AV_PIX_FMT_GRAY16;
 114     }
 115     /* rgb only */
 116     /* libjxl only supports packed RGB and gray output at the moment */
 117     if (basic_info->num_color_channels == 3) {
 118         if (basic_info->bits_per_sample <= 8) {
 119             format->data_type = JXL_TYPE_UINT8;
 120             return basic_info->alpha_bits ? AV_PIX_FMT_RGBA : AV_PIX_FMT_RGB24;
 121         }
 122         if (basic_info->bits_per_sample > 16)
 123             av_log(avctx, AV_LOG_WARNING, "Downsampling larger integer to 16-bit via libjxl\n");
 124         if (basic_info->exponent_bits_per_sample)
 125             av_log(avctx, AV_LOG_WARNING, "Downsampling float to 16-bit integer via libjxl\n");
 126         format->data_type = JXL_TYPE_UINT16;
 127         return basic_info->alpha_bits ? AV_PIX_FMT_RGBA64 : AV_PIX_FMT_RGB48;
 128     }
 129
 130     return AV_PIX_FMT_NONE;
 131 }
 132
 133 static enum AVColorPrimaries libjxl_get_primaries(void *avctx, const JxlColorEncoding *jxl_color)
 134 {
 135     AVColorPrimariesDesc desc;
 136     enum AVColorPrimaries prim;
 137
 138     /* libjxl populates these double values even if it uses an enum space */
 139     desc.prim.r.x = av_d2q(jxl_color->primaries_red_xy[0], 300000);
 140     desc.prim.r.y = av_d2q(jxl_color->primaries_red_xy[1], 300000);
 141     desc.prim.g.x = av_d2q(jxl_color->primaries_green_xy[0], 300000);
 142     desc.prim.g.y = av_d2q(jxl_color->primaries_green_xy[1], 300000);
 143     desc.prim.b.x = av_d2q(jxl_color->primaries_blue_xy[0], 300000);
 144     desc.prim.b.y = av_d2q(jxl_color->primaries_blue_xy[1], 300000);
 145     desc.wp.x = av_d2q(jxl_color->white_point_xy[0], 300000);
 146     desc.wp.y = av_d2q(jxl_color->white_point_xy[1], 300000);
 147
 148     prim = av_csp_primaries_id_from_desc(&desc);
 149     if (prim == AVCOL_PRI_UNSPECIFIED) {
 150         /* try D65 with the same primaries */
 151         /* BT.709 uses D65 white point */
 152         desc.wp = av_csp_primaries_desc_from_id(AVCOL_PRI_BT709)->wp;
 153         av_log(avctx, AV_LOG_WARNING, "Changing unknown white point to D65\n");
 154         prim = av_csp_primaries_id_from_desc(&desc);
 155     }
 156
 157     return prim;
 158 }
 159
 160 static enum AVColorTransferCharacteristic libjxl_get_trc(void *avctx, const JxlColorEncoding *jxl_color)
 161 {
 162     switch (jxl_color->transfer_function) {
 163     case JXL_TRANSFER_FUNCTION_709: return AVCOL_TRC_BT709;
 164     case JXL_TRANSFER_FUNCTION_LINEAR: return AVCOL_TRC_LINEAR;
 165     case JXL_TRANSFER_FUNCTION_SRGB: return AVCOL_TRC_IEC61966_2_1;
 166     case JXL_TRANSFER_FUNCTION_PQ: return AVCOL_TRC_SMPTE2084;
 167     case JXL_TRANSFER_FUNCTION_DCI: return AVCOL_TRC_SMPTE428;
 168     case JXL_TRANSFER_FUNCTION_HLG: return AVCOL_TRC_ARIB_STD_B67;
 169     case JXL_TRANSFER_FUNCTION_GAMMA:
 170         if (jxl_color->gamma > 2.199 && jxl_color->gamma < 2.201)
 171             return AVCOL_TRC_GAMMA22;
 172         else if (jxl_color->gamma > 2.799 && jxl_color->gamma < 2.801)
 173             return AVCOL_TRC_GAMMA28;
 174         else
 175             av_log(avctx, AV_LOG_WARNING, "Unsupported gamma transfer: %f\n", jxl_color->gamma);
 176         break;
 177     default:
 178         av_log(avctx, AV_LOG_WARNING, "Unknown transfer function: %d\n", jxl_color->transfer_function);
 179     }
 180
 181     return AVCOL_TRC_UNSPECIFIED;
 182 }
 183
 184 static int libjxl_get_icc(AVCodecContext *avctx)
 185 {
 186     LibJxlDecodeContext *ctx = avctx->priv_data;
 187     size_t icc_len;
 188     JxlDecoderStatus jret;
 189     /* an ICC profile is present, and we can meaningfully get it,
 190      * because the pixel data is not XYB-encoded */
 191     jret = JxlDecoderGetICCProfileSize(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA, &icc_len);
 192     if (jret == JXL_DEC_SUCCESS && icc_len > 0) {
 193         av_buffer_unref(&ctx->iccp);
 194         ctx->iccp = av_buffer_alloc(icc_len);
 195         if (!ctx->iccp)
 196             return AVERROR(ENOMEM);
 197         jret = JxlDecoderGetColorAsICCProfile(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA,
 198                                                 ctx->iccp->data, icc_len);
 199         if (jret != JXL_DEC_SUCCESS) {
 200             av_log(avctx, AV_LOG_WARNING, "Unable to obtain ICC Profile\n");
 201             av_buffer_unref(&ctx->iccp);
 202         }
 203     }
 204
 205     return 0;
 206 }
 207
 208 /*
 209  * There's generally four cases when it comes to decoding a libjxl image
 210  * with regard to color encoding:
 211  * (a) There is an embedded ICC profile in the image, and the image is XYB-encoded.
 212  * (b) There is an embedded ICC profile in the image, and the image is not XYB-encoded.
 213  * (c) There is no embedded ICC profile, and FFmpeg supports the tagged colorspace.
 214  * (d) There is no embedded ICC profile, and FFmpeg does not support the tagged colorspace.
 215  *
 216  * In case (b), we forward the pixel data as is and forward the ICC Profile as-is.
 217  * In case (c), we request the pixel data in the space it's tagged as,
 218  *     and tag the space accordingly.
 219  * In case (a), libjxl does not support getting the pixel data in the space described by the ICC
 220  *     profile, so instead we request the pixel data in BT.2020/PQ as it is the widest
 221  *     space that FFmpeg supports.
 222  * In case (d), we also request wide-gamut pixel data as a fallback since FFmpeg doesn't support
 223  *     the custom primaries tagged in the space.
 224  */
 225 static int libjxl_color_encoding_event(AVCodecContext *avctx, AVFrame *frame)
 226 {
 227     LibJxlDecodeContext *ctx = avctx->priv_data;
 228     JxlDecoderStatus jret;
 229     int ret;
 230     JxlColorEncoding jxl_color;
 231     /* set this flag if we need to fall back on wide gamut */
 232     int fallback = 0;
 233
 234     jret = JxlDecoderGetColorAsEncodedProfile(ctx->decoder, NULL, JXL_COLOR_PROFILE_TARGET_ORIGINAL, &jxl_color);
 235     if (jret == JXL_DEC_SUCCESS) {
 236         /* enum values describe the colors of this image */
 237         jret = JxlDecoderSetPreferredColorProfile(ctx->decoder, &jxl_color);
 238         if (jret == JXL_DEC_SUCCESS)
 239             jret = JxlDecoderGetColorAsEncodedProfile(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA, &jxl_color);
 240         /* if we couldn't successfully request the pixel data space, we fall back on wide gamut */
 241         /* this code path is very unlikely to happen in practice */
 242         if (jret != JXL_DEC_SUCCESS)
 243             fallback = 1;
 244     } else {
 245         /* an ICC Profile is present in the stream */
 246         if (ctx->basic_info.uses_original_profile) {
 247             /* uses_original_profile is the same as !xyb_encoded */
 248             av_log(avctx, AV_LOG_VERBOSE, "Using embedded ICC Profile\n");
 249             if ((ret = libjxl_get_icc(avctx)) < 0)
 250                 return ret;
 251         } else {
 252             /*
 253              * an XYB-encoded image with an embedded ICC profile can't always have the
 254              * pixel data requested in the original space, so libjxl has no feature
 255              * to allow this to happen, so we fall back on wide gamut
 256              */
 257             fallback = 1;
 258         }
 259     }
 260
 261     avctx->color_range = frame->color_range = AVCOL_RANGE_JPEG;
 262     if (ctx->jxl_pixfmt.num_channels >= 3)
 263         avctx->colorspace = AVCOL_SPC_RGB;
 264     avctx->color_primaries = AVCOL_PRI_UNSPECIFIED;
 265     avctx->color_trc = AVCOL_TRC_UNSPECIFIED;
 266
 267     if (!ctx->iccp) {
 268         /* checking enum values */
 269         if (!fallback) {
 270             if (avctx->colorspace == AVCOL_SPC_RGB)
 271                 avctx->color_primaries = libjxl_get_primaries(avctx, &jxl_color);
 272             avctx->color_trc = libjxl_get_trc(avctx, &jxl_color);
 273         }
 274         /* fall back on wide gamut if enum values fail */
 275         if (avctx->color_primaries == AVCOL_PRI_UNSPECIFIED) {
 276             if (avctx->colorspace == AVCOL_SPC_RGB) {
 277                 av_log(avctx, AV_LOG_WARNING, "Falling back on wide gamut output\n");
 278                 jxl_color.primaries = JXL_PRIMARIES_2100;
 279                 avctx->color_primaries = AVCOL_PRI_BT2020;
 280             }
 281             /* libjxl requires this set even for grayscale */
 282             jxl_color.white_point = JXL_WHITE_POINT_D65;
 283         }
 284         if (avctx->color_trc == AVCOL_TRC_UNSPECIFIED) {
 285             if (ctx->jxl_pixfmt.data_type == JXL_TYPE_FLOAT
 286                 || ctx->jxl_pixfmt.data_type == JXL_TYPE_FLOAT16) {
 287                 av_log(avctx, AV_LOG_WARNING, "Falling back on Linear Light transfer\n");
 288                 jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
 289                 avctx->color_trc = AVCOL_TRC_LINEAR;
 290             } else {
 291                 av_log(avctx, AV_LOG_WARNING, "Falling back on iec61966-2-1/sRGB transfer\n");
 292                 jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
 293                 avctx->color_trc = AVCOL_TRC_IEC61966_2_1;
 294             }
 295         }
 296         /* all colors will be in-gamut so we want accurate colors */
 297         jxl_color.rendering_intent = JXL_RENDERING_INTENT_RELATIVE;
 298         jxl_color.color_space = avctx->colorspace == AVCOL_SPC_RGB ? JXL_COLOR_SPACE_RGB : JXL_COLOR_SPACE_GRAY;
 299         jret = JxlDecoderSetPreferredColorProfile(ctx->decoder, &jxl_color);
 300         if (jret != JXL_DEC_SUCCESS) {
 301             av_log(avctx, AV_LOG_WARNING, "Unable to set fallback color encoding\n");
 302             /*
 303              * This should only happen if there's a non-XYB encoded image with custom primaries
 304              * embedded as enums and no embedded ICC Profile.
 305              * In this case, libjxl will synthesize an ICC Profile for us.
 306              */
 307             avctx->color_trc = AVCOL_TRC_UNSPECIFIED;
 308             avctx->color_primaries = AVCOL_PRI_UNSPECIFIED;
 309             if ((ret = libjxl_get_icc(avctx)) < 0)
 310                 return ret;
 311         }
 312     }
 313
 314     frame->color_trc = avctx->color_trc;
 315     frame->color_primaries = avctx->color_primaries;
 316     frame->colorspace = avctx->colorspace;
 317
 318     return 0;
 319 }
 320
 321 static int libjxl_decode_frame(AVCodecContext *avctx, AVFrame *frame, int *got_frame, AVPacket *avpkt)
 322 {
 323     LibJxlDecodeContext *ctx = avctx->priv_data;
 324     const uint8_t *buf = avpkt->data;
 325     size_t remaining = avpkt->size;
 326     JxlDecoderStatus jret;
 327     int ret;
 328     *got_frame = 0;
 329
 330     while (1) {
 331
 332         jret = JxlDecoderSetInput(ctx->decoder, buf, remaining);
 333
 334         if (jret == JXL_DEC_ERROR) {
 335             /* this should never happen here unless there's a bug in libjxl */
 336             av_log(avctx, AV_LOG_ERROR, "Unknown libjxl decode error\n");
 337             return AVERROR_EXTERNAL;
 338         }
 339
 340         jret = JxlDecoderProcessInput(ctx->decoder);
 341         /*
 342          * JxlDecoderReleaseInput returns the number
 343          * of bytes remaining to be read, rather than
 344          * the number of bytes that it did read
 345          */
 346         remaining = JxlDecoderReleaseInput(ctx->decoder);
 347         buf = avpkt->data + avpkt->size - remaining;
 348
 349         switch(jret) {
 350         case JXL_DEC_ERROR:
 351             av_log(avctx, AV_LOG_ERROR, "Unknown libjxl decode error\n");
 352             return AVERROR_INVALIDDATA;
 353         case JXL_DEC_NEED_MORE_INPUT:
 354             if (remaining == 0) {
 355                 av_log(avctx, AV_LOG_ERROR, "Unexpected end of JXL codestream\n");
 356                 return AVERROR_INVALIDDATA;
 357             }
 358             av_log(avctx, AV_LOG_DEBUG, "NEED_MORE_INPUT event emitted\n");
 359             continue;
 360         case JXL_DEC_BASIC_INFO:
 361             av_log(avctx, AV_LOG_DEBUG, "BASIC_INFO event emitted\n");
 362             if (JxlDecoderGetBasicInfo(ctx->decoder, &ctx->basic_info) != JXL_DEC_SUCCESS) {
 363                 /*
 364                  * this should never happen
 365                  * if it does it is likely a libjxl decoder bug
 366                  */
 367                 av_log(avctx, AV_LOG_ERROR, "Bad libjxl basic info event\n");
 368                 return AVERROR_EXTERNAL;
 369             }
 370             avctx->pix_fmt = libjxl_get_pix_fmt(avctx, &ctx->basic_info, &ctx->jxl_pixfmt);
 371             if (avctx->pix_fmt == AV_PIX_FMT_NONE) {
 372                 av_log(avctx, AV_LOG_ERROR, "Bad libjxl pixel format\n");
 373                 return AVERROR_EXTERNAL;
 374             }
 375             if ((ret = ff_set_dimensions(avctx, ctx->basic_info.xsize, ctx->basic_info.ysize)) < 0)
 376                 return ret;
 377             continue;
 378         case JXL_DEC_COLOR_ENCODING:
 379             av_log(avctx, AV_LOG_DEBUG, "COLOR_ENCODING event emitted\n");
 380             if ((ret = libjxl_color_encoding_event(avctx, frame)) < 0)
 381                 return ret;
 382             continue;
 383         case JXL_DEC_NEED_IMAGE_OUT_BUFFER:
 384             av_log(avctx, AV_LOG_DEBUG, "NEED_IMAGE_OUT_BUFFER event emitted\n");
 385             if ((ret = ff_get_buffer(avctx, frame, 0)) < 0)
 386                 return ret;
 387             ctx->jxl_pixfmt.align = frame->linesize[0];
 388             if (JxlDecoderSetImageOutBuffer(ctx->decoder, &ctx->jxl_pixfmt, frame->data[0], frame->buf[0]->size)
 389                 != JXL_DEC_SUCCESS) {
 390                 av_log(avctx, AV_LOG_ERROR, "Bad libjxl dec need image out buffer event\n");
 391                 return AVERROR_EXTERNAL;
 392             }
 393             continue;
 394         case JXL_DEC_FULL_IMAGE:
 395             /* full image is one frame, even if animated */
 396             av_log(avctx, AV_LOG_DEBUG, "FULL_IMAGE event emitted\n");
 397             frame->pict_type = AV_PICTURE_TYPE_I;
 398             frame->key_frame = 1;
 399             if (ctx->iccp) {
 400                 AVFrameSideData *sd = av_frame_new_side_data_from_buf(frame, AV_FRAME_DATA_ICC_PROFILE, ctx->iccp);
 401                 if (!sd)
 402                     return AVERROR(ENOMEM);
 403                 /* ownership is transfered, and it is not ref-ed */
 404                 ctx->iccp = NULL;
 405             }
 406             *got_frame = 1;
 407             return avpkt->size - remaining;
 408         case JXL_DEC_SUCCESS:
 409             av_log(avctx, AV_LOG_DEBUG, "SUCCESS event emitted\n");
 410             /*
 411              * The SUCCESS event isn't fired until after JXL_DEC_FULL_IMAGE. If this
 412              * stream only contains one JXL image then JXL_DEC_SUCCESS will never fire.
 413              * If the image2 sequence being decoded contains several JXL files, then
 414              * libjxl will fire this event after the next AVPacket has been passed,
 415              * which means the current packet is actually the next image in the sequence.
 416              * This is why we reset the decoder and populate the packet data now, since
 417              * this is the next packet and it has not been decoded yet. The decoder does
 418              * have to be reset to allow us to use it for the next image, or libjxl
 419              * will become very confused if the header information is not identical.
 420              */
 421             JxlDecoderReset(ctx->decoder);
 422             libjxl_init_jxl_decoder(avctx);
 423             buf = avpkt->data;
 424             remaining = avpkt->size;
 425             continue;
 426         default:
 427              av_log(avctx, AV_LOG_ERROR, "Bad libjxl event: %d\n", jret);
 428              return AVERROR_EXTERNAL;
 429         }
 430     }
 431 }
 432
 433 static av_cold int libjxl_decode_close(AVCodecContext *avctx)
 434 {
 435     LibJxlDecodeContext *ctx = avctx->priv_data;
 436
 437     if (ctx->runner)
 438         JxlThreadParallelRunnerDestroy(ctx->runner);
 439     ctx->runner = NULL;
 440     if (ctx->decoder)
 441         JxlDecoderDestroy(ctx->decoder);
 442     ctx->decoder = NULL;
 443     av_buffer_unref(&ctx->iccp);
 444
 445     return 0;
 446 }
 447
 448 const FFCodec ff_libjxl_decoder = {
 449     .p.name           = "libjxl",
 450     .p.long_name      = NULL_IF_CONFIG_SMALL("libjxl JPEG XL"),
 451     .p.type           = AVMEDIA_TYPE_VIDEO,
 452     .p.id             = AV_CODEC_ID_JPEGXL,
 453     .priv_data_size   = sizeof(LibJxlDecodeContext),
 454     .init             = libjxl_decode_init,
 455     FF_CODEC_DECODE_CB(libjxl_decode_frame),
 456     .close            = libjxl_decode_close,
 457     .p.capabilities   = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_OTHER_THREADS,
 458     .caps_internal    = FF_CODEC_CAP_AUTO_THREADS | FF_CODEC_CAP_INIT_CLEANUP,
 459     .p.wrapper_name   = "libjxl",
 460 };