From 3d6c6214e8e388cd866819f41716c6480b42db2b Mon Sep 17 00:00:00 2001
From: Matthieu Bouron <matthieu.bouron@stupeflix.com>
Date: Mon, 5 Sep 2016 17:08:41 +0200
Subject: [PATCH] lavc: add hevc mediacodec decoder

---
 configure                                          |   3 +
 libavcodec/Makefile                                |   3 +-
 libavcodec/allcodecs.c                             |   2 +
 libavcodec/hevc_parse.c                            | 134 ++++++++++++++
 libavcodec/hevc_parse.h                            |  33 ++++
 libavcodec/mediacodec_wrapper.c                    |  33 +++-
 libavcodec/mediacodecdec.c                         |   7 +
 ...{mediacodecdec_h264.c => mediacodecdec_h2645.c} | 198 ++++++++++++++++++---
 8 files changed, 376 insertions(+), 37 deletions(-)
 create mode 100644 libavcodec/hevc_parse.c
 create mode 100644 libavcodec/hevc_parse.h
 rename libavcodec/{mediacodecdec_h264.c => mediacodecdec_h2645.c} (68%)

diff --git a/configure b/configure
index b11ca7f..af3fbf4 100755
--- a/configure
+++ b/configure
@@ -2585,6 +2585,9 @@ h264_videotoolbox_hwaccel_select="h264_decoder"
 hevc_cuvid_hwaccel_deps="cuda cuvid CUVIDHEVCPICPARAMS"
 hevc_d3d11va_hwaccel_deps="d3d11va DXVA_PicParams_HEVC"
 hevc_d3d11va_hwaccel_select="hevc_decoder"
+hevc_mediacodec_decoder_deps="mediacodec"
+hevc_mediacodec_hwaccel_deps="mediacodec"
+hevc_mediacodec_decoder_select="hevc_mp4toannexb_bsf hevc_parser"
 hevc_dxva2_hwaccel_deps="dxva2 DXVA_PicParams_HEVC"
 hevc_dxva2_hwaccel_select="hevc_decoder"
 hevc_qsv_hwaccel_deps="libmfx"
diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index 7396468..71420fb 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -316,7 +316,7 @@ OBJS-$(CONFIG_H264_DECODER)            += h264dec.o h264_cabac.o h264_cavlc.o \
                                           h264_slice.o h264data.o h264_parse.o \
                                           h2645_parse.o
 OBJS-$(CONFIG_H264_CUVID_DECODER)      += cuvid.o
-OBJS-$(CONFIG_H264_MEDIACODEC_DECODER) += mediacodecdec_h264.o
+OBJS-$(CONFIG_H264_MEDIACODEC_DECODER) += mediacodecdec_h2645.o
 OBJS-$(CONFIG_H264_MMAL_DECODER)       += mmaldec.o
 OBJS-$(CONFIG_H264_NVENC_ENCODER)      += nvenc_h264.o
 OBJS-$(CONFIG_NVENC_ENCODER)           += nvenc_h264.o
@@ -333,6 +333,7 @@ OBJS-$(CONFIG_HEVC_DECODER)            += hevc.o hevc_mvs.o hevc_ps.o hevc_sei.o
                                           hevc_cabac.o hevc_refs.o hevcpred.o    \
                                           hevcdsp.o hevc_filter.o h2645_parse.o hevc_data.o
 OBJS-$(CONFIG_HEVC_CUVID_DECODER)      += cuvid.o
+OBJS-$(CONFIG_HEVC_MEDIACODEC_DECODER) += mediacodecdec_h2645.o hevc_parse.o
 OBJS-$(CONFIG_HEVC_NVENC_ENCODER)      += nvenc_hevc.o
 OBJS-$(CONFIG_NVENC_HEVC_ENCODER)      += nvenc_hevc.o
 OBJS-$(CONFIG_HEVC_QSV_DECODER)        += qsvdec_h2645.o
