aboutsummaryrefslogtreecommitdiffstats
path: root/libavcodec/libjxldec.c
blob: de48bea4b2ce02b4553ca6255d2f96abf4d513ff (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
/*
 * JPEG XL decoding 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 decoder using libjxl
 */

#include "libavutil/avassert.h"
#include "libavutil/buffer.h"
#include "libavutil/common.h"
#include "libavutil/csp.h"
#include "libavutil/error.h"
#include "libavutil/mem.h"
#include "libavutil/pixdesc.h"
#include "libavutil/pixfmt.h"
#include "libavutil/frame.h"

#include "avcodec.h"
#include "codec_internal.h"
#include "decode.h"

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

typedef struct LibJxlDecodeContext {
    void *runner;
    JxlDecoder *decoder;
    JxlBasicInfo basic_info;
    JxlPixelFormat jxl_pixfmt;
    JxlDecoderStatus events;
    AVBufferRef *iccp;
} LibJxlDecodeContext;

static int libjxl_init_jxl_decoder(AVCodecContext *avctx)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;

    ctx->events = JXL_DEC_BASIC_INFO | JXL_DEC_FULL_IMAGE | JXL_DEC_COLOR_ENCODING;
    if (JxlDecoderSubscribeEvents(ctx->decoder, ctx->events) != JXL_DEC_SUCCESS) {
        av_log(avctx, AV_LOG_ERROR, "Error subscribing to JXL events\n");
        return AVERROR_EXTERNAL;
    }

    if (JxlDecoderSetParallelRunner(ctx->decoder, JxlThreadParallelRunner, ctx->runner) != JXL_DEC_SUCCESS) {
        av_log(avctx, AV_LOG_ERROR, "Failed to set JxlThreadParallelRunner\n");
        return AVERROR_EXTERNAL;
    }

    memset(&ctx->basic_info, 0, sizeof(JxlBasicInfo));
    memset(&ctx->jxl_pixfmt, 0, sizeof(JxlPixelFormat));

    return 0;
}

static av_cold int libjxl_decode_init(AVCodecContext *avctx)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;
    JxlMemoryManager manager;

    ff_libjxl_init_memory_manager(&manager);
    ctx->decoder = JxlDecoderCreate(&manager);
    if (!ctx->decoder) {
        av_log(avctx, AV_LOG_ERROR, "Failed to create JxlDecoder\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;
    }

    return libjxl_init_jxl_decoder(avctx);
}

static enum AVPixelFormat libjxl_get_pix_fmt(void *avctx, const JxlBasicInfo *basic_info, JxlPixelFormat *format)
{
    format->endianness = JXL_NATIVE_ENDIAN;
    format->num_channels = basic_info->num_color_channels + (basic_info->alpha_bits > 0);
    /* Gray */
    if (basic_info->num_color_channels == 1) {
        if (basic_info->bits_per_sample <= 8) {
            format->data_type = JXL_TYPE_UINT8;
            return basic_info->alpha_bits ? AV_PIX_FMT_YA8 : AV_PIX_FMT_GRAY8;
        }
        if (basic_info->exponent_bits_per_sample || basic_info->bits_per_sample > 16) {
            if (basic_info->alpha_bits)
                return AV_PIX_FMT_NONE;
            format->data_type = JXL_TYPE_FLOAT;
            return AV_PIX_FMT_GRAYF32;
        }
        format->data_type = JXL_TYPE_UINT16;
        return basic_info->alpha_bits ? AV_PIX_FMT_YA16 : AV_PIX_FMT_GRAY16;
    }
    /* rgb only */
    /* libjxl only supports packed RGB and gray output at the moment */
    if (basic_info->num_color_channels == 3) {
        if (basic_info->bits_per_sample <= 8) {
            format->data_type = JXL_TYPE_UINT8;
            return basic_info->alpha_bits ? AV_PIX_FMT_RGBA : AV_PIX_FMT_RGB24;
        }
        if (basic_info->bits_per_sample > 16)
            av_log(avctx, AV_LOG_WARNING, "Downsampling larger integer to 16-bit via libjxl\n");
        if (basic_info->exponent_bits_per_sample)
            av_log(avctx, AV_LOG_WARNING, "Downsampling float to 16-bit integer via libjxl\n");
        format->data_type = JXL_TYPE_UINT16;
        return basic_info->alpha_bits ? AV_PIX_FMT_RGBA64 : AV_PIX_FMT_RGB48;
    }

