libavcodec/bsf/dovi_split.c - third_party/ffmpeg - Git at Google

 /*
  * 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
  */

 #include <stdbool.h>

 #include "libavutil/avassert.h"
 #include "libavutil/dovi_meta.h"
 #include "libavutil/intreadwrite.h"
 #include "libavutil/mem.h"
 #include "libavutil/opt.h"

 #include "libavcodec/bsf.h"
 #include "libavcodec/bsf_internal.h"
 #include "libavcodec/h2645_parse.h"
 #include "libavcodec/packet.h"

 #include "libavcodec/hevc/hevc.h"

 enum DOVISplitMode {
     DOVI_SPLIT_BL     = 0,
     DOVI_SPLIT_BL_RPU = 1,
     DOVI_SPLIT_EL     = 2,
     DOVI_SPLIT_EL_RPU = 3,
 };

 typedef struct DOVISplitContext {
     const AVClass *class;
     int mode;

     int nal_length_size;        /* 0 means Annex-B input */
     int out_nal_length_size;    /* 0 means Annex-B output */
     H2645Packet pkt;
 } DOVISplitContext;

 static int hvcc_nal_length_size(const uint8_t *data, int size)
 {
     if (size >= 23 && data[0] == 1 && AV_RB24(data) != 1 && AV_RB32(data) != 1)
         return (data[21] & 3) + 1;
     return 0;
 }

 static int dovi_split_init(AVBSFContext *ctx)
 {
     DOVISplitContext *s = ctx->priv_data;
     bool keep_bl  = s->mode == DOVI_SPLIT_BL || s->mode == DOVI_SPLIT_BL_RPU;
     bool keep_el  = s->mode == DOVI_SPLIT_EL || s->mode == DOVI_SPLIT_EL_RPU;
     bool keep_rpu = s->mode == DOVI_SPLIT_BL_RPU || s->mode == DOVI_SPLIT_EL_RPU;

     /* Profile 7 is currently the only supported variant with EL by Dolby.
      * Default to that in case there is no DOVI config. */
     uint8_t dv_profile = 7;

     for (int i = 0; i < ctx->par_out->nb_coded_side_data; i++) {
         AVPacketSideData *sd = &ctx->par_out->coded_side_data[i];
         AVDOVIDecoderConfigurationRecord *cfg;
         if (sd->type != AV_PKT_DATA_DOVI_CONF)
             continue;
         cfg = (AVDOVIDecoderConfigurationRecord *)sd->data;
         cfg->bl_present_flag  &= keep_bl;
         cfg->el_present_flag  &= keep_el;
         cfg->rpu_present_flag &= keep_rpu;
         dv_profile = cfg->dv_profile;
         break;
     }

     if (keep_el) {
         const AVPacketSideData *sd;
         sd = av_packet_side_data_get(ctx->par_out->coded_side_data,
                                      ctx->par_out->nb_coded_side_data,
                                      AV_PKT_DATA_HEVC_CONF);
         if (sd && sd->size >= 23) {
             uint8_t *new_ed = av_mallocz(sd->size + AV_INPUT_BUFFER_PADDING_SIZE);
             if (!new_ed)
                 return AVERROR(ENOMEM);
             memcpy(new_ed, sd->data, sd->size);
             av_freep(&ctx->par_out->extradata);
             ctx->par_out->extradata = new_ed;
             ctx->par_out->extradata_size = sd->size;
         }

         /* DV profile to EL size ratio */
         static const uint8_t el_div[] = {
             [2] = 2,
             [3] = 1,
             [4] = 2,
             [6] = 2,
             [7] = 2,
         };
         int div = dv_profile < FF_ARRAY_ELEMS(el_div) ? el_div[dv_profile] : 0;
         if (!div)
             av_log(ctx, AV_LOG_WARNING, "Unexpected DV Profile %d.\n", dv_profile);

         /* P7: EL is 1:1 for FHD BL */
         if (dv_profile == 7 && ctx->par_in->width <= 1920)
             div = 1;
         if (div > 1) {
             ctx->par_out->width = ctx->par_in->width / div;
             ctx->par_out->height = ctx->par_in->height / div;
         }
     }

     /* Drop AV_PKT_DATA_HEVC_CONF as it's no longer valid on output. It's
      * set as extradata for EL. */
     av_packet_side_data_remove(ctx->par_out->coded_side_data,
                                &ctx->par_out->nb_coded_side_data,
                                AV_PKT_DATA_HEVC_CONF);

     s->nal_length_size = hvcc_nal_length_size(ctx->par_in->extradata,
                                               ctx->par_in->extradata_size);
     s->out_nal_length_size = hvcc_nal_length_size(ctx->par_out->extradata,
                                                   ctx->par_out->extradata_size);

     return 0;
 }

 static void dovi_split_close(AVBSFContext *ctx)
 {
     DOVISplitContext *s = ctx->priv_data;
     ff_h2645_packet_uninit(&s->pkt);
 }

 static int nal_is_kept(const DOVISplitContext *s, const H2645NAL *nal,
                        const uint8_t **payload, int *payload_size)
 {
     bool keep_el  = s->mode == DOVI_SPLIT_EL || s->mode == DOVI_SPLIT_EL_RPU;
     bool keep_bl  = s->mode == DOVI_SPLIT_BL || s->mode == DOVI_SPLIT_BL_RPU;
     bool keep_rpu = s->mode == DOVI_SPLIT_BL_RPU || s->mode == DOVI_SPLIT_EL_RPU;

     switch (nal->type) {
     case HEVC_NAL_UNSPEC63:
         /* EL: keep only when extracting EL, strip two-bytes of outer NAL header */
         if (!keep_el || nal->raw_size <= 2)
             return 0;
         *payload = nal->raw_data + 2;
         *payload_size = nal->raw_size - 2;
         return 1;
     case HEVC_NAL_UNSPEC62:
         /* RPU: kept verbatim only when the selected mode opted in. */
         if (!keep_rpu)
             return 0;
         *payload = nal->raw_data;
         *payload_size = nal->raw_size;
         return 1;
     default:
         /* Anything else is a base-layer NAL. */
         if (!keep_bl)
             return 0;
         *payload = nal->raw_data;
         *payload_size = nal->raw_size;
         return 1;
     }
 }

 static int dovi_split_filter(AVBSFContext *ctx, AVPacket *out)
 {
     DOVISplitContext *s = ctx->priv_data;
     AVPacket *in = NULL;
     AVBufferRef *out_buf = NULL;
     uint8_t *dst;
     size_t out_size = 0;
     int kept_count = 0;
     int flags = (s->nal_length_size ? H2645_FLAG_IS_NALFF : 0) |
                 H2645_FLAG_SMALL_PADDING;
     int prefix_size = s->out_nal_length_size ? s->out_nal_length_size : 4;
     int ret;

     ret = ff_bsf_get_packet(ctx, &in);
     if (ret < 0)
         return ret;

     ret = ff_h2645_packet_split(&s->pkt, in->data, in->size, ctx,
                                 s->nal_length_size, AV_CODEC_ID_HEVC, flags);
     if (ret < 0)
         goto fail;

     for (int i = 0; i < s->pkt.nb_nals; i++) {
         const uint8_t *payload;
         int payload_size;
         if (!nal_is_kept(s, &s->pkt.nals[i], &payload, &payload_size))
             continue;
         out_size += prefix_size + payload_size;
         kept_count++;
     }

     if (!kept_count) {
         ret = AVERROR(EAGAIN);
         goto fail;
     }

     out_buf = av_buffer_alloc(out_size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!out_buf) {
         ret = AVERROR(ENOMEM);
         goto fail;
     }

     dst = out_buf->data;
     for (int i = 0; i < s->pkt.nb_nals; i++) {
         const uint8_t *payload;
         int payload_size;
         if (!nal_is_kept(s, &s->pkt.nals[i], &payload, &payload_size))
             continue;
         switch (s->out_nal_length_size) {
         case 0: AV_WB32(dst, 1);            break;
         case 1: AV_WB8 (dst, payload_size); break;
         case 2: AV_WB16(dst, payload_size); break;
         case 3: AV_WB24(dst, payload_size); break;
         case 4: AV_WB32(dst, payload_size); break;
         }
         dst += prefix_size;
         memcpy(dst, payload, payload_size);
         dst += payload_size;
     }
     memset(dst, 0, AV_INPUT_BUFFER_PADDING_SIZE);
     av_assert0(dst == out_buf->data + out_size);

     ret = av_packet_copy_props(out, in);
     if (ret < 0)
         goto fail;

     out->buf  = out_buf;
     out->data = out_buf->data;
     out->size = out_size;
     out_buf   = NULL;

 fail:
     av_buffer_unref(&out_buf);
     av_packet_free(&in);
     if (ret < 0 && ret != AVERROR(EAGAIN))
         av_packet_unref(out);
     return ret;
 }

 #define OFFSET(x) offsetof(DOVISplitContext, x)
 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_BSF_PARAM)
 static const AVOption dovi_split_options[] = {
     { "mode", "Which Dolby Vision components to keep in the output bitstream", OFFSET(mode), AV_OPT_TYPE_INT, { .i64 = DOVI_SPLIT_BL }, DOVI_SPLIT_BL, DOVI_SPLIT_EL_RPU, FLAGS, .unit = "mode" },
         { "bl", "Base layer only", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_BL }, .flags = FLAGS, .unit = "mode" },
         { "bl_rpu", "Base layer with the RPU NAL", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_BL_RPU }, .flags = FLAGS, .unit = "mode" },
         { "el", "Enhancement layer only", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_EL }, .flags = FLAGS, .unit = "mode" },
         { "el_rpu", "Enhancement layer with the RPU NAL", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_EL_RPU }, .flags = FLAGS, .unit = "mode" },
     { NULL },
 };

 static const AVClass dovi_split_class = {
     .class_name = "dovi_split_bsf",
     .item_name  = av_default_item_name,
     .option     = dovi_split_options,
     .version    = LIBAVUTIL_VERSION_INT,
 };

 static const enum AVCodecID dovi_split_codec_ids[] = {
     AV_CODEC_ID_HEVC, AV_CODEC_ID_NONE,
 };

 const FFBitStreamFilter ff_dovi_split_bsf = {
     .p.name         = "dovi_split",
     .p.codec_ids    = dovi_split_codec_ids,
     .p.priv_class   = &dovi_split_class,
     .priv_data_size = sizeof(DOVISplitContext),
     .init           = dovi_split_init,
     .close          = dovi_split_close,
     .filter         = dovi_split_filter,
 };
	/*
	* 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
	*/

	#include <stdbool.h>

	#include "libavutil/avassert.h"
	#include "libavutil/dovi_meta.h"
	#include "libavutil/intreadwrite.h"
	#include "libavutil/mem.h"
	#include "libavutil/opt.h"

	#include "libavcodec/bsf.h"
	#include "libavcodec/bsf_internal.h"
	#include "libavcodec/h2645_parse.h"
	#include "libavcodec/packet.h"

	#include "libavcodec/hevc/hevc.h"

	enum DOVISplitMode {
	DOVI_SPLIT_BL = 0,
	DOVI_SPLIT_BL_RPU = 1,
	DOVI_SPLIT_EL = 2,
	DOVI_SPLIT_EL_RPU = 3,
	};

	typedef struct DOVISplitContext {
	const AVClass *class;
	int mode;

	int nal_length_size; /* 0 means Annex-B input */
	int out_nal_length_size; /* 0 means Annex-B output */
	H2645Packet pkt;
	} DOVISplitContext;

	static int hvcc_nal_length_size(const uint8_t *data, int size)
	{
	if (size >= 23 && data[0] == 1 && AV_RB24(data) != 1 && AV_RB32(data) != 1)
	return (data[21] & 3) + 1;
	return 0;
	}

	static int dovi_split_init(AVBSFContext *ctx)
	{
	DOVISplitContext *s = ctx->priv_data;
	bool keep_bl = s->mode == DOVI_SPLIT_BL \|\| s->mode == DOVI_SPLIT_BL_RPU;
	bool keep_el = s->mode == DOVI_SPLIT_EL \|\| s->mode == DOVI_SPLIT_EL_RPU;
	bool keep_rpu = s->mode == DOVI_SPLIT_BL_RPU \|\| s->mode == DOVI_SPLIT_EL_RPU;

	/* Profile 7 is currently the only supported variant with EL by Dolby.
	* Default to that in case there is no DOVI config. */
	uint8_t dv_profile = 7;

	for (int i = 0; i < ctx->par_out->nb_coded_side_data; i++) {
	AVPacketSideData *sd = &ctx->par_out->coded_side_data[i];
	AVDOVIDecoderConfigurationRecord *cfg;
	if (sd->type != AV_PKT_DATA_DOVI_CONF)
	continue;
	cfg = (AVDOVIDecoderConfigurationRecord *)sd->data;
	cfg->bl_present_flag &= keep_bl;
	cfg->el_present_flag &= keep_el;
	cfg->rpu_present_flag &= keep_rpu;
	dv_profile = cfg->dv_profile;
	break;
	}

	if (keep_el) {
	const AVPacketSideData *sd;
	sd = av_packet_side_data_get(ctx->par_out->coded_side_data,
	ctx->par_out->nb_coded_side_data,
	AV_PKT_DATA_HEVC_CONF);
	if (sd && sd->size >= 23) {
	uint8_t *new_ed = av_mallocz(sd->size + AV_INPUT_BUFFER_PADDING_SIZE);
	if (!new_ed)
	return AVERROR(ENOMEM);
	memcpy(new_ed, sd->data, sd->size);
	av_freep(&ctx->par_out->extradata);
	ctx->par_out->extradata = new_ed;
	ctx->par_out->extradata_size = sd->size;
	}

	/* DV profile to EL size ratio */
	static const uint8_t el_div[] = {
	[2] = 2,
	[3] = 1,
	[4] = 2,
	[6] = 2,
	[7] = 2,
	};
	int div = dv_profile < FF_ARRAY_ELEMS(el_div) ? el_div[dv_profile] : 0;
	if (!div)
	av_log(ctx, AV_LOG_WARNING, "Unexpected DV Profile %d.\n", dv_profile);

	/* P7: EL is 1:1 for FHD BL */
	if (dv_profile == 7 && ctx->par_in->width <= 1920)
	div = 1;
	if (div > 1) {
	ctx->par_out->width = ctx->par_in->width / div;
	ctx->par_out->height = ctx->par_in->height / div;
	}
	}

	/* Drop AV_PKT_DATA_HEVC_CONF as it's no longer valid on output. It's
	* set as extradata for EL. */
	av_packet_side_data_remove(ctx->par_out->coded_side_data,
	&ctx->par_out->nb_coded_side_data,
	AV_PKT_DATA_HEVC_CONF);

	s->nal_length_size = hvcc_nal_length_size(ctx->par_in->extradata,
	ctx->par_in->extradata_size);
	s->out_nal_length_size = hvcc_nal_length_size(ctx->par_out->extradata,
	ctx->par_out->extradata_size);

	return 0;
	}

	static void dovi_split_close(AVBSFContext *ctx)
	{
	DOVISplitContext *s = ctx->priv_data;
	ff_h2645_packet_uninit(&s->pkt);
	}

	static int nal_is_kept(const DOVISplitContext s, const H2645NAL nal,
	const uint8_t *payload, int payload_size)
	{
	bool keep_el = s->mode == DOVI_SPLIT_EL \|\| s->mode == DOVI_SPLIT_EL_RPU;
	bool keep_bl = s->mode == DOVI_SPLIT_BL \|\| s->mode == DOVI_SPLIT_BL_RPU;
	bool keep_rpu = s->mode == DOVI_SPLIT_BL_RPU \|\| s->mode == DOVI_SPLIT_EL_RPU;

	switch (nal->type) {
	case HEVC_NAL_UNSPEC63:
	/* EL: keep only when extracting EL, strip two-bytes of outer NAL header */
	if (!keep_el \|\| nal->raw_size <= 2)
	return 0;
	*payload = nal->raw_data + 2;
	*payload_size = nal->raw_size - 2;
	return 1;
	case HEVC_NAL_UNSPEC62:
	/* RPU: kept verbatim only when the selected mode opted in. */
	if (!keep_rpu)
	return 0;
	*payload = nal->raw_data;
	*payload_size = nal->raw_size;
	return 1;
	default:
	/* Anything else is a base-layer NAL. */
	if (!keep_bl)
	return 0;
	*payload = nal->raw_data;
	*payload_size = nal->raw_size;
	return 1;
	}
	}

	static int dovi_split_filter(AVBSFContext ctx, AVPacket out)
	{
	DOVISplitContext *s = ctx->priv_data;
	AVPacket *in = NULL;
	AVBufferRef *out_buf = NULL;
	uint8_t *dst;
	size_t out_size = 0;
	int kept_count = 0;
	int flags = (s->nal_length_size ? H2645_FLAG_IS_NALFF : 0) \|
	H2645_FLAG_SMALL_PADDING;
	int prefix_size = s->out_nal_length_size ? s->out_nal_length_size : 4;
	int ret;

	ret = ff_bsf_get_packet(ctx, &in);
	if (ret < 0)
	return ret;

	ret = ff_h2645_packet_split(&s->pkt, in->data, in->size, ctx,
	s->nal_length_size, AV_CODEC_ID_HEVC, flags);
	if (ret < 0)
	goto fail;

	for (int i = 0; i < s->pkt.nb_nals; i++) {
	const uint8_t *payload;
	int payload_size;
	if (!nal_is_kept(s, &s->pkt.nals[i], &payload, &payload_size))
	continue;
	out_size += prefix_size + payload_size;
	kept_count++;
	}

	if (!kept_count) {
	ret = AVERROR(EAGAIN);
	goto fail;
	}

	out_buf = av_buffer_alloc(out_size + AV_INPUT_BUFFER_PADDING_SIZE);
	if (!out_buf) {
	ret = AVERROR(ENOMEM);
	goto fail;
	}

	dst = out_buf->data;
	for (int i = 0; i < s->pkt.nb_nals; i++) {
	const uint8_t *payload;
	int payload_size;
	if (!nal_is_kept(s, &s->pkt.nals[i], &payload, &payload_size))
	continue;
	switch (s->out_nal_length_size) {
	case 0: AV_WB32(dst, 1); break;
	case 1: AV_WB8 (dst, payload_size); break;
	case 2: AV_WB16(dst, payload_size); break;
	case 3: AV_WB24(dst, payload_size); break;
	case 4: AV_WB32(dst, payload_size); break;
	}
	dst += prefix_size;
	memcpy(dst, payload, payload_size);
	dst += payload_size;
	}
	memset(dst, 0, AV_INPUT_BUFFER_PADDING_SIZE);
	av_assert0(dst == out_buf->data + out_size);

	ret = av_packet_copy_props(out, in);
	if (ret < 0)
	goto fail;

	out->buf = out_buf;
	out->data = out_buf->data;
	out->size = out_size;
	out_buf = NULL;

	fail:
	av_buffer_unref(&out_buf);
	av_packet_free(&in);
	if (ret < 0 && ret != AVERROR(EAGAIN))
	av_packet_unref(out);
	return ret;
	}

	#define OFFSET(x) offsetof(DOVISplitContext, x)
	#define FLAGS (AV_OPT_FLAG_VIDEO_PARAM \| AV_OPT_FLAG_BSF_PARAM)
	static const AVOption dovi_split_options[] = {
	{ "mode", "Which Dolby Vision components to keep in the output bitstream", OFFSET(mode), AV_OPT_TYPE_INT, { .i64 = DOVI_SPLIT_BL }, DOVI_SPLIT_BL, DOVI_SPLIT_EL_RPU, FLAGS, .unit = "mode" },
	{ "bl", "Base layer only", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_BL }, .flags = FLAGS, .unit = "mode" },
	{ "bl_rpu", "Base layer with the RPU NAL", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_BL_RPU }, .flags = FLAGS, .unit = "mode" },
	{ "el", "Enhancement layer only", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_EL }, .flags = FLAGS, .unit = "mode" },
	{ "el_rpu", "Enhancement layer with the RPU NAL", 0, AV_OPT_TYPE_CONST, { .i64 = DOVI_SPLIT_EL_RPU }, .flags = FLAGS, .unit = "mode" },
	{ NULL },
	};

	static const AVClass dovi_split_class = {
	.class_name = "dovi_split_bsf",
	.item_name = av_default_item_name,
	.option = dovi_split_options,
	.version = LIBAVUTIL_VERSION_INT,
	};

	static const enum AVCodecID dovi_split_codec_ids[] = {
	AV_CODEC_ID_HEVC, AV_CODEC_ID_NONE,
	};

	const FFBitStreamFilter ff_dovi_split_bsf = {
	.p.name = "dovi_split",
	.p.codec_ids = dovi_split_codec_ids,
	.p.priv_class = &dovi_split_class,
	.priv_data_size = sizeof(DOVISplitContext),
	.init = dovi_split_init,
	.close = dovi_split_close,
	.filter = dovi_split_filter,
	};