avcodec/hevc: Add alpha layer support

Signed-off-by: Zhao Zhili <zhilizhao@tencent.com>

libavcodec/hevc/hevcdec.c

@@ -466,6 +466,25 @@ static int export_multilayer(HEVCContext *s, const HEVCVPS *vps)
     return 0;
 }
 
+int ff_hevc_is_alpha_video(const HEVCContext *s)
+{
+    const HEVCVPS *vps = s->vps;
+    int ret = 0;
+
+    if (vps->nb_layers != 2 || !vps->layer_id_in_nuh[1])
+        return 0;
+
+    /* decode_vps_ext() guarantees that SCALABILITY_AUXILIARY with AuxId other
+     * than alpha cannot reach here.
+     */
+    ret = (s->vps->scalability_mask_flag & HEVC_SCALABILITY_AUXILIARY);
+
+    av_log(s->avctx, AV_LOG_DEBUG, "Multi layer video, %s alpha video\n",
+           ret ? "is" : "not");
+
+    return ret;
+}
+
 static int setup_multilayer(HEVCContext *s, const HEVCVPS *vps)
 {
     unsigned layers_active_output = 0, highest_layer;
@@ -473,6 +492,18 @@ static int setup_multilayer(HEVCContext *s, const HEVCVPS *vps)
     s->layers_active_output = 1;
     s->layers_active_decode = 1;
 
+    if (ff_hevc_is_alpha_video(s)) {
+        const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(s->avctx->pix_fmt);
+
+        if (!(desc->flags & AV_PIX_FMT_FLAG_ALPHA))
+            return 0;
+
+        s->layers_active_decode = (1 << vps->nb_layers) - 1;
+        s->layers_active_output = 1;
+
+        return 0;
+    }
+
     // nothing requested - decode base layer only
     if (!s->nb_view_ids)
         return 0;
@@ -530,6 +561,34 @@ static int setup_multilayer(HEVCContext *s, const HEVCVPS *vps)
     return 0;
 }
 
+static enum AVPixelFormat map_to_alpha_format(HEVCContext *s,
+                                              enum AVPixelFormat pix_fmt)
+{
+    switch (pix_fmt) {
+    case AV_PIX_FMT_YUV420P:
+    case AV_PIX_FMT_YUVJ420P:
+        return AV_PIX_FMT_YUVA420P;
+    case AV_PIX_FMT_YUV420P10:
+        return AV_PIX_FMT_YUVA420P10;
+    case AV_PIX_FMT_YUV444P:
+        return AV_PIX_FMT_YUVA444P;
+    case AV_PIX_FMT_YUV422P:
+        return AV_PIX_FMT_YUVA422P;
+    case AV_PIX_FMT_YUV422P10LE:
+        return AV_PIX_FMT_YUVA422P10LE;
+    case AV_PIX_FMT_YUV444P10:
+        return AV_PIX_FMT_YUVA444P10;
+    case AV_PIX_FMT_YUV444P12:
+        return AV_PIX_FMT_YUVA444P12;
+    case AV_PIX_FMT_YUV422P12:
+        return AV_PIX_FMT_YUVA422P12;
+    default:
+        av_log(s->avctx, AV_LOG_WARNING, "No alpha pixel format map for %s\n",
+               av_get_pix_fmt_name(pix_fmt));
+        return AV_PIX_FMT_NONE;
+    }
+}
+
 static enum AVPixelFormat get_format(HEVCContext *s, const HEVCSPS *sps)
 {
 #define HWACCEL_MAX (CONFIG_HEVC_DXVA2_HWACCEL + \
@@ -540,9 +599,13 @@ static enum AVPixelFormat get_format(HEVCContext *s, const HEVCSPS *sps)
                      CONFIG_HEVC_VIDEOTOOLBOX_HWACCEL + \
                      CONFIG_HEVC_VDPAU_HWACCEL + \
                      CONFIG_HEVC_VULKAN_HWACCEL)
-    enum AVPixelFormat pix_fmts[HWACCEL_MAX + 2], *fmt = pix_fmts;
+    enum AVPixelFormat pix_fmts[HWACCEL_MAX + 3], *fmt = pix_fmts;
+    enum AVPixelFormat alpha_fmt = AV_PIX_FMT_NONE;
     int ret;
 