    return AV_PIX_FMT_NONE;
}

static enum AVColorPrimaries libjxl_get_primaries(void *avctx, const JxlColorEncoding *jxl_color)
{
    AVColorPrimariesDesc desc;
    enum AVColorPrimaries prim;

    /* libjxl populates these double values even if it uses an enum space */
    desc.prim.r.x = av_d2q(jxl_color->primaries_red_xy[0], 300000);
    desc.prim.r.y = av_d2q(jxl_color->primaries_red_xy[1], 300000);
    desc.prim.g.x = av_d2q(jxl_color->primaries_green_xy[0], 300000);
    desc.prim.g.y = av_d2q(jxl_color->primaries_green_xy[1], 300000);
    desc.prim.b.x = av_d2q(jxl_color->primaries_blue_xy[0], 300000);
    desc.prim.b.y = av_d2q(jxl_color->primaries_blue_xy[1], 300000);
    desc.wp.x = av_d2q(jxl_color->white_point_xy[0], 300000);
    desc.wp.y = av_d2q(jxl_color->white_point_xy[1], 300000);

    prim = av_csp_primaries_id_from_desc(&desc);
    if (prim == AVCOL_PRI_UNSPECIFIED) {
        /* try D65 with the same primaries */
        /* BT.709 uses D65 white point */
        desc.wp = av_csp_primaries_desc_from_id(AVCOL_PRI_BT709)->wp;
        av_log(avctx, AV_LOG_WARNING, "Changing unknown white point to D65\n");
        prim = av_csp_primaries_id_from_desc(&desc);
    }

    return prim;
}

static enum AVColorTransferCharacteristic libjxl_get_trc(void *avctx, const JxlColorEncoding *jxl_color)
{
    switch (jxl_color->transfer_function) {
    case JXL_TRANSFER_FUNCTION_709: return AVCOL_TRC_BT709;
    case JXL_TRANSFER_FUNCTION_LINEAR: return AVCOL_TRC_LINEAR;
    case JXL_TRANSFER_FUNCTION_SRGB: return AVCOL_TRC_IEC61966_2_1;
    case JXL_TRANSFER_FUNCTION_PQ: return AVCOL_TRC_SMPTE2084;
    case JXL_TRANSFER_FUNCTION_DCI: return AVCOL_TRC_SMPTE428;
    case JXL_TRANSFER_FUNCTION_HLG: return AVCOL_TRC_ARIB_STD_B67;
    case JXL_TRANSFER_FUNCTION_GAMMA:
        if (jxl_color->gamma > 2.199 && jxl_color->gamma < 2.201)
            return AVCOL_TRC_GAMMA22;
        else if (jxl_color->gamma > 2.799 && jxl_color->gamma < 2.801)
            return AVCOL_TRC_GAMMA28;
        else
            av_log(avctx, AV_LOG_WARNING, "Unsupported gamma transfer: %f\n", jxl_color->gamma);
        break;
    default:
        av_log(avctx, AV_LOG_WARNING, "Unknown transfer function: %d\n", jxl_color->transfer_function);
    }

    return AVCOL_TRC_UNSPECIFIED;
}

static int libjxl_get_icc(AVCodecContext *avctx)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;
    size_t icc_len;
    JxlDecoderStatus jret;
    /* an ICC profile is present, and we can meaningfully get it,
     * because the pixel data is not XYB-encoded */
    jret = JxlDecoderGetICCProfileSize(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA, &icc_len);
    if (jret == JXL_DEC_SUCCESS && icc_len > 0) {
        av_buffer_unref(&ctx->iccp);
        ctx->iccp = av_buffer_alloc(icc_len);
        if (!ctx->iccp)
            return AVERROR(ENOMEM);
        jret = JxlDecoderGetColorAsICCProfile(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA,
                                                ctx->iccp->data, icc_len);
        if (jret != JXL_DEC_SUCCESS) {
            av_log(avctx, AV_LOG_WARNING, "Unable to obtain ICC Profile\n");
            av_buffer_unref(&ctx->iccp);
        }
    }

    return 0;
}

