1 /*
2  * JPEG XL encoding 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 encoder using libjxl
25  */
26 
27 #include <string.h>
28 
29 #include "libavutil/avutil.h"
30 #include "libavutil/csp.h"
31 #include "libavutil/error.h"
32 #include "libavutil/frame.h"
33 #include "libavutil/libm.h"
34 #include "libavutil/opt.h"
35 #include "libavutil/pixdesc.h"
36 #include "libavutil/pixfmt.h"
37 #include "libavutil/version.h"
38 
39 #include "avcodec.h"
40 #include "encode.h"
41 #include "codec_internal.h"
42 
43 #include <jxl/encode.h>
44 #include <jxl/thread_parallel_runner.h>
45 #include "libjxl.h"
46 
47 typedef struct LibJxlEncodeContext {
48     AVClass *class;
49     void *runner;
50     JxlEncoder *encoder;
51     JxlEncoderFrameSettings *options;
52     int effort;
53     float distance;
54     int modular;
55     uint8_t *buffer;
56     size_t buffer_size;
57 } LibJxlEncodeContext;
58 
59 /**
60  * Map a quality setting for -qscale roughly from libjpeg
61  * quality numbers to libjxl's butteraugli distance for
62  * photographic content.
63  *
64  * Setting distance explicitly is preferred, but this will
65  * allow qscale to be used as a fallback.
66  *
67  * This function is continuous and injective on [0, 100] which
68  * makes it monotonic.
69  *
70  * @param  quality 0.0 to 100.0 quality setting, libjpeg quality
71  * @return         Butteraugli distance between 0.0 and 15.0
72  */
quality_to_distance(float quality)73 static float quality_to_distance(float quality)
74 {
75     if (quality >= 100.0)
76         return 0.0;
77     else if (quality >= 90.0)
78         return (100.0 - quality) * 0.10;
79     else if (quality >= 30.0)
80         return 0.1 + (100.0 - quality) * 0.09;
81     else if (quality > 0.0)
82         return 15.0 + (59.0 * quality - 4350.0) * quality / 9000.0;
83     else
84         return 15.0;
85 }
86 
87 /**
88  * Initalize the decoder on a per-frame basis. All of these need to be set
89  * once each time the decoder is reset, which it must be each frame to make
90  * the image2 muxer work.
91  *
92  * @return       0 upon success, negative on failure.
93  */
libjxl_init_jxl_encoder(AVCodecContext *avctx)94 static int libjxl_init_jxl_encoder(AVCodecContext *avctx)
95 {
96     LibJxlEncodeContext *ctx = avctx->priv_data;
97 
98     /* reset the encoder every frame for image2 muxer */
99     JxlEncoderReset(ctx->encoder);
100 
101     ctx->options = JxlEncoderFrameSettingsCreate(ctx->encoder, NULL);
102     if (!ctx->options) {
103         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoderOptions\n");
104         return AVERROR_EXTERNAL;
105     }
106 
107     /* This needs to be set each time the decoder is reset */
108     if (JxlEncoderSetParallelRunner(ctx->encoder, JxlThreadParallelRunner, ctx->runner)
109             != JXL_ENC_SUCCESS) {
110         av_log(avctx, AV_LOG_ERROR, "Failed to set JxlThreadParallelRunner\n");
111         return AVERROR_EXTERNAL;
112     }
113 
114     /* these shouldn't fail, libjxl bug notwithstanding */
115     if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_EFFORT, ctx->effort)
116             != JXL_ENC_SUCCESS) {
117         av_log(avctx, AV_LOG_ERROR, "Failed to set effort to: %d\n", ctx->effort);
118         return AVERROR_EXTERNAL;
119     }
120 
121     /* check for negative, our default */
122     if (ctx->distance < 0.0) {
123         /* use ffmpeg.c -q option if passed */
124         if (avctx->flags & AV_CODEC_FLAG_QSCALE)
125             ctx->distance = quality_to_distance((float)avctx->global_quality / FF_QP2LAMBDA);
126         else
127             /* default 1.0 matches cjxl */
128             ctx->distance = 1.0;
129     }
130 
131     /*
132      * 0.01 is the minimum distance accepted for lossy
133      * interpreting any positive value less than this as minimum
134      */
135     if (ctx->distance > 0.0 && ctx->distance < 0.01)
136         ctx->distance = 0.01;
137     if (JxlEncoderSetFrameDistance(ctx->options, ctx->distance) != JXL_ENC_SUCCESS) {
138         av_log(avctx, AV_LOG_ERROR, "Failed to set distance: %f\n", ctx->distance);
139         return AVERROR_EXTERNAL;
140     }
141 
142     /*
143      * In theory the library should automatically enable modular if necessary,
144      * but it appears it won't at the moment due to a bug. This will still
145      * work even if that is patched.
146      */
147     if (JxlEncoderFrameSettingsSetOption(ctx->options, JXL_ENC_FRAME_SETTING_MODULAR,
148             ctx->modular || ctx->distance <= 0.0 ? 1 : -1) != JXL_ENC_SUCCESS) {
149         av_log(avctx, AV_LOG_ERROR, "Failed to set modular\n");
150         return AVERROR_EXTERNAL;
151     }
152 
153     return 0;
154 }
155 
156 /**
157  * Global encoder initialization. This only needs to be run once,
158  * not every frame.
159  */
libjxl_encode_init(AVCodecContext *avctx)160 static av_cold int libjxl_encode_init(AVCodecContext *avctx)
161 {
162     LibJxlEncodeContext *ctx = avctx->priv_data;
163     JxlMemoryManager manager;
164 
165     ff_libjxl_init_memory_manager(&manager);
166     ctx->encoder = JxlEncoderCreate(&manager);
167     if (!ctx->encoder) {
168         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlEncoder\n");
169         return AVERROR_EXTERNAL;
170     }
171 
172     ctx->runner = JxlThreadParallelRunnerCreate(&manager, ff_libjxl_get_threadcount(avctx->thread_count));
173     if (!ctx->runner) {
174         av_log(avctx, AV_LOG_ERROR, "Failed to create JxlThreadParallelRunner\n");
175         return AVERROR_EXTERNAL;
176     }
177 
178     ctx->buffer_size = 4096;
179     ctx->buffer = av_realloc(NULL, ctx->buffer_size);
180 
181     if (!ctx->buffer) {
182         av_log(avctx, AV_LOG_ERROR, "Could not allocate encoding buffer\n");
183         return AVERROR(ENOMEM);
184     }
185 
186     return 0;
187 }
188 
189 /**
190  * Populate a JxlColorEncoding with the given enum AVColorPrimaries.
191  * @return < 0 upon failure, >= 0 upon success
192  */
libjxl_populate_primaries(void *avctx, JxlColorEncoding *jxl_color, enum AVColorPrimaries prm)193 static int libjxl_populate_primaries(void *avctx, JxlColorEncoding *jxl_color, enum AVColorPrimaries prm)
194 {
195     const AVColorPrimariesDesc *desc;
196 
197     switch (prm) {
198     case AVCOL_PRI_BT709:
199         jxl_color->primaries = JXL_PRIMARIES_SRGB;
200         jxl_color->white_point = JXL_WHITE_POINT_D65;
201         return 0;
202     case AVCOL_PRI_BT2020:
203         jxl_color->primaries = JXL_PRIMARIES_2100;
204         jxl_color->white_point = JXL_WHITE_POINT_D65;
205         return 0;
206     case AVCOL_PRI_SMPTE431:
207         jxl_color->primaries = JXL_PRIMARIES_P3;
208         jxl_color->white_point = JXL_WHITE_POINT_DCI;
209         return 0;
210     case AVCOL_PRI_SMPTE432:
211         jxl_color->primaries = JXL_PRIMARIES_P3;
212         jxl_color->white_point = JXL_WHITE_POINT_D65;
213         return 0;
214     case AVCOL_PRI_UNSPECIFIED:
215         av_log(avctx, AV_LOG_WARNING, "Unknown primaries, assuming BT.709/sRGB. Colors may be wrong.\n");
216         jxl_color->primaries = JXL_PRIMARIES_SRGB;
217         jxl_color->white_point = JXL_WHITE_POINT_D65;
218         return 0;
219     }
220 
221     desc = av_csp_primaries_desc_from_id(prm);
222     if (!desc)
223         return AVERROR(EINVAL);
224 
225     jxl_color->primaries = JXL_PRIMARIES_CUSTOM;
226     jxl_color->white_point = JXL_WHITE_POINT_CUSTOM;
227 
228     jxl_color->primaries_red_xy[0] = av_q2d(desc->prim.r.x);
229     jxl_color->primaries_red_xy[1] = av_q2d(desc->prim.r.y);
230     jxl_color->primaries_green_xy[0] = av_q2d(desc->prim.g.x);
231     jxl_color->primaries_green_xy[1] = av_q2d(desc->prim.g.y);
232     jxl_color->primaries_blue_xy[0] = av_q2d(desc->prim.b.x);
233     jxl_color->primaries_blue_xy[1] = av_q2d(desc->prim.b.y);
234     jxl_color->white_point_xy[0] = av_q2d(desc->wp.x);
235     jxl_color->white_point_xy[1] = av_q2d(desc->wp.y);
236 
237     return 0;
238 }
239 
240 /**
241  * Encode an entire frame. Currently animation, is not supported by
242  * this encoder, so this will always reinitialize a new still image
243  * and encode a one-frame image (for image2 and image2pipe).
244  */
libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFrame *frame, int *got_packet)245 static int libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFrame *frame, int *got_packet)
246 {
247     LibJxlEncodeContext *ctx = avctx->priv_data;
248     AVFrameSideData *sd;
249     const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(frame->format);
250     JxlBasicInfo info;
251     JxlColorEncoding jxl_color;
252     JxlPixelFormat jxl_fmt;
253     JxlEncoderStatus jret;
254     int ret;
255     size_t available = ctx->buffer_size;
256     size_t bytes_written = 0;
257     uint8_t *next_out = ctx->buffer;
258 
259     ret = libjxl_init_jxl_encoder(avctx);
260     if (ret) {
261         av_log(avctx, AV_LOG_ERROR, "Error frame-initializing JxlEncoder\n");
262         return ret;
263     }
264 
265     /* populate the basic info settings */
266     JxlEncoderInitBasicInfo(&info);
267     jxl_fmt.num_channels = pix_desc->nb_components;
268     info.xsize = frame->width;
269     info.ysize = frame->height;
270     info.num_extra_channels = (jxl_fmt.num_channels + 1) % 2;
271     info.num_color_channels = jxl_fmt.num_channels - info.num_extra_channels;
272     info.bits_per_sample = av_get_bits_per_pixel(pix_desc) / jxl_fmt.num_channels;
273     info.alpha_bits = (info.num_extra_channels > 0) * info.bits_per_sample;
274     if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
275         info.exponent_bits_per_sample = info.bits_per_sample > 16 ? 8 : 5;
276         info.alpha_exponent_bits = info.alpha_bits ? info.exponent_bits_per_sample : 0;
277         jxl_fmt.data_type = info.bits_per_sample > 16 ? JXL_TYPE_FLOAT : JXL_TYPE_FLOAT16;
278     } else {
279         info.exponent_bits_per_sample = 0;
280         info.alpha_exponent_bits = 0;
281         jxl_fmt.data_type = info.bits_per_sample <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16;
282     }
283 
284     /* JPEG XL format itself does not support limited range */
285     if (avctx->color_range == AVCOL_RANGE_MPEG ||
286         avctx->color_range == AVCOL_RANGE_UNSPECIFIED && frame->color_range == AVCOL_RANGE_MPEG)
287         av_log(avctx, AV_LOG_WARNING, "This encoder does not support limited (tv) range, colors will be wrong!\n");
288     else if (avctx->color_range != AVCOL_RANGE_JPEG && frame->color_range != AVCOL_RANGE_JPEG)
289         av_log(avctx, AV_LOG_WARNING, "Unknown color range, assuming full (pc)\n");
290 
291     /* bitexact lossless requires there to be no XYB transform */
292     info.uses_original_profile = ctx->distance == 0.0;
293 
294     if (JxlEncoderSetBasicInfo(ctx->encoder, &info) != JXL_ENC_SUCCESS) {
295         av_log(avctx, AV_LOG_ERROR, "Failed to set JxlBasicInfo\n");
296         return AVERROR_EXTERNAL;
297     }
298 
299     /* rendering intent doesn't matter here
300      * but libjxl will whine if we don't set it */
301     jxl_color.rendering_intent = JXL_RENDERING_INTENT_RELATIVE;
302 
303     switch (frame->color_trc && frame->color_trc != AVCOL_TRC_UNSPECIFIED
304             ? frame->color_trc : avctx->color_trc) {
305     case AVCOL_TRC_BT709:
306         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_709;
307         break;
308     case AVCOL_TRC_LINEAR:
309         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
310         break;
311     case AVCOL_TRC_IEC61966_2_1:
312         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
313         break;
314     case AVCOL_TRC_SMPTE428:
315         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_DCI;
316         break;
317     case AVCOL_TRC_SMPTE2084:
318         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_PQ;
319         break;
320     case AVCOL_TRC_ARIB_STD_B67:
321         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_HLG;
322         break;
323     case AVCOL_TRC_GAMMA22:
324         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
325         jxl_color.gamma = 2.2;
326         break;
327     case AVCOL_TRC_GAMMA28:
328         jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_GAMMA;
329         jxl_color.gamma = 2.8;
330         break;
331     default:
332         if (pix_desc->flags & AV_PIX_FMT_FLAG_FLOAT) {
333             av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming Linear Light. Colors may be wrong.\n");
334             jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
335         } else {
336             av_log(avctx, AV_LOG_WARNING, "Unknown transfer function, assuming IEC61966-2-1/sRGB. Colors may be wrong.\n");
337             jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
338         }
339     }
340 
341     /* This should be implied to be honest
342      * but a libjxl bug makes it fail otherwise */
343     if (info.num_color_channels == 1)
344         jxl_color.color_space = JXL_COLOR_SPACE_GRAY;
345     else
346         jxl_color.color_space = JXL_COLOR_SPACE_RGB;
347 
348     ret = libjxl_populate_primaries(avctx, &jxl_color,
349             frame->color_primaries && frame->color_primaries != AVCOL_PRI_UNSPECIFIED
350             ? frame->color_primaries : avctx->color_primaries);
351     if (ret < 0)
352         return ret;
353 
354     sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE);
355     if (sd && sd->size && JxlEncoderSetICCProfile(ctx->encoder, sd->data, sd->size) != JXL_ENC_SUCCESS)
356         av_log(avctx, AV_LOG_WARNING, "Could not set ICC Profile\n");
357     if (JxlEncoderSetColorEncoding(ctx->encoder, &jxl_color) != JXL_ENC_SUCCESS)
358         av_log(avctx, AV_LOG_WARNING, "Failed to set JxlColorEncoding\n");
359 
360     /* depending on basic info, level 10 might
361      * be required instead of level 5 */
362     if (JxlEncoderGetRequiredCodestreamLevel(ctx->encoder) > 5) {
363         if (JxlEncoderSetCodestreamLevel(ctx->encoder, 10) != JXL_ENC_SUCCESS)
364             av_log(avctx, AV_LOG_WARNING, "Could not increase codestream level\n");
365     }
366 
367     jxl_fmt.endianness = JXL_NATIVE_ENDIAN;
368     jxl_fmt.align = frame->linesize[0];
369 
370     if (JxlEncoderAddImageFrame(ctx->options, &jxl_fmt, frame->data[0], jxl_fmt.align * info.ysize) != JXL_ENC_SUCCESS) {
371         av_log(avctx, AV_LOG_ERROR, "Failed to add Image Frame\n");
372         return AVERROR_EXTERNAL;
373     }
374 
375     /*
376      * Run this after the last frame in the image has been passed.
377      * TODO support animation
378      */
379     JxlEncoderCloseInput(ctx->encoder);
380 
381     while (1) {
382         jret = JxlEncoderProcessOutput(ctx->encoder, &next_out, &available);
383         if (jret == JXL_ENC_ERROR) {
384             av_log(avctx, AV_LOG_ERROR, "Unspecified libjxl error occurred\n");
385             return AVERROR_EXTERNAL;
386         }
387         bytes_written = ctx->buffer_size - available;
388         /* all data passed has been encoded */
389         if (jret == JXL_ENC_SUCCESS)
390             break;
391         if (jret == JXL_ENC_NEED_MORE_OUTPUT) {
392             /*
393              * at the moment, libjxl has no way to
394              * tell us how much space it actually needs
395              * so we need to malloc loop
396              */
397             uint8_t *temp;
398             size_t new_size = ctx->buffer_size * 2;
399             temp = av_realloc(ctx->buffer, new_size);
400             if (!temp)
401                 return AVERROR(ENOMEM);
402             ctx->buffer = temp;
403             ctx->buffer_size = new_size;
404             next_out = ctx->buffer + bytes_written;
405             available = new_size - bytes_written;
406             continue;
407         }
408         av_log(avctx, AV_LOG_ERROR, "Bad libjxl event: %d\n", jret);
409         return AVERROR_EXTERNAL;
410     }
411 
412     ret = ff_get_encode_buffer(avctx, pkt, bytes_written, 0);
413     if (ret < 0)
414         return ret;
415 
416     memcpy(pkt->data, ctx->buffer, bytes_written);
417     *got_packet = 1;
418 
419     return 0;
420 }
421 
libjxl_encode_close(AVCodecContext *avctx)422 static av_cold int libjxl_encode_close(AVCodecContext *avctx)
423 {
424     LibJxlEncodeContext *ctx = avctx->priv_data;
425 
426     if (ctx->runner)
427         JxlThreadParallelRunnerDestroy(ctx->runner);
428     ctx->runner = NULL;
429 
430     /*
431      * destroying the decoder also frees
432      * ctx->options so we don't need to
433      */
434     if (ctx->encoder)
435         JxlEncoderDestroy(ctx->encoder);
436     ctx->encoder = NULL;
437 
438     av_freep(&ctx->buffer);
439 
440     return 0;
441 }
442 
443 #define OFFSET(x) offsetof(LibJxlEncodeContext, x)
444 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
445 
446 static const AVOption libjxl_encode_options[] = {
447     { "effort",        "Encoding effort",                                  OFFSET(effort),     AV_OPT_TYPE_INT,    { .i64 =    7 },    1,     9, VE },
448     { "distance",      "Maximum Butteraugli distance (quality setting, "
449                         "lower = better, zero = lossless, default 1.0)",   OFFSET(distance),   AV_OPT_TYPE_FLOAT,  { .dbl = -1.0 }, -1.0,  15.0, VE },
450     { "modular",       "Force modular mode",                               OFFSET(modular),    AV_OPT_TYPE_INT,    { .i64 =    0 },    0,     1, VE },
451     { NULL },
452 };
453 
454 static const AVClass libjxl_encode_class = {
455     .class_name = "libjxl",
456     .item_name  = av_default_item_name,
457     .option     = libjxl_encode_options,
458     .version    = LIBAVUTIL_VERSION_INT,
459 };
460 
461 const FFCodec ff_libjxl_encoder = {
462     .p.name           = "libjxl",
463     .p.long_name      = NULL_IF_CONFIG_SMALL("libjxl JPEG XL"),
464     .p.type           = AVMEDIA_TYPE_VIDEO,
465     .p.id             = AV_CODEC_ID_JPEGXL,
466     .priv_data_size   = sizeof(LibJxlEncodeContext),
467     .init             = libjxl_encode_init,
468     FF_CODEC_ENCODE_CB(libjxl_encode_frame),
469     .close            = libjxl_encode_close,
470     .p.capabilities   = AV_CODEC_CAP_OTHER_THREADS,
471     .caps_internal    = FF_CODEC_CAP_AUTO_THREADS | FF_CODEC_CAP_INIT_CLEANUP,
472     .p.pix_fmts       = (const enum AVPixelFormat[]) {
473         AV_PIX_FMT_RGB24, AV_PIX_FMT_RGBA,
474         AV_PIX_FMT_RGB48, AV_PIX_FMT_RGBA64,
475         AV_PIX_FMT_GRAY8, AV_PIX_FMT_YA8,
476         AV_PIX_FMT_GRAY16, AV_PIX_FMT_YA16,
477         AV_PIX_FMT_GRAYF32,
478         AV_PIX_FMT_NONE
479     },
480     .p.priv_class     = &libjxl_encode_class,
481     .p.wrapper_name   = "libjxl",
482 };
483