diff --git a/libavcodec/allcodecs.c b/libavcodec/allcodecs.c
index a26a80e..142ccba 100644
--- a/libavcodec/allcodecs.c
+++ b/libavcodec/allcodecs.c
@@ -84,6 +84,7 @@ void avcodec_register_all(void)
     REGISTER_HWACCEL(HEVC_CUVID,        hevc_cuvid);
     REGISTER_HWACCEL(HEVC_D3D11VA,      hevc_d3d11va);
     REGISTER_HWACCEL(HEVC_DXVA2,        hevc_dxva2);
+    REGISTER_HWACCEL(HEVC_MEDIACODEC,   hevc_mediacodec);
     REGISTER_HWACCEL(HEVC_QSV,          hevc_qsv);
     REGISTER_HWACCEL(HEVC_VAAPI,        hevc_vaapi);
     REGISTER_HWACCEL(HEVC_VDPAU,        hevc_vdpau);
@@ -644,6 +645,7 @@ void avcodec_register_all(void)
     REGISTER_ENCODER(NVENC_HEVC,        nvenc_hevc);
 #endif
     REGISTER_DECODER(HEVC_CUVID,        hevc_cuvid);
+    REGISTER_DECODER(HEVC_MEDIACODEC,   hevc_mediacodec);
     REGISTER_ENCODER(HEVC_NVENC,        hevc_nvenc);
     REGISTER_ENCODER(HEVC_QSV,          hevc_qsv);
     REGISTER_ENCODER(HEVC_VAAPI,        hevc_vaapi);