/*
 * There's generally four cases when it comes to decoding a libjxl image
 * with regard to color encoding:
 * (a) There is an embedded ICC profile in the image, and the image is XYB-encoded.
 * (b) There is an embedded ICC profile in the image, and the image is not XYB-encoded.
 * (c) There is no embedded ICC profile, and FFmpeg supports the tagged colorspace.
 * (d) There is no embedded ICC profile, and FFmpeg does not support the tagged colorspace.
 *
 * In case (b), we forward the pixel data as is and forward the ICC Profile as-is.
 * In case (c), we request the pixel data in the space it's tagged as,
 *     and tag the space accordingly.
 * In case (a), libjxl does not support getting the pixel data in the space described by the ICC
 *     profile, so instead we request the pixel data in BT.2020/PQ as it is the widest
 *     space that FFmpeg supports.
 * In case (d), we also request wide-gamut pixel data as a fallback since FFmpeg doesn't support
 *     the custom primaries tagged in the space.
 */
static int libjxl_color_encoding_event(AVCodecContext *avctx, AVFrame *frame)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;
    JxlDecoderStatus jret;
    int ret;
    JxlColorEncoding jxl_color;
    /* set this flag if we need to fall back on wide gamut */
    int fallback = 0;

    jret = JxlDecoderGetColorAsEncodedProfile(ctx->decoder, NULL, JXL_COLOR_PROFILE_TARGET_ORIGINAL, &jxl_color);
    if (jret == JXL_DEC_SUCCESS) {
        /* enum values describe the colors of this image */
        jret = JxlDecoderSetPreferredColorProfile(ctx->decoder, &jxl_color);
        if (jret == JXL_DEC_SUCCESS)
            jret = JxlDecoderGetColorAsEncodedProfile(ctx->decoder, &ctx->jxl_pixfmt, JXL_COLOR_PROFILE_TARGET_DATA, &jxl_color);
        /* if we couldn't successfully request the pixel data space, we fall back on wide gamut */
        /* this code path is very unlikely to happen in practice */
        if (jret != JXL_DEC_SUCCESS)
            fallback = 1;
    } else {
        /* an ICC Profile is present in the stream */
        if (ctx->basic_info.uses_original_profile) {
            /* uses_original_profile is the same as !xyb_encoded */
            av_log(avctx, AV_LOG_VERBOSE, "Using embedded ICC Profile\n");
            if ((ret = libjxl_get_icc(avctx)) < 0)
                return ret;
        } else {
            /*
             * an XYB-encoded image with an embedded ICC profile can't always have the
             * pixel data requested in the original space, so libjxl has no feature
             * to allow this to happen, so we fall back on wide gamut
             */
            fallback = 1;
        }
    }

    avctx->color_range = frame->color_range = AVCOL_RANGE_JPEG;
    if (ctx->jxl_pixfmt.num_channels >= 3)
        avctx->colorspace = AVCOL_SPC_RGB;
    avctx->color_primaries = AVCOL_PRI_UNSPECIFIED;
    avctx->color_trc = AVCOL_TRC_UNSPECIFIED;

    if (!ctx->iccp) {
        /* checking enum values */
        if (!fallback) {
            if (avctx->colorspace == AVCOL_SPC_RGB)
                avctx->color_primaries = libjxl_get_primaries(avctx, &jxl_color);
            avctx->color_trc = libjxl_get_trc(avctx, &jxl_color);
        }
        /* fall back on wide gamut if enum values fail */
        if (avctx->color_primaries == AVCOL_PRI_UNSPECIFIED) {
            if (avctx->colorspace == AVCOL_SPC_RGB) {
                av_log(avctx, AV_LOG_WARNING, "Falling back on wide gamut output\n");
                jxl_color.primaries = JXL_PRIMARIES_2100;
                avctx->color_primaries = AVCOL_PRI_BT2020;
            }
            /* libjxl requires this set even for grayscale */
            jxl_color.white_point = JXL_WHITE_POINT_D65;
        }
        if (avctx->color_trc == AVCOL_TRC_UNSPECIFIED) {
            if (ctx->jxl_pixfmt.data_type == JXL_TYPE_FLOAT
                || ctx->jxl_pixfmt.data_type == JXL_TYPE_FLOAT16) {
                av_log(avctx, AV_LOG_WARNING, "Falling back on Linear Light transfer\n");
                jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_LINEAR;
                avctx->color_trc = AVCOL_TRC_LINEAR;
            } else {
                av_log(avctx, AV_LOG_WARNING, "Falling back on iec61966-2-1/sRGB transfer\n");
                jxl_color.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;
                avctx->color_trc = AVCOL_TRC_IEC61966_2_1;
            }
        }
        /* all colors will be in-gamut so we want accurate colors */
        jxl_color.rendering_intent = JXL_RENDERING_INTENT_RELATIVE;
        jxl_color.color_space = avctx->colorspace == AVCOL_SPC_RGB ? JXL_COLOR_SPACE_RGB : JXL_COLOR_SPACE_GRAY;
        jret = JxlDecoderSetPreferredColorProfile(ctx->decoder, &jxl_color);
        if (jret != JXL_DEC_SUCCESS) {
            av_log(avctx, AV_LOG_WARNING, "Unable to set fallback color encoding\n");
            /*
             * This should only happen if there's a non-XYB encoded image with custom primaries
             * embedded as enums and no embedded ICC Profile.
             * In this case, libjxl will synthesize an ICC Profile for us.
             */
            avctx->color_trc = AVCOL_TRC_UNSPECIFIED;
            avctx->color_primaries = AVCOL_PRI_UNSPECIFIED;
            if ((ret = libjxl_get_icc(avctx)) < 0)
                return ret;
        }
    }

    frame->color_trc = avctx->color_trc;
    frame->color_primaries = avctx->color_primaries;
    frame->colorspace = avctx->colorspace;

    return 0;
}