+    if (ff_hevc_is_alpha_video(s))
+        alpha_fmt = map_to_alpha_format(s, sps->pix_fmt);
+
     switch (sps->pix_fmt) {
     case AV_PIX_FMT_YUV420P:
     case AV_PIX_FMT_YUVJ420P:
@@ -664,6 +727,8 @@ static enum AVPixelFormat get_format(HEVCContext *s, const HEVCSPS *sps)
         break;
     }
 
+    if (alpha_fmt != AV_PIX_FMT_NONE)
+        *fmt++ = alpha_fmt;
     *fmt++ = sps->pix_fmt;
     *fmt = AV_PIX_FMT_NONE;
 
@@ -3194,6 +3259,12 @@ static int hevc_frame_start(HEVCContext *s, HEVCLayerContext *l,
                 !sps->vui.common.video_signal_type_present_flag)
                 pix_fmt = sps_base->pix_fmt;
 
+            // Ignore range mismatch between base layer and alpha layer
+            if (ff_hevc_is_alpha_video(s) &&
+                sps_base->pix_fmt == AV_PIX_FMT_YUV420P &&
+                pix_fmt == AV_PIX_FMT_YUVJ420P)
+                pix_fmt = sps_base->pix_fmt;
+
             if (pix_fmt     != sps_base->pix_fmt ||
                 sps->width  != sps_base->width   ||
                 sps->height != sps_base->height) {

libavcodec/hevc/hevcdec.h

@@ -714,6 +714,8 @@ void ff_hevc_hls_residual_coding(HEVCLocalContext *lc, const HEVCPPS *pps,
 
 void ff_hevc_hls_mvd_coding(HEVCLocalContext *lc, int x0, int y0, int log2_cb_size);
 
+int ff_hevc_is_alpha_video(const HEVCContext *s);
+
 extern const uint8_t ff_hevc_qpel_extra_before[4];
 extern const uint8_t ff_hevc_qpel_extra_after[4];
 extern const uint8_t ff_hevc_qpel_extra[4];

libavcodec/hevc/refs.c

@@ -79,6 +79,31 @@ void ff_hevc_flush_dpb(HEVCContext *s)
     }
 }
 
+static int replace_alpha_plane(AVFrame *alpha, AVFrame *base)
+{
+    AVBufferRef *base_a = av_frame_get_plane_buffer(base, 3);
+    uintptr_t data = (uintptr_t)alpha->data[0];
+    int ret;
+
+    for (int i = 0; i < FF_ARRAY_ELEMS(alpha->buf) && alpha->buf[i]; i++) {
+        AVBufferRef *buf = alpha->buf[i];
+        uintptr_t buf_begin = (uintptr_t)buf->data;
+
+        if (data >=  buf_begin && data < buf_begin + buf->size) {
+            ret = av_buffer_replace(&alpha->buf[i], base_a);
+            if (ret < 0)
+                return ret;
+
+            alpha->linesize[0] = base->linesize[3];
+            alpha->data[0] = base->data[3];
+
+            return 0;
+        }
+    }
+
+    return AVERROR_BUG;
+}
+
 static HEVCFrame *alloc_frame(HEVCContext *s, HEVCLayerContext *l)
 {
     const HEVCVPS *vps = l->sps->vps;
@@ -103,7 +128,7 @@ static HEVCFrame *alloc_frame(HEVCContext *s, HEVCLayerContext *l)
         }
 
         // add view ID side data if it's nontrivial
-        if (vps->nb_layers > 1 || view_id) {
+        if (!ff_hevc_is_alpha_video(s) && (vps->nb_layers > 1 || view_id)) {
             HEVCSEITDRDI *tdrdi = &s->sei.tdrdi;
             AVFrameSideData *sd = av_frame_side_data_new(&frame->f->side_data,
                                                          &frame->f->nb_side_data,
@@ -162,6 +187,13 @@ static HEVCFrame *alloc_frame(HEVCContext *s, HEVCLayerContext *l)
             goto fail;
 
         frame->pps = av_refstruct_ref_c(s->pps);
+        if (l != &s->layers[0] && ff_hevc_is_alpha_video(s)) {
+            AVFrame *alpha = frame->f;
+            AVFrame *base = s->layers[0].cur_frame->f;
+            ret = replace_alpha_plane(alpha, base);
+            if (ret < 0)
+                goto fail;
+        }
 
         return frame;
 fail: