prores: decode alpha plane when it's present

1 files changed, 103 insertions, 9 deletions

diff --git a/libavcodec/proresdec.c b/libavcodec/proresdec.c
index c42e444153..0d1eb7ab38 100644
--- a/libavcodec/proresdec.c
+++ b/libavcodec/proresdec.c
@@ -134,12 +134,21 @@ static int decode_frame_header(ProresContext *ctx, const uint8_t *buf,
     ctx->chroma_factor     = (buf[12] >> 6) & 3;
     ctx->mb_chroma_factor  = ctx->chroma_factor + 2;
     ctx->num_chroma_blocks = (1 << ctx->chroma_factor) >> 1;
+    ctx->alpha_info        = buf[17] & 0xf;
+
+    if (ctx->alpha_info > 2) {
+        av_log(avctx, AV_LOG_ERROR, "Invalid alpha mode %d\n", ctx->alpha_info);
+        return AVERROR_INVALIDDATA;
+    }
+
     switch (ctx->chroma_factor) {
     case 2:
-        avctx->pix_fmt = AV_PIX_FMT_YUV422P10;
+        avctx->pix_fmt = ctx->alpha_info ? AV_PIX_FMT_YUVA422P10
+                                         : AV_PIX_FMT_YUV422P10;
         break;
     case 3:
-        avctx->pix_fmt = AV_PIX_FMT_YUV444P10;
+        avctx->pix_fmt = ctx->alpha_info ? AV_PIX_FMT_YUVA444P10
+                                         : AV_PIX_FMT_YUV444P10;
         break;
     default:
         av_log(avctx, AV_LOG_ERROR,
@@ -168,10 +177,6 @@ static int decode_frame_header(ProresContext *ctx, const uint8_t *buf,
     avctx->color_trc       = buf[15];
     avctx->colorspace      = buf[16];
 
-    ctx->alpha_info = buf[17] & 0xf;
-    if (ctx->alpha_info)
-        avpriv_report_missing_feature(avctx, "Alpha channel");
-
     ctx->qmat_changed = 0;
     ptr   = buf + 20;
     flags = buf[19];
@@ -466,6 +471,78 @@ static void decode_slice_plane(ProresContext *ctx, ProresThreadData *td,
 }
 
 
+static void unpack_alpha(GetBitContext *gb, uint16_t *dst, int num_coeffs,
+                         const int num_bits)
+{
+    const int mask = (1 << num_bits) - 1;
+    int i, idx, val, alpha_val;
+
+    idx       = 0;
+    alpha_val = mask;
+    do {
+        do {
+            if (get_bits1(gb))
+                val = get_bits(gb, num_bits);
+            else {
+                int sign;
+                val  = get_bits(gb, num_bits == 16 ? 7 : 4);
+                sign = val & 1;
+                val  = (val + 2) >> 1;
+                if (sign)
+                    val = -val;
+            }
+            alpha_val = (alpha_val + val) & mask;
+            if (num_bits == 16)
+                dst[idx++] = alpha_val >> 6;
+            else
+                dst[idx++] = (alpha_val << 2) | (alpha_val >> 6);
+            if (idx == num_coeffs - 1)
+                break;
+        } while (get_bits1(gb));
+        val = get_bits(gb, 4);
+        if (!val)
+            val = get_bits(gb, 11);
+        if (idx + val > num_coeffs)
+            val = num_coeffs - idx;
+        if (num_bits == 16)
+            for (i = 0; i < val; i++)
+                dst[idx++] = alpha_val >> 6;
+        else
+            for (i = 0; i < val; i++)
+                dst[idx++] = (alpha_val << 2) | (alpha_val >> 6);
+    } while (idx < num_coeffs);
+}
+
+/**
+ * Decode alpha slice plane.
+ */
+static void decode_alpha_plane(ProresContext *ctx, ProresThreadData *td,
+                               const uint8_t *buf, int data_size,
+                               uint16_t *out_ptr, int linesize,
+                               int mbs_per_slice)
+{
+    GetBitContext gb;
+    int i;
+    uint16_t *block_ptr;
+
+    memset(td->blocks, 0, 8 * 4 * 64 * sizeof(*td->blocks));
+
+    init_get_bits(&gb, buf, data_size << 3);
+
+    if (ctx->alpha_info == 2)
+        unpack_alpha(&gb, td->blocks, mbs_per_slice * 4 * 64, 16);
+    else
+        unpack_alpha(&gb, td->blocks, mbs_per_slice * 4 * 64, 8);
+
+    block_ptr = td->blocks;
+
+    for (i = 0; i < 16; i++) {
+        memcpy(out_ptr, block_ptr, 16 * mbs_per_slice * sizeof(*out_ptr));
+        out_ptr   += linesize >> 1;
+        block_ptr += 16 * mbs_per_slice;
+    }
+}
+
 static int decode_slice(AVCodecContext *avctx, void *tdata)
 {
     ProresThreadData *td = tdata;
@@ -476,11 +553,12 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
     int slice_num = td->slice_num;
     int mbs_per_slice = td->slice_width;
     const uint8_t *buf;
-    uint8_t *y_data, *u_data, *v_data;
+    uint8_t *y_data, *u_data, *v_data, *a_data;
     AVFrame *pic = ctx->frame;
     int i, sf, slice_width_factor;
-    int slice_data_size, hdr_size, y_data_size, u_data_size, v_data_size;
-    int y_linesize, u_linesize, v_linesize;
+    int slice_data_size, hdr_size;
+    int y_data_size, u_data_size, v_data_size, a_data_size;
+    int y_linesize, u_linesize, v_linesize, a_linesize;
 
     buf             = ctx->slice_data[slice_num].index;
     slice_data_size = ctx->slice_data[slice_num + 1].index - buf;
@@ -490,19 +568,23 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
     y_data     = pic->data[0];
     u_data     = pic->data[1];
     v_data     = pic->data[2];
+    a_data     = pic->data[3];
     y_linesize = pic->linesize[0];
     u_linesize = pic->linesize[1];
     v_linesize = pic->linesize[2];
+    a_linesize = pic->linesize[3];
 
     if (pic->interlaced_frame) {
         if (!(pic_num ^ pic->top_field_first)) {
             y_data += y_linesize;
             u_data += u_linesize;
             v_data += v_linesize;
+            a_data += a_linesize;
         }
         y_linesize <<= 1;
         u_linesize <<= 1;
         v_linesize <<= 1;
+        a_linesize <<= 1;
     }
 
     if (slice_data_size < 6) {
@@ -516,6 +598,8 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
     u_data_size = AV_RB16(buf + 4);
     v_data_size = hdr_size > 7 ? AV_RB16(buf + 6) :
         slice_data_size - y_data_size - u_data_size - hdr_size;
+    a_data_size = slice_data_size - y_data_size - u_data_size -
+                  v_data_size - hdr_size;
 
     if (hdr_size + y_data_size + u_data_size + v_data_size > slice_data_size ||
         v_data_size < 0 || hdr_size < 6) {
@@ -560,6 +644,16 @@ static int decode_slice(AVCodecContext *avctx, void *tdata)
                        slice_width_factor + ctx->chroma_factor - 1,
                        td->qmat_chroma_scaled, 1);
 
+    /* decode alpha plane if available */
+    if (a_data && a_data_size)
+        decode_alpha_plane(ctx, td,
+                           buf + hdr_size + y_data_size +
+                               u_data_size + v_data_size,
+                           a_data_size,
+                           (uint16_t*) (a_data + (mb_y_pos << 4) * a_linesize +
+                                        (mb_x_pos << 5)), a_linesize,
+                           mbs_per_slice);
+
     return 0;
 }