static int libjxl_decode_frame(AVCodecContext *avctx, AVFrame *frame, int *got_frame, AVPacket *avpkt)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;
    const uint8_t *buf = avpkt->data;
    size_t remaining = avpkt->size;
    JxlDecoderStatus jret;
    int ret;
    *got_frame = 0;

    while (1) {

        jret = JxlDecoderSetInput(ctx->decoder, buf, remaining);

        if (jret == JXL_DEC_ERROR) {
            /* this should never happen here unless there's a bug in libjxl */
            av_log(avctx, AV_LOG_ERROR, "Unknown libjxl decode error\n");
            return AVERROR_EXTERNAL;
        }

        jret = JxlDecoderProcessInput(ctx->decoder);
        /*
         * JxlDecoderReleaseInput returns the number
         * of bytes remaining to be read, rather than
         * the number of bytes that it did read
         */
        remaining = JxlDecoderReleaseInput(ctx->decoder);
        buf = avpkt->data + avpkt->size - remaining;

        switch(jret) {
        case JXL_DEC_ERROR:
            av_log(avctx, AV_LOG_ERROR, "Unknown libjxl decode error\n");
            return AVERROR_INVALIDDATA;
        case JXL_DEC_NEED_MORE_INPUT:
            if (remaining == 0) {
                av_log(avctx, AV_LOG_ERROR, "Unexpected end of JXL codestream\n");
                return AVERROR_INVALIDDATA;
            }
            av_log(avctx, AV_LOG_DEBUG, "NEED_MORE_INPUT event emitted\n");
            continue;
        case JXL_DEC_BASIC_INFO:
            av_log(avctx, AV_LOG_DEBUG, "BASIC_INFO event emitted\n");
            if (JxlDecoderGetBasicInfo(ctx->decoder, &ctx->basic_info) != JXL_DEC_SUCCESS) {
                /*
                 * this should never happen
                 * if it does it is likely a libjxl decoder bug
                 */
                av_log(avctx, AV_LOG_ERROR, "Bad libjxl basic info event\n");
                return AVERROR_EXTERNAL;
            }
            avctx->pix_fmt = libjxl_get_pix_fmt(avctx, &ctx->basic_info, &ctx->jxl_pixfmt);
            if (avctx->pix_fmt == AV_PIX_FMT_NONE) {
                av_log(avctx, AV_LOG_ERROR, "Bad libjxl pixel format\n");
                return AVERROR_EXTERNAL;
            }
            if ((ret = ff_set_dimensions(avctx, ctx->basic_info.xsize, ctx->basic_info.ysize)) < 0)
                return ret;
            continue;
        case JXL_DEC_COLOR_ENCODING:
            av_log(avctx, AV_LOG_DEBUG, "COLOR_ENCODING event emitted\n");
            if ((ret = libjxl_color_encoding_event(avctx, frame)) < 0)
                return ret;
            continue;
        case JXL_DEC_NEED_IMAGE_OUT_BUFFER:
            av_log(avctx, AV_LOG_DEBUG, "NEED_IMAGE_OUT_BUFFER event emitted\n");
            if ((ret = ff_get_buffer(avctx, frame, 0)) < 0)
                return ret;
            ctx->jxl_pixfmt.align = frame->linesize[0];
            if (JxlDecoderSetImageOutBuffer(ctx->decoder, &ctx->jxl_pixfmt, frame->data[0], frame->buf[0]->size)
                != JXL_DEC_SUCCESS) {
                av_log(avctx, AV_LOG_ERROR, "Bad libjxl dec need image out buffer event\n");
                return AVERROR_EXTERNAL;
            }
            continue;
        case JXL_DEC_FULL_IMAGE:
            /* full image is one frame, even if animated */
            av_log(avctx, AV_LOG_DEBUG, "FULL_IMAGE event emitted\n");
            frame->pict_type = AV_PICTURE_TYPE_I;
            frame->key_frame = 1;
            if (ctx->iccp) {
                AVFrameSideData *sd = av_frame_new_side_data_from_buf(frame, AV_FRAME_DATA_ICC_PROFILE, ctx->iccp);
                if (!sd)
                    return AVERROR(ENOMEM);
                /* ownership is transfered, and it is not ref-ed */
                ctx->iccp = NULL;
            }
            *got_frame = 1;
            return avpkt->size - remaining;
        case JXL_DEC_SUCCESS:
            av_log(avctx, AV_LOG_DEBUG, "SUCCESS event emitted\n");
            /*
             * The SUCCESS event isn't fired until after JXL_DEC_FULL_IMAGE. If this
             * stream only contains one JXL image then JXL_DEC_SUCCESS will never fire.
             * If the image2 sequence being decoded contains several JXL files, then
             * libjxl will fire this event after the next AVPacket has been passed,
             * which means the current packet is actually the next image in the sequence.
             * This is why we reset the decoder and populate the packet data now, since
             * this is the next packet and it has not been decoded yet. The decoder does
             * have to be reset to allow us to use it for the next image, or libjxl
             * will become very confused if the header information is not identical.
             */
            JxlDecoderReset(ctx->decoder);
            libjxl_init_jxl_decoder(avctx);
            buf = avpkt->data;
            remaining = avpkt->size;
            continue;
        default:
             av_log(avctx, AV_LOG_ERROR, "Bad libjxl event: %d\n", jret);
             return AVERROR_EXTERNAL;
        }
    }
}

