libavcodec/libjxlenc.c - third_party/ffmpeg - Git at Google

 /*
  * JPEG XL encoding support via libjxl
  * Copyright (c) 2021 Leo Izen <leo.izen@gmail.com>
  *
  * This file is part of FFmpeg.
  *
  * FFmpeg is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
  * License as published by the Free Software Foundation; either
  * version 2.1 of the License, or (at your option) any later version.
  *
  * FFmpeg is distributed in the hope that it will be useful,
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  * Lesser General Public License for more details.
  *
  * You should have received a copy of the GNU Lesser General Public
  * License along with FFmpeg; if not, write to the Free Software
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */

 /**
  * @file
  * JPEG XL encoder using libjxl
  */

 #include <string.h>

 #include "libavutil/avutil.h"
 #include "libavutil/csp.h"
 #include "libavutil/error.h"
 #include "libavutil/frame.h"
 #include "libavutil/libm.h"
 #include "libavutil/mem.h"
 #include "libavutil/opt.h"
 #include "libavutil/pixdesc.h"
 #include "libavutil/pixfmt.h"
 #include "libavutil/version.h"

 #include "avcodec.h"
 #include "encode.h"
 #include "codec_internal.h"

 #include <jxl/encode.h>
 #include <jxl/thread_parallel_runner.h>
 #include "libjxl.h"

 typedef struct LibJxlEncodeContext {
     AVClass *class;
     void *runner;
     JxlEncoder *encoder;
     JxlEncoderFrameSettings *options;
     int effort;
     float distance;
     int modular;
     int xyb;
     uint8_t *buffer;
     size_t buffer_size;
 } LibJxlEncodeContext;

 /**
  * Map a quality setting for -qscale roughly from libjpeg
  * quality numbers to libjxl's butteraugli distance for
  * photographic content.
  *
  * Setting distance explicitly is preferred, but this will
  * allow qscale to be used as a fallback.
  *
  * This function is continuous and injective on [0, 100] which
  * makes it monotonic.
  *
  * @param  quality 0.0 to 100.0 quality setting, libjpeg quality
  * @return         Butteraugli distance between 0.0 and 15.0
  */
 static float quality_to_distance(float quality)
 {
     if (quality >= 100.0)
         return 0.0;
     else if (quality >= 90.0)
         return (100.0 - quality) * 0.10;
     else if (quality >= 30.0)
         return 0.1 + (100.0 - quality) * 0.09;
     else if (quality > 0.0)
         return 15.0 + (59.0 * quality - 4350.0) * quality / 9000.0;
     else
         return 15.0;
 }

 /**
  * Initalize the encoder on a per-frame basis. All of these need to be set
  * once each time the encoder is reset, which it must be each frame to make
  * the image2 muxer work.
  *
  * @return       0 upon success, negative on failure.
  */
 static int libjxl_init_jxl_encoder(AVCodecContext *avctx)
 {
     LibJxlEncodeContext *ctx = avctx->priv_data;

     /* reset the encoder every frame for image2 muxer */
     JxlEncoderReset(ctx->encoder);

     ctx->options = JxlEncoderFrameSettingsCreate(ctx->encoder, NULL);
     if (!ctx->options) {
         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoderOptions\n");
         return AVERROR_EXTERNAL;
     }

     /* This needs to be set each time the encoder is reset */
     if (JxlEncoderSetParallelRunner(ctx->encoder, JxlThreadParallelRunner, ctx->runner)
             != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to set JxlThreadParallelRunner\n");
         return AVERROR_EXTERNAL;
     }

     /* these shouldn't fail, libjxl bug notwithstanding */
     if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_EFFORT, ctx->effort)
             != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to set effort to: %d\n", ctx->effort);
         return AVERROR_EXTERNAL;
     }

     /* check for negative, our default */
     if (ctx->distance < 0.0) {
         /* use ffmpeg.c -q option if passed */
         if (avctx->flags & AV_CODEC_FLAG_QSCALE)
             ctx->distance = quality_to_distance((float)avctx->global_quality / FF_QP2LAMBDA);
         else
             /* default 1.0 matches cjxl */
             ctx->distance = 1.0;
     }

     /*
      * 0.01 is the minimum distance accepted for lossy
      * interpreting any positive value less than this as minimum
      */
     if (ctx->distance > 0.0 && ctx->distance < 0.01)
         ctx->distance = 0.01;
     if (JxlEncoderSetFrameDistance(ctx->options, ctx->distance) != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to set distance: %f\n", ctx->distance);
         return AVERROR_EXTERNAL;
     }

     /*
      * In theory the library should automatically enable modular if necessary,
      * but it appears it won't at the moment due to a bug. This will still
      * work even if that is patched.
      */
     if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_MODULAR,
             ctx->modular || ctx->distance <= 0.0 ? 1 : -1) != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to set modular\n");
         return AVERROR_EXTERNAL;
     }

     return 0;
 }

 /**
  * Global encoder initialization. This only needs to be run once,
  * not every frame.
  */
 static av_cold int libjxl_encode_init(AVCodecContext *avctx)
 {
     LibJxlEncodeContext *ctx = avctx->priv_data;
     JxlMemoryManager manager;

     ff_libjxl_init_memory_manager(&manager);
     ctx->encoder = JxlEncoderCreate(&manager);
     if (!ctx->encoder) {
         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoder\n");
         return AVERROR_EXTERNAL;
     }

     ctx->runner = JxlThreadParallelRunnerCreate(&manager, ff_libjxl_get_threadcount(avctx->thread_count));
     if (!ctx->runner) {
         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlThreadParallelRunner\n");
         return AVERROR_EXTERNAL;
     }

     ctx->buffer_size = 4096;
     ctx->buffer = av_realloc(NULL, ctx->buffer_size);

     if (!ctx->buffer) {
         av_log(avctx, AV_LOG_ERROR, "Could not allocate encoding buffer\n");
         return AVERROR(ENOMEM);
     }

     return 0;
 }

 /**
  * Populate a JxlColorEncoding with the given enum AVColorPrimaries.
  * @return < 0 upon failure, >= 0 upon success
  */
 static int libjxl_populate_primaries(void *avctx, JxlColorEncoding *jxl_color, enum AVColorPrimaries prm)
 {
     const AVColorPrimariesDesc *desc;

     switch (prm) {
     case AVCOL_PRI_BT709:
         jxl_color->primaries = JXL_PRIMARIES_SRGB;
         jxl_color->white_point = JXL_WHITE_POINT_D65;
         return 0;
     case AVCOL_PRI_BT2020:
         jxl_color->primaries = JXL_PRIMARIES_2100;
         jxl_color->white_point = JXL_WHITE_POINT_D65;
         return 0;
     case AVCOL_PRI_SMPTE431:
         jxl_color->primaries = JXL_PRIMARIES_P3;
         jxl_color->white_point = JXL_WHITE_POINT_DCI;
         return 0;
     case AVCOL_PRI_SMPTE432:
         jxl_color->primaries = JXL_PRIMARIES_P3;
         jxl_color->white_point = JXL_WHITE_POINT_D65;
         return 0;
     case AVCOL_PRI_UNSPECIFIED:
         av_log(avctx, AV_LOG_WARNING, "Unknown primaries, assuming BT.709/sRGB. Colors may be wrong.\n");
         jxl_color->primaries = JXL_PRIMARIES_SRGB;
         jxl_color->white_point = JXL_WHITE_POINT_D65;
         return 0;
     }

     desc = av_csp_primaries_desc_from_id(prm);
     if (!desc)
         return AVERROR(EINVAL);

     jxl_color->primaries = JXL_PRIMARIES_CUSTOM;
     jxl_color->white_point = JXL_WHITE_POINT_CUSTOM;

     jxl_color->primaries_red_xy[0] = av_q2d(desc->prim.r.x);
     jxl_color->primaries_red_xy[1] = av_q2d(desc->prim.r.y);
     jxl_color->primaries_green_xy[0] = av_q2d(desc->prim.g.x);
     jxl_color->primaries_green_xy[1] = av_q2d(desc->prim.g.y);
     jxl_color->primaries_blue_xy[0] = av_q2d(desc->prim.b.x);
     jxl_color->primaries_blue_xy[1] = av_q2d(desc->prim.b.y);
     jxl_color->white_point_xy[0] = av_q2d(desc->wp.x);
     jxl_color->white_point_xy[1] = av_q2d(desc->wp.y);

     return 0;
 }

 /**
  * Encode an entire frame. Currently animation, is not supported by
  * this encoder, so this will always reinitialize a new still image
  * and encode a one-frame image (for image2 and image2pipe).
  */
 static int libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFrame *frame, int *got_packet)
 {
     LibJxlEncodeContext *ctx = avctx->priv_data;
     AVFrameSideData *sd;
     const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(frame->format);
     JxlBasicInfo info;
     JxlColorEncoding jxl_color;
     JxlPixelFormat jxl_fmt;
     int bits_per_sample;
 #if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
     JxlBitDepth jxl_bit_depth;
 #endif
     JxlEncoderStatus jret;
     int ret;
     size_t available = ctx->buffer_size;
     size_t bytes_written = 0;
     uint8_t *next_out = ctx->buffer;
     const uint8_t *data;

     ret = libjxl_init_jxl_encoder(avctx);
     if (ret) {
         av_log(avctx, AV_LOG_ERROR, "Error frame-initializing JxlEncoder\n");
         return ret;
     }

     /* populate the basic info settings */
     JxlEncoderInitBasicInfo(&info);
     jxl_fmt.num_channels = pix_desc->nb_components;
     info.xsize = frame->width;
     info.ysize = frame->height;
     info.num_extra_channels = (jxl_fmt.num_channels + 1) % 2;
     info.num_color_channels = jxl_fmt.num_channels - info.num_extra_channels;
     bits_per_sample = av_get_bits_per_pixel(pix_desc) / jxl_fmt.num_channels;
     info.bits_per_sample = avctx->bits_per_raw_sample > 0 && !(pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT)
                            ? avctx->bits_per_raw_sample : bits_per_sample;
     info.alpha_bits = (info.num_extra_channels > 0) * info.bits_per_sample;
     if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
         info.exponent_bits_per_sample = info.bits_per_sample > 16 ? 8 : 5;
         info.alpha_exponent_bits = info.alpha_bits ? info.exponent_bits_per_sample : 0;
         jxl_fmt.data_type = info.bits_per_sample > 16 ? JXL_TYPE_FLOAT : JXL_TYPE_FLOAT16;
     } else {
         info.exponent_bits_per_sample = 0;
         info.alpha_exponent_bits = 0;
         jxl_fmt.data_type = info.bits_per_sample <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16;
     }

 #if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
     jxl_bit_depth.bits_per_sample = bits_per_sample;
     jxl_bit_depth.type = JXL_BIT_DEPTH_FROM_PIXEL_FORMAT;
     jxl_bit_depth.exponent_bits_per_sample = pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT ?
                                              info.exponent_bits_per_sample : 0;
 #endif

     /* JPEG XL format itself does not support limited range */
     if (avctx->color_range == AVCOL_RANGE_MPEG ||
         avctx->color_range == AVCOL_RANGE_UNSPECIFIED && frame->color_range == AVCOL_RANGE_MPEG)
         av_log(avctx, AV_LOG_WARNING, "This encoder does not support limited (tv) range, colors will be wrong!\n");
     else if (avctx->color_range != AVCOL_RANGE_JPEG && frame->color_range != AVCOL_RANGE_JPEG)
         av_log(avctx, AV_LOG_WARNING, "Unknown color range, assuming full (pc)\n");

     /* bitexact lossless requires there to be no XYB transform */
     info.uses_original_profile = ctx->distance == 0.0 || !ctx->xyb;
     info.orientation = frame->linesize[0] >= 0 ? JXL_ORIENT_IDENTITY : JXL_ORIENT_FLIP_VERTICAL;

     if (JxlEncoderSetBasicInfo(ctx->encoder, &info) != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to set JxlBasicInfo\n");
         return AVERROR_EXTERNAL;
     }

     /* rendering intent doesn't matter here
      * but libjxl will whine if we don't set it */
     jxl_color.rendering_intent = JXL_RENDERING_INTENT_RELATIVE;

     switch (frame->color_trc && frame->color_trc != AVCOL_TRC_UNSPECIFIED
             ? frame->color_trc : avctx->color_trc) {
     case AVCOL_TRC_BT709:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_709;
         break;
     case AVCOL_TRC_LINEAR:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
         break;
     case AVCOL_TRC_IEC61966_2_1:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
         break;
     case AVCOL_TRC_SMPTE428:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_DCI;
         break;
     case AVCOL_TRC_SMPTE2084:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_PQ;
         break;
     case AVCOL_TRC_ARIB_STD_B67:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_HLG;
         break;
     case AVCOL_TRC_GAMMA22:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
         jxl_color.gamma = 1/2.2f;
         break;
     case AVCOL_TRC_GAMMA28:
         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
         jxl_color.gamma = 1/2.8f;
         break;
     default:
         if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
             av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming Linear Light. Colors may be wrong.\n");
             jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
         } else {
             av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming IEC61966-2-1/sRGB. Colors may be wrong.\n");
             jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
         }
     }

     /* This should be implied to be honest
      * but a libjxl bug makes it fail otherwise */
     if (info.num_color_channels == 1)
         jxl_color.color_space = JXL_COLOR_SPACE_GRAY;
     else
         jxl_color.color_space = JXL_COLOR_SPACE_RGB;

     ret = libjxl_populate_primaries(avctx, &jxl_color,
             frame->color_primaries && frame->color_primaries != AVCOL_PRI_UNSPECIFIED
             ? frame->color_primaries : avctx->color_primaries);
     if (ret < 0)
         return ret;

     sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE);
     if (sd && sd->size && JxlEncoderSetICCProfile(ctx->encoder, sd->data, sd->size) != JXL_ENC_SUCCESS)
         av_log(avctx, AV_LOG_WARNING, "Could not set ICC Profile\n");
     if (JxlEncoderSetColorEncoding(ctx->encoder, &jxl_color) != JXL_ENC_SUCCESS)
         av_log(avctx, AV_LOG_WARNING, "Failed to set JxlColorEncoding\n");

 #if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
     if (JxlEncoderSetFrameBitDepth(ctx->options, &jxl_bit_depth) != JXL_ENC_SUCCESS)
         av_log(avctx, AV_LOG_WARNING, "Failed to set JxlBitDepth\n");
 #endif

     /* depending on basic info, level 10 might
      * be required instead of level 5 */
     if (JxlEncoderGetRequiredCodestreamLevel(ctx->encoder) > 5) {
         if (JxlEncoderSetCodestreamLevel(ctx->encoder, 10) != JXL_ENC_SUCCESS)
             av_log(avctx, AV_LOG_WARNING, "Could not increase codestream level\n");
     }

     jxl_fmt.endianness = JXL_NATIVE_ENDIAN;
     if (frame->linesize[0] >= 0) {
         jxl_fmt.align = frame->linesize[0];
         data = frame->data[0];
     } else {
         jxl_fmt.align = -frame->linesize[0];
         data = frame->data[0] + frame->linesize[0] * (info.ysize - 1);
     }

     if (JxlEncoderAddImageFrame(ctx->options, &jxl_fmt, data, jxl_fmt.align * info.ysize) != JXL_ENC_SUCCESS) {
         av_log(avctx, AV_LOG_ERROR, "Failed to add Image Frame\n");
         return AVERROR_EXTERNAL;
     }

     /*
      * Run this after the last frame in the image has been passed.
      * TODO support animation
      */
     JxlEncoderCloseInput(ctx->encoder);

     while (1) {
         jret = JxlEncoderProcessOutput(ctx->encoder, &next_out, &available);
         if (jret == JXL_ENC_ERROR) {
             av_log(avctx, AV_LOG_ERROR, "Unspecified libjxl error occurred\n");
             return AVERROR_EXTERNAL;
         }
         bytes_written = ctx->buffer_size - available;
         /* all data passed has been encoded */
         if (jret == JXL_ENC_SUCCESS)
             break;
         if (jret == JXL_ENC_NEED_MORE_OUTPUT) {
             /*
              * at the moment, libjxl has no way to
              * tell us how much space it actually needs
              * so we need to malloc loop
              */
             uint8_t *temp;
             size_t new_size = ctx->buffer_size * 2;
             temp = av_realloc(ctx->buffer, new_size);
             if (!temp)
                 return AVERROR(ENOMEM);
             ctx->buffer = temp;
             ctx->buffer_size = new_size;
             next_out = ctx->buffer + bytes_written;
             available = new_size - bytes_written;
             continue;
         }
         av_log(avctx, AV_LOG_ERROR, "Bad libjxl event: %d\n", jret);
         return AVERROR_EXTERNAL;
     }

     ret = ff_get_encode_buffer(avctx, pkt, bytes_written, 0);
     if (ret < 0)
         return ret;

     memcpy(pkt->data, ctx->buffer, bytes_written);
     *got_packet = 1;

     return 0;
 }

 static av_cold int libjxl_encode_close(AVCodecContext *avctx)
 {
     LibJxlEncodeContext *ctx = avctx->priv_data;

     if (ctx->runner)
         JxlThreadParallelRunnerDestroy(ctx->runner);
     ctx->runner = NULL;

     /*
      * destroying the encoder also frees
      * ctx->options so we don't need to
      */
     if (ctx->encoder)
         JxlEncoderDestroy(ctx->encoder);
     ctx->encoder = NULL;

     av_freep(&ctx->buffer);

     return 0;
 }

 #define OFFSET(x) offsetof(LibJxlEncodeContext, x)
 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM

 static const AVOption libjxl_encode_options[] = {
     { "effort",        "Encoding effort",                                  OFFSET(effort),     AV_OPT_TYPE_INT,    { .i64 =    7 },    1,     9, VE },
     { "distance",      "Maximum Butteraugli distance (quality setting, "
                         "lower = better, zero = lossless, default 1.0)",   OFFSET(distance),   AV_OPT_TYPE_FLOAT,  { .dbl = -1.0 }, -1.0,  15.0, VE },
     { "modular",       "Force modular mode",                               OFFSET(modular),    AV_OPT_TYPE_INT,    { .i64 =    0 },    0,     1, VE },
     { "xyb",           "Use XYB-encoding for lossy images",                OFFSET(xyb),
         AV_OPT_TYPE_INT,   { .i64 =    1 },    0,     1, VE },
     { NULL },
 };

 static const AVClass libjxl_encode_class = {
     .class_name = "libjxl",
     .item_name  = av_default_item_name,
     .option     = libjxl_encode_options,
     .version    = LIBAVUTIL_VERSION_INT,
 };

 const FFCodec ff_libjxl_encoder = {
     .p.name           = "libjxl",
     CODEC_LONG_NAME("libjxl JPEG XL"),
     .p.type           = AVMEDIA_TYPE_VIDEO,
     .p.id             = AV_CODEC_ID_JPEGXL,
     .priv_data_size   = sizeof(LibJxlEncodeContext),
     .init             = libjxl_encode_init,
     FF_CODEC_ENCODE_CB(libjxl_encode_frame),
     .close            = libjxl_encode_close,
     .p.capabilities   = AV_CODEC_CAP_OTHER_THREADS |
                         AV_CODEC_CAP_DR1 |
                         AV_CODEC_CAP_ENCODER_REORDERED_OPAQUE,
     .caps_internal    = FF_CODEC_CAP_NOT_INIT_THREADSAFE |
                         FF_CODEC_CAP_AUTO_THREADS | FF_CODEC_CAP_INIT_CLEANUP |
                         FF_CODEC_CAP_ICC_PROFILES,
     .p.pix_fmts       = (const enum AVPixelFormat[]) {
         AV_PIX_FMT_RGB24, AV_PIX_FMT_RGBA,
         AV_PIX_FMT_RGB48, AV_PIX_FMT_RGBA64,
         AV_PIX_FMT_RGBF32, AV_PIX_FMT_RGBAF32,
         AV_PIX_FMT_GRAY8, AV_PIX_FMT_YA8,
         AV_PIX_FMT_GRAY16, AV_PIX_FMT_YA16,
         AV_PIX_FMT_GRAYF32,
         AV_PIX_FMT_NONE
     },
     .p.priv_class     = &libjxl_encode_class,
     .p.wrapper_name   = "libjxl",
 };
	/*
	* JPEG XL encoding support via libjxl
	* Copyright (c) 2021 Leo Izen <leo.izen@gmail.com>
	*
	* This file is part of FFmpeg.
	*
	* FFmpeg is free software; you can redistribute it and/or
	* modify it under the terms of the GNU Lesser General Public
	* License as published by the Free Software Foundation; either
	* version 2.1 of the License, or (at your option) any later version.
	*
	* FFmpeg is distributed in the hope that it will be useful,
	* but WITHOUT ANY WARRANTY; without even the implied warranty of
	* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
	* Lesser General Public License for more details.
	*
	* You should have received a copy of the GNU Lesser General Public
	* License along with FFmpeg; if not, write to the Free Software
	* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
	*/

	/**
	* @file
	* JPEG XL encoder using libjxl
	*/

	#include <string.h>

	#include "libavutil/avutil.h"
	#include "libavutil/csp.h"
	#include "libavutil/error.h"
	#include "libavutil/frame.h"
	#include "libavutil/libm.h"
	#include "libavutil/mem.h"
	#include "libavutil/opt.h"
	#include "libavutil/pixdesc.h"
	#include "libavutil/pixfmt.h"
	#include "libavutil/version.h"

	#include "avcodec.h"
	#include "encode.h"
	#include "codec_internal.h"

	#include <jxl/encode.h>
	#include <jxl/thread_parallel_runner.h>
	#include "libjxl.h"

	typedef struct LibJxlEncodeContext {
	AVClass *class;
	void *runner;
	JxlEncoder *encoder;
	JxlEncoderFrameSettings *options;
	int effort;
	float distance;
	int modular;
	int xyb;
	uint8_t *buffer;
	size_t buffer_size;
	} LibJxlEncodeContext;

	/**
	* Map a quality setting for -qscale roughly from libjpeg
	* quality numbers to libjxl's butteraugli distance for
	* photographic content.
	*
	* Setting distance explicitly is preferred, but this will
	* allow qscale to be used as a fallback.
	*
	* This function is continuous and injective on [0, 100] which
	* makes it monotonic.
	*
	* @param quality 0.0 to 100.0 quality setting, libjpeg quality
	* @return Butteraugli distance between 0.0 and 15.0
	*/
	static float quality_to_distance(float quality)
	{
	if (quality >= 100.0)
	return 0.0;
	else if (quality >= 90.0)
	return (100.0 - quality) * 0.10;
	else if (quality >= 30.0)
	return 0.1 + (100.0 - quality) * 0.09;
	else if (quality > 0.0)
	return 15.0 + (59.0 * quality - 4350.0) * quality / 9000.0;
	else
	return 15.0;
	}

	/**
	* Initalize the encoder on a per-frame basis. All of these need to be set
	* once each time the encoder is reset, which it must be each frame to make
	* the image2 muxer work.
	*
	* @return 0 upon success, negative on failure.
	*/
	static int libjxl_init_jxl_encoder(AVCodecContext *avctx)
	{
	LibJxlEncodeContext *ctx = avctx->priv_data;

	/* reset the encoder every frame for image2 muxer */
	JxlEncoderReset(ctx->encoder);

	ctx->options = JxlEncoderFrameSettingsCreate(ctx->encoder, NULL);
	if (!ctx->options) {
	av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoderOptions\n");
	return AVERROR_EXTERNAL;
	}

	/* This needs to be set each time the encoder is reset */
	if (JxlEncoderSetParallelRunner(ctx->encoder, JxlThreadParallelRunner, ctx->runner)
	!= JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to set JxlThreadParallelRunner\n");
	return AVERROR_EXTERNAL;
	}

	/* these shouldn't fail, libjxl bug notwithstanding */
	if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_EFFORT, ctx->effort)
	!= JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to set effort to: %d\n", ctx->effort);
	return AVERROR_EXTERNAL;
	}

	/* check for negative, our default */
	if (ctx->distance < 0.0) {
	/* use ffmpeg.c -q option if passed */
	if (avctx->flags & AV_CODEC_FLAG_QSCALE)
	ctx->distance = quality_to_distance((float)avctx->global_quality / FF_QP2LAMBDA);
	else
	/* default 1.0 matches cjxl */
	ctx->distance = 1.0;
	}

	/*
	* 0.01 is the minimum distance accepted for lossy
	* interpreting any positive value less than this as minimum
	*/
	if (ctx->distance > 0.0 && ctx->distance < 0.01)
	ctx->distance = 0.01;
	if (JxlEncoderSetFrameDistance(ctx->options, ctx->distance) != JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to set distance: %f\n", ctx->distance);
	return AVERROR_EXTERNAL;
	}

	/*
	* In theory the library should automatically enable modular if necessary,
	* but it appears it won't at the moment due to a bug. This will still
	* work even if that is patched.
	*/
	if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_MODULAR,
	ctx->modular \|\| ctx->distance <= 0.0 ? 1 : -1) != JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to set modular\n");
	return AVERROR_EXTERNAL;
	}

	return 0;
	}

	/**
	* Global encoder initialization. This only needs to be run once,
	* not every frame.
	*/
	static av_cold int libjxl_encode_init(AVCodecContext *avctx)
	{
	LibJxlEncodeContext *ctx = avctx->priv_data;
	JxlMemoryManager manager;

	ff_libjxl_init_memory_manager(&manager);
	ctx->encoder = JxlEncoderCreate(&manager);
	if (!ctx->encoder) {
	av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoder\n");
	return AVERROR_EXTERNAL;
	}

	ctx->runner = JxlThreadParallelRunnerCreate(&manager, ff_libjxl_get_threadcount(avctx->thread_count));
	if (!ctx->runner) {
	av_log(avctx, AV_LOG_ERROR, "Failed to create JxlThreadParallelRunner\n");
	return AVERROR_EXTERNAL;
	}

	ctx->buffer_size = 4096;
	ctx->buffer = av_realloc(NULL, ctx->buffer_size);

	if (!ctx->buffer) {
	av_log(avctx, AV_LOG_ERROR, "Could not allocate encoding buffer\n");
	return AVERROR(ENOMEM);
	}

	return 0;
	}

	/**
	* Populate a JxlColorEncoding with the given enum AVColorPrimaries.
	* @return < 0 upon failure, >= 0 upon success
	*/
	static int libjxl_populate_primaries(void avctx, JxlColorEncoding jxl_color, enum AVColorPrimaries prm)
	{
	const AVColorPrimariesDesc *desc;

	switch (prm) {
	case AVCOL_PRI_BT709:
	jxl_color->primaries = JXL_PRIMARIES_SRGB;
	jxl_color->white_point = JXL_WHITE_POINT_D65;
	return 0;
	case AVCOL_PRI_BT2020:
	jxl_color->primaries = JXL_PRIMARIES_2100;
	jxl_color->white_point = JXL_WHITE_POINT_D65;
	return 0;
	case AVCOL_PRI_SMPTE431:
	jxl_color->primaries = JXL_PRIMARIES_P3;
	jxl_color->white_point = JXL_WHITE_POINT_DCI;
	return 0;
	case AVCOL_PRI_SMPTE432:
	jxl_color->primaries = JXL_PRIMARIES_P3;
	jxl_color->white_point = JXL_WHITE_POINT_D65;
	return 0;
	case AVCOL_PRI_UNSPECIFIED:
	av_log(avctx, AV_LOG_WARNING, "Unknown primaries, assuming BT.709/sRGB. Colors may be wrong.\n");
	jxl_color->primaries = JXL_PRIMARIES_SRGB;
	jxl_color->white_point = JXL_WHITE_POINT_D65;
	return 0;
	}

	desc = av_csp_primaries_desc_from_id(prm);
	if (!desc)
	return AVERROR(EINVAL);

	jxl_color->primaries = JXL_PRIMARIES_CUSTOM;
	jxl_color->white_point = JXL_WHITE_POINT_CUSTOM;

	jxl_color->primaries_red_xy[0] = av_q2d(desc->prim.r.x);
	jxl_color->primaries_red_xy[1] = av_q2d(desc->prim.r.y);
	jxl_color->primaries_green_xy[0] = av_q2d(desc->prim.g.x);
	jxl_color->primaries_green_xy[1] = av_q2d(desc->prim.g.y);
	jxl_color->primaries_blue_xy[0] = av_q2d(desc->prim.b.x);
	jxl_color->primaries_blue_xy[1] = av_q2d(desc->prim.b.y);
	jxl_color->white_point_xy[0] = av_q2d(desc->wp.x);
	jxl_color->white_point_xy[1] = av_q2d(desc->wp.y);

	return 0;
	}

	/**
	* Encode an entire frame. Currently animation, is not supported by
	* this encoder, so this will always reinitialize a new still image
	* and encode a one-frame image (for image2 and image2pipe).
	*/
	static int libjxl_encode_frame(AVCodecContext avctx, AVPacket pkt, const AVFrame frame, int got_packet)
	{
	LibJxlEncodeContext *ctx = avctx->priv_data;
	AVFrameSideData *sd;
	const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(frame->format);
	JxlBasicInfo info;
	JxlColorEncoding jxl_color;
	JxlPixelFormat jxl_fmt;
	int bits_per_sample;
	#if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
	JxlBitDepth jxl_bit_depth;
	#endif
	JxlEncoderStatus jret;
	int ret;
	size_t available = ctx->buffer_size;
	size_t bytes_written = 0;
	uint8_t *next_out = ctx->buffer;
	const uint8_t *data;

	ret = libjxl_init_jxl_encoder(avctx);
	if (ret) {
	av_log(avctx, AV_LOG_ERROR, "Error frame-initializing JxlEncoder\n");
	return ret;
	}

	/* populate the basic info settings */
	JxlEncoderInitBasicInfo(&info);
	jxl_fmt.num_channels = pix_desc->nb_components;
	info.xsize = frame->width;
	info.ysize = frame->height;
	info.num_extra_channels = (jxl_fmt.num_channels + 1) % 2;
	info.num_color_channels = jxl_fmt.num_channels - info.num_extra_channels;
	bits_per_sample = av_get_bits_per_pixel(pix_desc) / jxl_fmt.num_channels;
	info.bits_per_sample = avctx->bits_per_raw_sample > 0 && !(pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT)
	? avctx->bits_per_raw_sample : bits_per_sample;
	info.alpha_bits = (info.num_extra_channels > 0) * info.bits_per_sample;
	if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
	info.exponent_bits_per_sample = info.bits_per_sample > 16 ? 8 : 5;
	info.alpha_exponent_bits = info.alpha_bits ? info.exponent_bits_per_sample : 0;
	jxl_fmt.data_type = info.bits_per_sample > 16 ? JXL_TYPE_FLOAT : JXL_TYPE_FLOAT16;
	} else {
	info.exponent_bits_per_sample = 0;
	info.alpha_exponent_bits = 0;
	jxl_fmt.data_type = info.bits_per_sample <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16;
	}

	#if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
	jxl_bit_depth.bits_per_sample = bits_per_sample;
	jxl_bit_depth.type = JXL_BIT_DEPTH_FROM_PIXEL_FORMAT;
	jxl_bit_depth.exponent_bits_per_sample = pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT ?
	info.exponent_bits_per_sample : 0;
	#endif

	/* JPEG XL format itself does not support limited range */
	if (avctx->color_range == AVCOL_RANGE_MPEG \|\|
	avctx->color_range == AVCOL_RANGE_UNSPECIFIED && frame->color_range == AVCOL_RANGE_MPEG)
	av_log(avctx, AV_LOG_WARNING, "This encoder does not support limited (tv) range, colors will be wrong!\n");
	else if (avctx->color_range != AVCOL_RANGE_JPEG && frame->color_range != AVCOL_RANGE_JPEG)
	av_log(avctx, AV_LOG_WARNING, "Unknown color range, assuming full (pc)\n");

	/* bitexact lossless requires there to be no XYB transform */
	info.uses_original_profile = ctx->distance == 0.0 \|\| !ctx->xyb;
	info.orientation = frame->linesize[0] >= 0 ? JXL_ORIENT_IDENTITY : JXL_ORIENT_FLIP_VERTICAL;

	if (JxlEncoderSetBasicInfo(ctx->encoder, &info) != JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to set JxlBasicInfo\n");
	return AVERROR_EXTERNAL;
	}

	/* rendering intent doesn't matter here
	* but libjxl will whine if we don't set it */
	jxl_color.rendering_intent = JXL_RENDERING_INTENT_RELATIVE;

	switch (frame->color_trc && frame->color_trc != AVCOL_TRC_UNSPECIFIED
	? frame->color_trc : avctx->color_trc) {
	case AVCOL_TRC_BT709:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_709;
	break;
	case AVCOL_TRC_LINEAR:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
	break;
	case AVCOL_TRC_IEC61966_2_1:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
	break;
	case AVCOL_TRC_SMPTE428:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_DCI;
	break;
	case AVCOL_TRC_SMPTE2084:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_PQ;
	break;
	case AVCOL_TRC_ARIB_STD_B67:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_HLG;
	break;
	case AVCOL_TRC_GAMMA22:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
	jxl_color.gamma = 1/2.2f;
	break;
	case AVCOL_TRC_GAMMA28:
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
	jxl_color.gamma = 1/2.8f;
	break;
	default:
	if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
	av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming Linear Light. Colors may be wrong.\n");
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
	} else {
	av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming IEC61966-2-1/sRGB. Colors may be wrong.\n");
	jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
	}
	}

	/* This should be implied to be honest
	* but a libjxl bug makes it fail otherwise */
	if (info.num_color_channels == 1)
	jxl_color.color_space = JXL_COLOR_SPACE_GRAY;
	else
	jxl_color.color_space = JXL_COLOR_SPACE_RGB;

	ret = libjxl_populate_primaries(avctx, &jxl_color,
	frame->color_primaries && frame->color_primaries != AVCOL_PRI_UNSPECIFIED
	? frame->color_primaries : avctx->color_primaries);
	if (ret < 0)
	return ret;

	sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE);
	if (sd && sd->size && JxlEncoderSetICCProfile(ctx->encoder, sd->data, sd->size) != JXL_ENC_SUCCESS)
	av_log(avctx, AV_LOG_WARNING, "Could not set ICC Profile\n");
	if (JxlEncoderSetColorEncoding(ctx->encoder, &jxl_color) != JXL_ENC_SUCCESS)
	av_log(avctx, AV_LOG_WARNING, "Failed to set JxlColorEncoding\n");

	#if JPEGXL_NUMERIC_VERSION >= JPEGXL_COMPUTE_NUMERIC_VERSION(0, 8, 0)
	if (JxlEncoderSetFrameBitDepth(ctx->options, &jxl_bit_depth) != JXL_ENC_SUCCESS)
	av_log(avctx, AV_LOG_WARNING, "Failed to set JxlBitDepth\n");
	#endif

	/* depending on basic info, level 10 might
	* be required instead of level 5 */
	if (JxlEncoderGetRequiredCodestreamLevel(ctx->encoder) > 5) {
	if (JxlEncoderSetCodestreamLevel(ctx->encoder, 10) != JXL_ENC_SUCCESS)
	av_log(avctx, AV_LOG_WARNING, "Could not increase codestream level\n");
	}

	jxl_fmt.endianness = JXL_NATIVE_ENDIAN;
	if (frame->linesize[0] >= 0) {
	jxl_fmt.align = frame->linesize[0];
	data = frame->data[0];
	} else {
	jxl_fmt.align = -frame->linesize[0];
	data = frame->data[0] + frame->linesize[0] * (info.ysize - 1);
	}

	if (JxlEncoderAddImageFrame(ctx->options, &jxl_fmt, data, jxl_fmt.align * info.ysize) != JXL_ENC_SUCCESS) {
	av_log(avctx, AV_LOG_ERROR, "Failed to add Image Frame\n");
	return AVERROR_EXTERNAL;
	}

	/*
	* Run this after the last frame in the image has been passed.
	* TODO support animation
	*/
	JxlEncoderCloseInput(ctx->encoder);

	while (1) {
	jret = JxlEncoderProcessOutput(ctx->encoder, &next_out, &available);
	if (jret == JXL_ENC_ERROR) {
	av_log(avctx, AV_LOG_ERROR, "Unspecified libjxl error occurred\n");
	return AVERROR_EXTERNAL;
	}
	bytes_written = ctx->buffer_size - available;
	/* all data passed has been encoded */
	if (jret == JXL_ENC_SUCCESS)
	break;
	if (jret == JXL_ENC_NEED_MORE_OUTPUT) {
	/*
	* at the moment, libjxl has no way to
	* tell us how much space it actually needs
	* so we need to malloc loop
	*/
	uint8_t *temp;
	size_t new_size = ctx->buffer_size * 2;
	temp = av_realloc(ctx->buffer, new_size);
	if (!temp)
	return AVERROR(ENOMEM);
	ctx->buffer = temp;
	ctx->buffer_size = new_size;
	next_out = ctx->buffer + bytes_written;
	available = new_size - bytes_written;
	continue;
	}
	av_log(avctx, AV_LOG_ERROR, "Bad libjxl event: %d\n", jret);
	return AVERROR_EXTERNAL;
	}

	ret = ff_get_encode_buffer(avctx, pkt, bytes_written, 0);
	if (ret < 0)
	return ret;

	memcpy(pkt->data, ctx->buffer, bytes_written);
	*got_packet = 1;

	return 0;
	}

	static av_cold int libjxl_encode_close(AVCodecContext *avctx)
	{
	LibJxlEncodeContext *ctx = avctx->priv_data;

	if (ctx->runner)
	JxlThreadParallelRunnerDestroy(ctx->runner);
	ctx->runner = NULL;

	/*
	* destroying the encoder also frees
	* ctx->options so we don't need to
	*/
	if (ctx->encoder)
	JxlEncoderDestroy(ctx->encoder);
	ctx->encoder = NULL;

	av_freep(&ctx->buffer);

	return 0;
	}

	#define OFFSET(x) offsetof(LibJxlEncodeContext, x)
	#define VE AV_OPT_FLAG_VIDEO_PARAM \| AV_OPT_FLAG_ENCODING_PARAM

	static const AVOption libjxl_encode_options[] = {
	{ "effort", "Encoding effort", OFFSET(effort), AV_OPT_TYPE_INT, { .i64 = 7 }, 1, 9, VE },
	{ "distance", "Maximum Butteraugli distance (quality setting, "
	"lower = better, zero = lossless, default 1.0)", OFFSET(distance), AV_OPT_TYPE_FLOAT, { .dbl = -1.0 }, -1.0, 15.0, VE },
	{ "modular", "Force modular mode", OFFSET(modular), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, VE },
	{ "xyb", "Use XYB-encoding for lossy images", OFFSET(xyb),
	AV_OPT_TYPE_INT, { .i64 = 1 }, 0, 1, VE },
	{ NULL },
	};

	static const AVClass libjxl_encode_class = {
	.class_name = "libjxl",
	.item_name = av_default_item_name,
	.option = libjxl_encode_options,
	.version = LIBAVUTIL_VERSION_INT,
	};

	const FFCodec ff_libjxl_encoder = {
	.p.name = "libjxl",
	CODEC_LONG_NAME("libjxl JPEG XL"),
	.p.type = AVMEDIA_TYPE_VIDEO,
	.p.id = AV_CODEC_ID_JPEGXL,
	.priv_data_size = sizeof(LibJxlEncodeContext),
	.init = libjxl_encode_init,
	FF_CODEC_ENCODE_CB(libjxl_encode_frame),
	.close = libjxl_encode_close,
	.p.capabilities = AV_CODEC_CAP_OTHER_THREADS \|
	AV_CODEC_CAP_DR1 \|
	AV_CODEC_CAP_ENCODER_REORDERED_OPAQUE,
	.caps_internal = FF_CODEC_CAP_NOT_INIT_THREADSAFE \|
	FF_CODEC_CAP_AUTO_THREADS \| FF_CODEC_CAP_INIT_CLEANUP \|
	FF_CODEC_CAP_ICC_PROFILES,
	.p.pix_fmts = (const enum AVPixelFormat[]) {
	AV_PIX_FMT_RGB24, AV_PIX_FMT_RGBA,
	AV_PIX_FMT_RGB48, AV_PIX_FMT_RGBA64,
	AV_PIX_FMT_RGBF32, AV_PIX_FMT_RGBAF32,
	AV_PIX_FMT_GRAY8, AV_PIX_FMT_YA8,
	AV_PIX_FMT_GRAY16, AV_PIX_FMT_YA16,
	AV_PIX_FMT_GRAYF32,
	AV_PIX_FMT_NONE
	},
	.p.priv_class = &libjxl_encode_class,
	.p.wrapper_name = "libjxl",
	};