diff --git a/libavcodec/hevc_parse.c b/libavcodec/hevc_parse.c
new file mode 100644
index 0000000..cf04bc2
--- /dev/null
+++ b/libavcodec/hevc_parse.c
@@ -0,0 +1,134 @@
+/*
+ * 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 "bytestream.h"
+#include "h2645_parse.h"
+#include "hevc.h"
+#include "hevc_parse.h"
+
+static int hevc_decode_nal_units(const uint8_t *buf, int buf_size, HEVCParamSets *ps,
+                                 int is_nalff, int nal_length_size, void *logctx)
+{
+    int i;
+    int ret = 0;
+    H2645Packet pkt = { 0 };
+
+    ret = ff_h2645_packet_split(&pkt, buf, buf_size, logctx, is_nalff, nal_length_size, AV_CODEC_ID_HEVC, 1);
+    if (ret < 0) {
+        goto done;
+    }
+
+    for (i = 0; i < pkt.nb_nals; i++) {
+        H2645NAL *nal = &pkt.nals[i];
+
+        /* ignore everything except parameter sets and VCL NALUs */
+        switch (nal->type) {
+        case NAL_VPS: ff_hevc_decode_nal_vps(&nal->gb, logctx, ps);    break;
+        case NAL_SPS: ff_hevc_decode_nal_sps(&nal->gb, logctx, ps, 1); break;
+        case NAL_PPS: ff_hevc_decode_nal_pps(&nal->gb, logctx, ps);    break;
+        case NAL_TRAIL_R:
+        case NAL_TRAIL_N:
+        case NAL_TSA_N:
+        case NAL_TSA_R:
+        case NAL_STSA_N:
+        case NAL_STSA_R:
+        case NAL_BLA_W_LP:
+        case NAL_BLA_W_RADL:
+        case NAL_BLA_N_LP:
+        case NAL_IDR_W_RADL:
+        case NAL_IDR_N_LP:
+        case NAL_CRA_NUT:
+        case NAL_RADL_N:
+        case NAL_RADL_R:
+        case NAL_RASL_N:
+        case NAL_RASL_R:
+            av_log(logctx, AV_LOG_ERROR, "Invalid NAL unit: %d\n", nal->type);
+            ret = AVERROR_INVALIDDATA;
+            goto done;
+            break;
+        }
+    }
+
+done:
+    ff_h2645_packet_uninit(&pkt);
+    return ret;
+}
+
+int ff_hevc_decode_extradata(const uint8_t *data, int size, HEVCParamSets *ps,
+                             int *is_nalff, int *nal_length_size,
+                             int err_recognition, void *logctx)
+{
+    int ret = 0;
+    GetByteContext gb;
+
+    bytestream2_init(&gb, data, size);
+
+    if (size > 3 && (data[0] || data[1] || data[2] > 1)) {
+        /* It seems the extradata is encoded as hvcC format.
+         * Temporarily, we support configurationVersion==0 until 14496-15 3rd
+         * is finalized. When finalized, configurationVersion will be 1 and we
+         * can recognize hvcC by checking if avctx->extradata[0]==1 or not. */
+        int i, j, num_arrays, nal_len_size;
+
+        *is_nalff = 1;
+
+        bytestream2_skip(&gb, 21);
+        nal_len_size = (bytestream2_get_byte(&gb) & 3) + 1;
+        num_arrays   = bytestream2_get_byte(&gb);
+
+        /* nal units in the hvcC always have length coded with 2 bytes,
+         * so put a fake nal_length_size = 2 while parsing them */
+        *nal_length_size = 2;
+
+        /* Decode nal units from hvcC. */
+        for (i = 0; i < num_arrays; i++) {
+            int type = bytestream2_get_byte(&gb) & 0x3f;
+            int cnt  = bytestream2_get_be16(&gb);
+
+            for (j = 0; j < cnt; j++) {
+                // +2 for the nal size field
+                int nalsize = bytestream2_peek_be16(&gb) + 2;
+                if (bytestream2_get_bytes_left(&gb) < nalsize) {
+                    av_log(logctx, AV_LOG_ERROR,
+                           "Invalid NAL unit size in extradata.\n");
+                    return AVERROR_INVALIDDATA;
+                }
+
+                ret = hevc_decode_nal_units(gb.buffer, nalsize, ps, *is_nalff, *nal_length_size, logctx);
+                if (ret < 0) {
+                    av_log(logctx, AV_LOG_ERROR,
+                           "Decoding nal unit %d %d from hvcC failed\n",
+                           type, i);
+                    return ret;
+                }
+                bytestream2_skip(&gb, nalsize);
+            }
+        }
+
+        /* Now store right nal length size, that will be used to parse
+         * all other nals */
+        *nal_length_size = nal_len_size;
+    } else {
+        *is_nalff = 0;
+        ret = hevc_decode_nal_units(data, size, ps, *is_nalff, *nal_length_size, logctx);
+        if (ret < 0)
+            return ret;
+    }
+
+    return ret;
+}
diff --git a/libavcodec/hevc_parse.h b/libavcodec/hevc_parse.h
new file mode 100644
index 0000000..877356a
--- /dev/null
+++ b/libavcodec/hevc_parse.h
@@ -0,0 +1,33 @@
+/*
+ * 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
+ * H.265 parser code
+ */
+
+#ifndef AVCODEC_HEVC_PARSE_H
+#define AVCODEC_HEVC_PARSE_H
+
+#include "hevc.h"
+
+int ff_hevc_decode_extradata(const uint8_t *data, int size, HEVCParamSets *ps,
+                             int *is_nalff, int *nal_length_size,
+                             int err_recognition, void *logctx);
+
+#endif /* AVCODEC_HEVC_PARSE_H */
diff --git a/libavcodec/mediacodec_wrapper.c b/libavcodec/mediacodec_wrapper.c
index 36f38bb..5fc69b1 100644
--- a/libavcodec/mediacodec_wrapper.c
+++ b/libavcodec/mediacodec_wrapper.c
@@ -62,6 +62,10 @@ struct JNIAMediaCodecListFields {
     jfieldID avc_profile_high422_id;
     jfieldID avc_profile_high444_id;
 
+    jfieldID hevc_profile_main_id;
+    jfieldID hevc_profile_main10_id;
+    jfieldID hevc_profile_main10_hdr10_id;
+
 } JNIAMediaCodecListFields;
 
 static const struct FFJniField jni_amediacodeclist_mapping[] = {
@@ -94,6 +98,10 @@ static const struct FFJniField jni_amediacodeclist_mapping[] = {
         { "android/media/MediaCodecInfo$CodecProfileLevel", "AVCProfileHigh422", "I", FF_JNI_STATIC_FIELD, offsetof(struct JNIAMediaCodecListFields, avc_profile_high422_id), 1 },
         { "android/media/MediaCodecInfo$CodecProfileLevel", "AVCProfileHigh444", "I", FF_JNI_STATIC_FIELD, offsetof(struct JNIAMediaCodecListFields, avc_profile_high444_id), 1 },
 
+        { "android/media/MediaCodecInfo$CodecProfileLevel", "HEVCProfileMain", "I", FF_JNI_STATIC_FIELD, offsetof(struct JNIAMediaCodecListFields, hevc_profile_main_id), 0 },
+        { "android/media/MediaCodecInfo$CodecProfileLevel", "HEVCProfileMain10", "I", FF_JNI_STATIC_FIELD, offsetof(struct JNIAMediaCodecListFields, hevc_profile_main10_id), 0 },
+        { "android/media/MediaCodecInfo$CodecProfileLevel", "HEVCProfileMain10HDR10", "I", FF_JNI_STATIC_FIELD, offsetof(struct JNIAMediaCodecListFields, hevc_profile_main10_hdr10_id), 0 },
+
     { NULL }
 };
 
@@ -303,6 +311,7 @@ int ff_AMediaCodecProfile_getProfileFromAVCodecContext(AVCodecContext *avctx)
 
     JNIEnv *env = NULL;
     struct JNIAMediaCodecListFields jfields = { 0 };
+    jfieldID field_id = 0;
 
     JNI_GET_ENV_OR_RETURN(env, avctx, -1);
 
@@ -311,8 +320,6 @@ int ff_AMediaCodecProfile_getProfileFromAVCodecContext(AVCodecContext *avctx)
     }
 
     if (avctx->codec_id == AV_CODEC_ID_H264) {
-        jfieldID field_id = 0;
-
         switch(avctx->profile) {
         case FF_PROFILE_H264_BASELINE:
         case FF_PROFILE_H264_CONSTRAINED_BASELINE:
@@ -341,13 +348,23 @@ int ff_AMediaCodecProfile_getProfileFromAVCodecContext(AVCodecContext *avctx)
             field_id = jfields.avc_profile_high444_id;
             break;
         }
+    } else if (avctx->codec_id == AV_CODEC_ID_HEVC) {
+        switch(avctx->profile) {
+        case FF_PROFILE_HEVC_MAIN:
+        case FF_PROFILE_HEVC_MAIN_STILL_PICTURE:
+            field_id = jfields.hevc_profile_main_id;
+            break;
+        case FF_PROFILE_HEVC_MAIN_10:
+            field_id = jfields.hevc_profile_main10_id;
+            break;
+        }
+    }
 