static av_cold int libjxl_decode_close(AVCodecContext *avctx)
{
    LibJxlDecodeContext *ctx = avctx->priv_data;

    if (ctx->runner)
        JxlThreadParallelRunnerDestroy(ctx->runner);
    ctx->runner = NULL;
    if (ctx->decoder)
        JxlDecoderDestroy(ctx->decoder);
    ctx->decoder = NULL;
    av_buffer_unref(&ctx->iccp);

    return 0;
}

const FFCodec ff_libjxl_decoder = {
    .p.name           = "libjxl",
    CODEC_LONG_NAME("libjxl JPEG XL"),
    .p.type           = AVMEDIA_TYPE_VIDEO,
    .p.id             = AV_CODEC_ID_JPEGXL,
    .priv_data_size   = sizeof(LibJxlDecodeContext),
    .init             = libjxl_decode_init,
    FF_CODEC_DECODE_CB(libjxl_decode_frame),
    .close            = libjxl_decode_close,
    .p.capabilities   = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_OTHER_THREADS,
    .caps_internal    = FF_CODEC_CAP_NOT_INIT_THREADSAFE |
                        FF_CODEC_CAP_AUTO_THREADS | FF_CODEC_CAP_INIT_CLEANUP |
                        FF_CODEC_CAP_ICC_PROFILES,
    .p.wrapper_name   = "libjxl",
};