avcodec/libjxl: respect AVCodecContext->bits_per_raw_sample

libjxl only accepts 16-bit buffers with its API, but it can
accept 9-bit to 15-bit input via a 16-bit buffer, provided the flag
is set declaring the buffer to be of the respective significant depth.

Likewise, it can only provide pixel data on decode as a 16-bit buffer
(if higher than 8) but does provide the metadata tagging the actual bit
depth.

This commit causes libjxlenc.c and libjxldec.c to respect this metadata
and tag/read it accordingly from AVCodecContext->bits_per_raw_sample.

Signed-off-by: Leo Izen <leo.izen@gmail.com>

2 files changed, 22 insertions, 5 deletions

diff --git a/libavcodec/libjxldec.c b/libavcodec/libjxldec.c
index c85cfa3d69..abe08eb400 100644
--- a/libavcodec/libjxldec.c
+++ b/libavcodec/libjxldec.c
@@ -47,6 +47,7 @@ typedef struct LibJxlDecodeContext {
     JxlDecoder *decoder;
     JxlBasicInfo basic_info;
     JxlPixelFormat jxl_pixfmt;
+    JxlBitDepth jxl_bit_depth;
     JxlDecoderStatus events;
     AVBufferRef *iccp;
 } LibJxlDecodeContext;
@@ -93,10 +94,14 @@ static av_cold int libjxl_decode_init(AVCodecContext *avctx)
     return libjxl_init_jxl_decoder(avctx);
 }
 
-static enum AVPixelFormat libjxl_get_pix_fmt(void *avctx, const JxlBasicInfo *basic_info, JxlPixelFormat *format)
+static enum AVPixelFormat libjxl_get_pix_fmt(AVCodecContext *avctx, const JxlBasicInfo *basic_info,
+                                             JxlPixelFormat *format, JxlBitDepth *depth)
 {
     format->endianness = JXL_NATIVE_ENDIAN;
     format->num_channels = basic_info->num_color_channels + (basic_info->alpha_bits > 0);
+    depth->bits_per_sample = avctx->bits_per_raw_sample = basic_info->bits_per_sample;
+    depth->type = JXL_BIT_DEPTH_FROM_PIXEL_FORMAT;
+    depth->exponent_bits_per_sample = basic_info->exponent_bits_per_sample;
     /* Gray */
     if (basic_info->num_color_channels == 1) {
         if (basic_info->bits_per_sample <= 8) {
@@ -119,10 +124,10 @@ static enum AVPixelFormat libjxl_get_pix_fmt(void *avctx, const JxlBasicInfo *ba
             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");
+        else if (basic_info->bits_per_sample > 16)
+            av_log(avctx, AV_LOG_WARNING, "Downsampling larger integer to 16-bit via libjxl\n");
         format->data_type = JXL_TYPE_UINT16;
         return basic_info->alpha_bits ? AV_PIX_FMT_RGBA64 : AV_PIX_FMT_RGB48;
     }
@@ -367,7 +372,7 @@ static int libjxl_decode_frame(AVCodecContext *avctx, AVFrame *frame, int *got_f
                 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);
+            avctx->pix_fmt = libjxl_get_pix_fmt(avctx, &ctx->basic_info, &ctx->jxl_pixfmt, &ctx->jxl_bit_depth);
             if (avctx->pix_fmt == AV_PIX_FMT_NONE) {
                 av_log(avctx, AV_LOG_ERROR, "Bad libjxl pixel format\n");
                 return AVERROR_EXTERNAL;
@@ -390,6 +395,10 @@ static int libjxl_decode_frame(AVCodecContext *avctx, AVFrame *frame, int *got_f
                 av_log(avctx, AV_LOG_ERROR, "Bad libjxl dec need image out buffer event\n");
                 return AVERROR_EXTERNAL;
             }
+            if (JxlDecoderSetImageOutBitDepth(ctx->decoder, &ctx->jxl_bit_depth) != JXL_DEC_SUCCESS) {
+                av_log(avctx, AV_LOG_ERROR, "Error setting output bit depth\n");
+                return AVERROR_EXTERNAL;
+            }
             continue;
         case JXL_DEC_FULL_IMAGE:
             /* full image is one frame, even if animated */
diff --git a/libavcodec/libjxlenc.c b/libavcodec/libjxlenc.c
index 06581b268d..c51024f180 100644
--- a/libavcodec/libjxlenc.c
+++ b/libavcodec/libjxlenc.c
@@ -250,6 +250,7 @@ static int libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFra
     JxlBasicInfo info;
     JxlColorEncoding jxl_color;
     JxlPixelFormat jxl_fmt;
+    JxlBitDepth jxl_bit_depth;
     JxlEncoderStatus jret;
     int ret;
     size_t available = ctx->buffer_size;
@@ -269,17 +270,22 @@ static int libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFra
     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;
-    info.bits_per_sample = av_get_bits_per_pixel(pix_desc) / jxl_fmt.num_channels;
+    jxl_bit_depth.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 : jxl_bit_depth.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;
+        jxl_bit_depth.exponent_bits_per_sample = info.exponent_bits_per_sample;
     } 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;
+        jxl_bit_depth.exponent_bits_per_sample = 0;
     }
+    jxl_bit_depth.type = JXL_BIT_DEPTH_FROM_PIXEL_FORMAT;
 
     /* JPEG XL format itself does not support limited range */
     if (avctx->color_range == AVCOL_RANGE_MPEG ||
@@ -356,6 +362,8 @@ static int libjxl_encode_frame(AVCodecContext *avctx, AVPacket *pkt, const AVFra
         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 (JxlEncoderSetFrameBitDepth(ctx->options, &jxl_bit_depth) != JXL_ENC_SUCCESS)
+        av_log(avctx, AV_LOG_WARNING, "Failed to set JxlBitDepth\n");
 
     /* depending on basic info, level 10 might
      * be required instead of level 5 */