-        if (field_id) {
-            ret = (*env)->GetStaticIntField(env, jfields.codec_profile_level_class, field_id);
-            if (ff_jni_exception_check(env, 1, avctx) < 0) {
-                ret = -1;
-                goto done;
-            }
+    if (field_id) {
+        ret = (*env)->GetStaticIntField(env, jfields.codec_profile_level_class, field_id);
+        if (ff_jni_exception_check(env, 1, avctx) < 0) {
+            ret = -1;
+            goto done;
         }
     }
 
diff --git a/libavcodec/mediacodecdec.c b/libavcodec/mediacodecdec.c
index df60104..0faa4cf 100644
--- a/libavcodec/mediacodecdec.c
+++ b/libavcodec/mediacodecdec.c
@@ -758,3 +758,10 @@ AVHWAccel ff_h264_mediacodec_hwaccel = {
     .id      = AV_CODEC_ID_H264,
     .pix_fmt = AV_PIX_FMT_MEDIACODEC,
 };
+
+AVHWAccel ff_hevc_mediacodec_hwaccel = {
+    .name    = "mediacodec",
+    .type    = AVMEDIA_TYPE_VIDEO,
+    .id      = AV_CODEC_ID_HEVC,
+    .pix_fmt = AV_PIX_FMT_MEDIACODEC,
+};
diff --git a/libavcodec/mediacodecdec_h264.c b/libavcodec/mediacodecdec_h2645.c
similarity index 68%
rename from libavcodec/mediacodecdec_h264.c
rename to libavcodec/mediacodecdec_h2645.c
index 4f9d737..d7761a5 100644
--- a/libavcodec/mediacodecdec_h264.c
+++ b/libavcodec/mediacodecdec_h2645.c
@@ -1,5 +1,5 @@
 /*
- * Android MediaCodec H.264 decoder
+ * Android MediaCodec H.264 / H.265 decoders
  *
  * Copyright (c) 2015-2016 Matthieu Bouron <matthieu.bouron stupeflix.com>
  *
@@ -33,12 +33,11 @@
 
 #include "avcodec.h"
 #include "h264_parse.h"
+#include "hevc_parse.h"
 #include "internal.h"
 #include "mediacodecdec.h"
 #include "mediacodec_wrapper.h"
 
-#define CODEC_MIME "video/avc"
-
 typedef struct MediaCodecH264DecContext {
 
     MediaCodecDecContext *ctx;
@@ -66,7 +65,7 @@ static av_cold int mediacodec_decode_close(AVCodecContext *avctx)
     return 0;
 }
 
-static int h264_ps_to_nalu(const uint8_t *src, int src_size, uint8_t **out, int *out_size)
+static int h2645_ps_to_nalu(const uint8_t *src, int src_size, uint8_t **out, int *out_size)
 {
     int i;
     int ret = 0;
@@ -118,7 +117,8 @@ done:
     return ret;
 }
 
-static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
+#if CONFIG_H264_MEDIACODEC_DECODER
+static int h264_set_extradata(AVCodecContext *avctx, FFAMediaFormat *format)
 {
     int i;
     int ret;
@@ -129,24 +129,8 @@ static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
     int is_avc = 0;
     int nal_length_size = 0;
 
-    const AVBitStreamFilter *bsf = NULL;
-
-    FFAMediaFormat *format = NULL;
-    MediaCodecH264DecContext *s = avctx->priv_data;
-
     memset(&ps, 0, sizeof(ps));
 
-    format = ff_AMediaFormat_new();
-    if (!format) {
-        av_log(avctx, AV_LOG_ERROR, "Failed to create media format\n");
-        ret = AVERROR_EXTERNAL;
-        goto done;
-    }
-
-    ff_AMediaFormat_setString(format, "mime", CODEC_MIME);
-    ff_AMediaFormat_setInt32(format, "width", avctx->width);
-    ff_AMediaFormat_setInt32(format, "height", avctx->height);
-
     ret = ff_h264_decode_extradata(avctx->extradata, avctx->extradata_size,
                                    &ps, &is_avc, &nal_length_size, 0, avctx);
     if (ret < 0) {
@@ -170,13 +154,13 @@ static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
         uint8_t *data = NULL;
         size_t data_size = 0;
 
-        if ((ret = h264_ps_to_nalu(sps->data, sps->data_size, &data, &data_size)) < 0) {
+        if ((ret = h2645_ps_to_nalu(sps->data, sps->data_size, &data, &data_size)) < 0) {
             goto done;
         }
         ff_AMediaFormat_setBuffer(format, "csd-0", (void*)data, data_size);
         av_freep(&data);
 
-        if ((ret = h264_ps_to_nalu(pps->data, pps->data_size, &data, &data_size)) < 0) {
+        if ((ret = h2645_ps_to_nalu(pps->data, pps->data_size, &data, &data_size)) < 0) {
             goto done;
         }
         ff_AMediaFormat_setBuffer(format, "csd-1", (void*)data, data_size);
@@ -184,9 +168,151 @@ static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
     } else {
         av_log(avctx, AV_LOG_ERROR, "Could not extract PPS/SPS from extradata");
         ret = AVERROR_INVALIDDATA;
+    }
+
+done:
+    ff_h264_ps_uninit(&ps);
+
+    return ret;
+}
+#endif
+
+#if CONFIG_HEVC_MEDIACODEC_DECODER
+static int hevc_set_extradata(AVCodecContext *avctx, FFAMediaFormat *format)
+{
+    int i;
+    int ret;
+
+    HEVCParamSets ps;
+
+    const HEVCVPS *vps = NULL;
+    const HEVCPPS *pps = NULL;
+    const HEVCSPS *sps = NULL;
+    int is_nalff = 0;
+    int nal_length_size = 0;
+
+    uint8_t *vps_data = NULL;
+    uint8_t *sps_data = NULL;
+    uint8_t *pps_data = NULL;
+    int vps_data_size = 0;
+    int sps_data_size = 0;
+    int pps_data_size = 0;
+
+    memset(&ps, 0, sizeof(ps));
+
+    ret = ff_hevc_decode_extradata(avctx->extradata, avctx->extradata_size,
+                                &ps, &is_nalff, &nal_length_size, 0, avctx);
+    if (ret < 0) {
+        goto done;
+    }
+
+
+    for (i = 0; i < MAX_VPS_COUNT; i++) {
+        if (ps.vps_list[i]) {
+            vps = (const HEVCVPS*)ps.vps_list[i]->data;
+            break;
+        }
+    }
+
+    for (i = 0; i < MAX_PPS_COUNT; i++) {
+        if (ps.pps_list[i]) {
+            pps = (const HEVCPPS*)ps.pps_list[i]->data;
+            break;
+        }
+    }
+
+    if (pps) {
+        if (ps.sps_list[pps->sps_id]) {
+            sps = (const HEVCSPS*)ps.sps_list[pps->sps_id]->data;
+        }
+    }
+
+    if (vps && pps && sps) {
+        uint8_t *data;
+        int data_size;
+
+        if ((ret = h2645_ps_to_nalu(vps->data, vps->data_size, &vps_data, &vps_data_size)) < 0 ||
+            (ret = h2645_ps_to_nalu(sps->data, sps->data_size, &sps_data, &sps_data_size)) < 0 ||
+            (ret = h2645_ps_to_nalu(pps->data, pps->data_size, &pps_data, &pps_data_size)) < 0) {
+            goto done;
+        }
+
+        data_size = vps_data_size + sps_data_size + pps_data_size;
+        data = av_mallocz(data_size);
+        if (!data) {
+            ret = AVERROR(ENOMEM);
+            goto done;
+        }
+
+        memcpy(data                                , vps_data, vps_data_size);
+        memcpy(data + vps_data_size                , sps_data, sps_data_size);
+        memcpy(data + vps_data_size + sps_data_size, pps_data, pps_data_size);
+
+        ff_AMediaFormat_setBuffer(format, "csd-0", data, data_size);
+
+        av_freep(&data);
+    } else {
+        av_log(avctx, AV_LOG_ERROR, "Could not extract VPS/PPS/SPS from extradata");
+        ret = AVERROR_INVALIDDATA;
+    }
+
+done:
+    av_freep(&vps_data);
+    av_freep(&sps_data);
+    av_freep(&pps_data);
+
+    return ret;
+}
+#endif
+
+static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
+{
+    int ret;
+
+    const char *codec_mime = NULL;
+
+    const char *bsf_name = NULL;
+    const AVBitStreamFilter *bsf = NULL;
+
+    FFAMediaFormat *format = NULL;
+    MediaCodecH264DecContext *s = avctx->priv_data;
+
+    format = ff_AMediaFormat_new();
+    if (!format) {
+        av_log(avctx, AV_LOG_ERROR, "Failed to create media format\n");
+        ret = AVERROR_EXTERNAL;
         goto done;
     }
 
+    switch(avctx->codec_id) {
+#if CONFIG_H264_MEDIACODEC_DECODER
+    case AV_CODEC_ID_H264:
+        codec_mime = "video/avc";
+        bsf_name = "h264_mp4toannexb";
+
+        ret = h264_set_extradata(avctx, format);
+        if (ret < 0)
+            goto done;
+        break;
+#endif
+#if CONFIG_HEVC_MEDIACODEC_DECODER
+    case AV_CODEC_ID_HEVC:
+        codec_mime = "video/hevc";
+        bsf_name = "hevc_mp4toannexb";
+
+        ret = hevc_set_extradata(avctx, format);
+        if (ret < 0)
+            goto done;
+        break;
+#endif
+    default:
+        av_assert0(0);
+    }
+
+    ff_AMediaFormat_setString(format, "mime", codec_mime);
+    ff_AMediaFormat_setInt32(format, "width", avctx->width);
+    ff_AMediaFormat_setInt32(format, "height", avctx->height);
+
     s->ctx = av_mallocz(sizeof(*s->ctx));
     if (!s->ctx) {
         av_log(avctx, AV_LOG_ERROR, "Failed to allocate MediaCodecDecContext\n");
@@ -194,7 +320,7 @@ static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
         goto done;
     }
 
-    if ((ret = ff_mediacodec_dec_init(avctx, s->ctx, CODEC_MIME, format)) < 0) {
+    if ((ret = ff_mediacodec_dec_init(avctx, s->ctx, codec_mime, format)) < 0) {
         s->ctx = NULL;
         goto done;
     }
@@ -207,7 +333,7 @@ static av_cold int mediacodec_decode_init(AVCodecContext *avctx)
         goto done;
     }
 
-    bsf = av_bsf_get_by_name("h264_mp4toannexb");
+    bsf = av_bsf_get_by_name(bsf_name);
     if(!bsf) {
         ret = AVERROR_BSF_NOT_FOUND;
         goto done;
@@ -233,8 +359,6 @@ done:
         mediacodec_decode_close(avctx);
     }
 
-    ff_h264_ps_uninit(&ps);
-
     return ret;
 }
 
@@ -323,7 +447,7 @@ static int mediacodec_decode_frame(AVCodecContext *avctx, void *data,
                 goto done;
             }
 
-            /* h264_mp4toannexb is used here and does not requires flushing */
+            /* {h264,hevc}_mp4toannexb are used here and do not require flushing */
             av_assert0(ret != AVERROR_EOF);
 
             if (ret < 0) {
@@ -358,6 +482,7 @@ static void mediacodec_decode_flush(AVCodecContext *avctx)
     ff_mediacodec_dec_flush(avctx, s->ctx);
 }
 
+#if CONFIG_H264_MEDIACODEC_DECODER
 AVCodec ff_h264_mediacodec_decoder = {
     .name           = "h264_mediacodec",
     .long_name      = NULL_IF_CONFIG_SMALL("H.264 Android MediaCodec decoder"),
@@ -371,3 +496,20 @@ AVCodec ff_h264_mediacodec_decoder = {
     .capabilities   = CODEC_CAP_DELAY,
     .caps_internal  = FF_CODEC_CAP_SETS_PKT_DTS,
 };
+#endif
+
+#if CONFIG_HEVC_MEDIACODEC_DECODER
+AVCodec ff_hevc_mediacodec_decoder = {
+    .name           = "hevc_mediacodec",
+    .long_name      = NULL_IF_CONFIG_SMALL("H.265 Android MediaCodec decoder"),
+    .type           = AVMEDIA_TYPE_VIDEO,
+    .id             = AV_CODEC_ID_HEVC,
+    .priv_data_size = sizeof(MediaCodecH264DecContext),
+    .init           = mediacodec_decode_init,
+    .decode         = mediacodec_decode_frame,
+    .flush          = mediacodec_decode_flush,
+    .close          = mediacodec_decode_close,
+    .capabilities   = CODEC_CAP_DELAY,
+    .caps_internal  = FF_CODEC_CAP_SETS_PKT_DTS,
+};
+#endif
-- 
2.9.3

