avcodec/mediacodecenc: Add global header support

The extradata is generated by encoding a dummy frame, then reset the encoder state by mediacodec flush(). It only works for pixel format other than AV_PIX_FMT_MEDIACODEC, since I'm not sure how to create a dummy frame safely with AV_PIX_FMT_MEDIACODEC. Signed-off-by: Zhao Zhili <zhilizhao@tencent.com>
11 months ago · 9e49915195
parent 3cfea6993a
commit 9e49915195
2 changed files with 155 additions and 17 deletions
--- a/6
+++ b/6
@ -3315,6 +3315,7 @@ ac3_mf_encoder_deps="mediafoundation"
 av1_cuvid_decoder_deps="cuvid CUVIDAV1PICPARAMS"
 av1_mediacodec_decoder_deps="mediacodec"
 av1_mediacodec_encoder_deps="mediacodec"
+av1_mediacodec_encoder_select="extract_extradata_bsf"
 av1_nvenc_encoder_deps="nvenc NV_ENC_PIC_PARAMS_AV1"
 av1_nvenc_encoder_select="atsc_a53"
 h263_v4l2m2m_decoder_deps="v4l2_m2m h263_v4l2_m2m"
@ -3325,7 +3326,7 @@ h264_cuvid_decoder_select="h264_mp4toannexb_bsf"
 h264_mediacodec_decoder_deps="mediacodec"
 h264_mediacodec_decoder_select="h264_mp4toannexb_bsf h264_parser"
 h264_mediacodec_encoder_deps="mediacodec"
-h264_mediacodec_encoder_select="h264_metadata"
+h264_mediacodec_encoder_select="extract_extradata_bsf h264_metadata"
 h264_mf_encoder_deps="mediafoundation"
 h264_mmal_decoder_deps="mmal"
 h264_nvenc_encoder_deps="nvenc"
@ -3345,7 +3346,7 @@ hevc_cuvid_decoder_select="hevc_mp4toannexb_bsf"
 hevc_mediacodec_decoder_deps="mediacodec"
 hevc_mediacodec_decoder_select="hevc_mp4toannexb_bsf hevc_parser"
 hevc_mediacodec_encoder_deps="mediacodec"
-hevc_mediacodec_encoder_select="hevc_metadata"
+hevc_mediacodec_encoder_select="extract_extradata_bsf hevc_metadata"
 hevc_mf_encoder_deps="mediafoundation"
 hevc_nvenc_encoder_deps="nvenc"
 hevc_nvenc_encoder_select="atsc_a53"
@ -3377,6 +3378,7 @@ mpeg2_v4l2m2m_decoder_deps="v4l2_m2m mpeg2_v4l2_m2m"
 mpeg4_cuvid_decoder_deps="cuvid"
 mpeg4_mediacodec_decoder_deps="mediacodec"
 mpeg4_mediacodec_encoder_deps="mediacodec"
+mpeg4_mediacodec_encoder_select="extract_extradata_bsf"
 mpeg4_mmal_decoder_deps="mmal"
 mpeg4_omx_encoder_deps="omx"
 mpeg4_v4l2m2m_decoder_deps="v4l2_m2m mpeg4_v4l2_m2m"
--- a/libavcodec/mediacodecenc.c
+++ b/libavcodec/mediacodecenc.c
@ -23,6 +23,7 @@
 #include "config_components.h"

 #include "libavutil/avassert.h"
+#include "libavutil/avstring.h"
 #include "libavutil/hwcontext_mediacodec.h"
 #include "libavutil/imgutils.h"
 #include "libavutil/mem.h"
@ -74,6 +75,7 @@ typedef struct MediaCodecEncContext {
    int bitrate_mode;
    int level;
    int pts_as_dts;
+    int extract_extradata;
 } MediaCodecEncContext;

 enum {
@ -112,6 +114,23 @@ static void mediacodec_output_format(AVCodecContext *avctx)
    ff_AMediaFormat_delete(out_format);
 }

+static int extract_extradata_support(AVCodecContext *avctx)
+{
+    const AVBitStreamFilter *bsf = av_bsf_get_by_name("extract_extradata");
+
+    if (!bsf) {
+        av_log(avctx, AV_LOG_WARNING, "extract_extradata bsf not found\n");
+        return 0;
+    }
+
+    for (int i = 0; bsf->codec_ids[i] != AV_CODEC_ID_NONE; i++) {
+        if (bsf->codec_ids[i] == avctx->codec_id)
+            return 1;
+    }
+
+    return 0;
+}
+
 static int mediacodec_init_bsf(AVCodecContext *avctx)
 {
    MediaCodecEncContext *s = avctx->priv_data;
@ -120,20 +139,32 @@ static int mediacodec_init_bsf(AVCodecContext *avctx)
    int crop_right = s->width - avctx->width;
    int crop_bottom = s->height - avctx->height;

-    if (!crop_right && !crop_bottom)
+    /* Nothing can be done for this format now */
+    if (avctx->pix_fmt == AV_PIX_FMT_MEDIACODEC)
+        return 0;
+
+    s->extract_extradata = (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) &&
+                           extract_extradata_support(avctx);
+    if (!crop_right && !crop_bottom && !s->extract_extradata)
        return 0;

+    ret = 0;
+    if (crop_right || crop_bottom) {
        if (avctx->codec_id == AV_CODEC_ID_H264)
            ret = snprintf(str, sizeof(str), "h264_metadata=crop_right=%d:crop_bottom=%d",
                           crop_right, crop_bottom);
        else if (avctx->codec_id == AV_CODEC_ID_HEVC)
            ret = snprintf(str, sizeof(str), "hevc_metadata=crop_right=%d:crop_bottom=%d",
                           crop_right, crop_bottom);
-    else
-        return 0;
+        if (ret >= sizeof(str))
+            return AVERROR_BUFFER_TOO_SMALL;
+    }

+    if (s->extract_extradata) {
+        ret = av_strlcatf(str, sizeof(str), "%sextract_extradata", ret ? "," : "");
        if (ret >= sizeof(str))
            return AVERROR_BUFFER_TOO_SMALL;
+    }

    ret = av_bsf_list_parse_str(str, &s->bsf);
    if (ret < 0)
@ -148,6 +179,8 @@ static int mediacodec_init_bsf(AVCodecContext *avctx)
    return ret;
 }

+static int mediacodec_generate_extradata(AVCodecContext *avctx);
+
 static av_cold int mediacodec_init(AVCodecContext *avctx)
 {
    const char *codec_mime = NULL;
@ -337,14 +370,14 @@ static av_cold int mediacodec_init(AVCodecContext *avctx)
        goto bailout;

    mediacodec_output_format(avctx);
-    if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER)
-        av_log(avctx, AV_LOG_WARNING,
-                "Mediacodec encoder doesn't support AV_CODEC_FLAG_GLOBAL_HEADER. "
-                "Use extract_extradata bsf when necessary.\n");

    s->frame = av_frame_alloc();
-    if (!s->frame)
+    if (!s->frame) {
        ret = AVERROR(ENOMEM);
+        goto bailout;
+    }
+
+    ret = mediacodec_generate_extradata(avctx);

 bailout:
    if (format)
@ -549,6 +582,109 @@ static int mediacodec_encode(AVCodecContext *avctx, AVPacket *pkt)
    return 0;
 }

+static int mediacodec_send_dummy_frame(AVCodecContext *avctx)
+{
+    MediaCodecEncContext *s = avctx->priv_data;
+    int ret;
+
+    s->frame->width = avctx->width;
+    s->frame->height = avctx->height;
+    s->frame->format = avctx->pix_fmt;
+    s->frame->pts = 0;
+
+    ret = av_frame_get_buffer(s->frame, 0);
+    if (ret < 0)
+        return ret;
+
+    do {
+        ret = mediacodec_send(avctx, s->frame);
+    } while (ret == AVERROR(EAGAIN));
+    av_frame_unref(s->frame);
+
+    if (ret < 0)
+        return ret;
+
+    ret = mediacodec_send(avctx, NULL);
+    if (ret < 0) {
+        av_log(avctx, AV_LOG_ERROR, "Flush failed: %s\n", av_err2str(ret));
+        return ret;
+    }
+
+    return 0;
+}
+
+static int mediacodec_receive_dummy_pkt(AVCodecContext *avctx, AVPacket *pkt)
+{
+    MediaCodecEncContext *s = avctx->priv_data;
+    int ret;
+
+    do {
+        ret = mediacodec_receive(avctx, pkt);
+    } while (ret == AVERROR(EAGAIN));
+
+    if (ret < 0)
+        return ret;
+
+    do {
+        ret = av_bsf_send_packet(s->bsf, pkt);
+        if (ret < 0)
+            return ret;
+        ret = av_bsf_receive_packet(s->bsf, pkt);
+    } while (ret == AVERROR(EAGAIN));
+
+    return ret;
+}
+
+static int mediacodec_generate_extradata(AVCodecContext *avctx)
+{
+    MediaCodecEncContext *s = avctx->priv_data;
+    AVPacket *pkt = NULL;
+    int ret;
+    size_t side_size;
+    uint8_t *side;
+
+    if (!(avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER))
+        return 0;
+
+    if (!s->extract_extradata) {
+        av_log(avctx, AV_LOG_WARNING,
+               "Mediacodec encoder doesn't support AV_CODEC_FLAG_GLOBAL_HEADER. "
+               "Use extract_extradata bsf when necessary.\n");
+        return 0;
+    }
+
+    pkt = av_packet_alloc();
+    if (!pkt)
+        return AVERROR(ENOMEM);
+
+    ret = mediacodec_send_dummy_frame(avctx);
+    if (ret < 0)
+        goto bailout;
+    ret = mediacodec_receive_dummy_pkt(avctx, pkt);
+    if (ret < 0)
+        goto bailout;
+
+    side = av_packet_get_side_data(pkt, AV_PKT_DATA_NEW_EXTRADATA, &side_size);
+    if (side && side_size > 0) {
+        avctx->extradata = av_mallocz(side_size + AV_INPUT_BUFFER_PADDING_SIZE);
+        if (!avctx->extradata) {
+            ret = AVERROR(ENOMEM);
+            goto bailout;
+        }
+
+        memcpy(avctx->extradata, side, side_size);
+        avctx->extradata_size = side_size;
+    }
+
+bailout:
+    if (s->eof_sent) {
+        s->eof_sent = 0;
+        ff_AMediaCodec_flush(s->codec);
+    }
+    av_packet_free(&pkt);
+    return ret;
+}
+
 static av_cold int mediacodec_close(AVCodecContext *avctx)
 {
    MediaCodecEncContext *s = avctx->priv_data;