prores: interlaced ProRes encoding

Signed-off-by: Luca Barbato <lu_zero@gentoo.org>

1 files changed, 115 insertions, 68 deletions

diff --git a/libavcodec/proresenc.c b/libavcodec/proresenc.c
index 40644bb3b1..a24b7118c2 100644
--- a/libavcodec/proresenc.c
+++ b/libavcodec/proresenc.c
@@ -188,7 +188,9 @@ typedef struct ProresContext {
     int mbs_per_slice;
     int num_chroma_blocks, chroma_factor;
     int slices_width;
-    int num_slices;
+    int slices_per_picture;
+    int pictures_per_frame; // 1 for progressive, 2 for interlaced
+    int cur_picture_idx;
     int num_planes;
     int bits_per_mb;
     int force_quant;
@@ -196,7 +198,7 @@ typedef struct ProresContext {
     char *vendor;
     int quant_sel;
 
-    int frame_size;
+    int frame_size_upper_bound;
 
     int profile;
     const struct prores_profile *profile_info;
@@ -402,10 +404,15 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
     int total_size = 0;
     const uint16_t *src;
     int slice_width_factor = av_log2(mbs_per_slice);
-    int num_cblocks, pwidth;
+    int num_cblocks, pwidth, linesize, line_add;
     int plane_factor, is_chroma;
     uint16_t *qmat;
 
+    if (ctx->pictures_per_frame == 1)
+        line_add = 0;
+    else
+        line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
+
     if (ctx->force_quant) {
         qmat = ctx->quants[0];
     } else if (quant < MAX_STORED_Q) {
@@ -432,12 +439,16 @@ static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
             num_cblocks = 2;
             pwidth      = avctx->width >> 1;
         }
-        src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
 
-        get_slice_data(ctx, src, pic->linesize[i], xp, yp,
-                       pwidth, avctx->height, ctx->blocks[0], ctx->emu_buf,
+        linesize = pic->linesize[i] * ctx->pictures_per_frame;
+        src = (const uint16_t*)(pic->data[i] + yp * linesize +
+                                line_add * pic->linesize[i]) + xp;
+
+        get_slice_data(ctx, src, linesize, xp, yp,
+                       pwidth, avctx->height / ctx->pictures_per_frame,
+                       ctx->blocks[0], ctx->emu_buf,
                        mbs_per_slice, num_cblocks, is_chroma);
-        sizes[i] = encode_slice_plane(ctx, pb, src, pic->linesize[i],
+        sizes[i] = encode_slice_plane(ctx, pb, src, linesize,
                                       mbs_per_slice, ctx->blocks[0],
                                       num_cblocks, plane_factor,
                                       qmat);
@@ -570,7 +581,12 @@ static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
     int overquant;
     uint16_t *qmat;
+    int linesize[4], line_add;
 
+    if (ctx->pictures_per_frame == 1)
+        line_add = 0;
+    else
+        line_add = ctx->cur_picture_idx ^ !pic->top_field_first;
     mbs = x + mbs_per_slice;
 
     for (i = 0; i < ctx->num_planes; i++) {
@@ -589,10 +605,14 @@ static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
             num_cblocks[i] = 2;
             pwidth         = avctx->width >> 1;
         }
-        src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
 
-        get_slice_data(ctx, src, pic->linesize[i], xp, yp,
-                       pwidth, avctx->height, td->blocks[i], td->emu_buf,
+        linesize[i] = pic->linesize[i] * ctx->pictures_per_frame;
+        src = (const uint16_t*)(pic->data[i] + yp * linesize[i] +
+                                line_add * pic->linesize[i]) + xp;
+
+        get_slice_data(ctx, src, linesize[i], xp, yp,
+                       pwidth, avctx->height / ctx->pictures_per_frame,
+                       td->blocks[i], td->emu_buf,
                        mbs_per_slice, num_cblocks[i], is_chroma[i]);
     }
 
@@ -607,7 +627,7 @@ static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
         error = 0;
         for (i = 0; i < ctx->num_planes; i++) {
             bits += estimate_slice_plane(ctx, &error, i,
-                                         src, pic->linesize[i],
+                                         src, linesize[i],
                                          mbs_per_slice,
                                          num_cblocks[i], plane_factor[i],
                                          ctx->quants[q], td);
@@ -636,7 +656,7 @@ static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
             }
             for (i = 0; i < ctx->num_planes; i++) {
                 bits += estimate_slice_plane(ctx, &error, i,
-                                             src, pic->linesize[i],
+                                             src, linesize[i],
                                              mbs_per_slice,
                                              num_cblocks[i], plane_factor[i],
                                              qmat, td);
@@ -725,12 +745,13 @@ static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
     int slice_hdr_size = 2 + 2 * (ctx->num_planes - 1);
     int frame_size, picture_size, slice_size;
     int pkt_size, ret;
+    uint8_t frame_flags;
 
     *avctx->coded_frame           = *pic;
     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
     avctx->coded_frame->key_frame = 1;
 
-    pkt_size = ctx->frame_size + FF_MIN_BUFFER_SIZE;
+    pkt_size = ctx->frame_size_upper_bound + FF_MIN_BUFFER_SIZE;
 
     if ((ret = ff_alloc_packet(pkt, pkt_size)) < 0) {
         av_log(avctx, AV_LOG_ERROR, "Error getting output packet.\n");
@@ -751,7 +772,12 @@ static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
     bytestream_put_buffer(&buf, ctx->vendor, 4);
     bytestream_put_be16  (&buf, avctx->width);
     bytestream_put_be16  (&buf, avctx->height);
-    bytestream_put_byte  (&buf, ctx->chroma_factor << 6); // frame flags
+
+    frame_flags = ctx->chroma_factor << 6;
+    if (avctx->flags & CODEC_FLAG_INTERLACED_DCT)
+        frame_flags |= pic->top_field_first ? 0x04 : 0x08;
+    bytestream_put_byte  (&buf, frame_flags);
+
     bytestream_put_byte  (&buf, 0);             // reserved
     bytestream_put_byte  (&buf, avctx->color_primaries);
     bytestream_put_byte  (&buf, avctx->color_trc);
@@ -771,56 +797,64 @@ static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
     }
     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
 
-    // picture header
-    picture_size_pos = buf + 1;
-    bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
-    buf += 4;                                   // picture data size will be stored here
-    bytestream_put_be16  (&buf, ctx->num_slices); // total number of slices
-    bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
-
-    // seek table - will be filled during slice encoding
-    slice_sizes = buf;
-    buf += ctx->num_slices * 2;
-
-    // slices
-    if (!ctx->force_quant) {
-        ret = avctx->execute2(avctx, find_quant_thread, NULL, NULL,
-                              ctx->mb_height);
-        if (ret)
-            return ret;
-    }
+    for (ctx->cur_picture_idx = 0;
+         ctx->cur_picture_idx < ctx->pictures_per_frame;
+         ctx->cur_picture_idx++) {
+        // picture header
+        picture_size_pos = buf + 1;
+        bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
+        buf += 4;                                   // picture data size will be stored here
+        bytestream_put_be16  (&buf, ctx->slices_per_picture);
+        bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
+
+        // seek table - will be filled during slice encoding
+        slice_sizes = buf;
+        buf += ctx->slices_per_picture * 2;
+
+        // slices
+        if (!ctx->force_quant) {
+            ret = avctx->execute2(avctx, find_quant_thread, NULL, NULL,
+                                  ctx->mb_height);
+            if (ret)
+                return ret;
+        }
 
-    for (y = 0; y < ctx->mb_height; y++) {
-        int mbs_per_slice = ctx->mbs_per_slice;
-        for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
-            q = ctx->force_quant ? ctx->force_quant
-                                 : ctx->slice_q[mb + y * ctx->slices_width];
-
-            while (ctx->mb_width - x < mbs_per_slice)
-                mbs_per_slice >>= 1;
-
-            bytestream_put_byte(&buf, slice_hdr_size << 3);
-            slice_hdr = buf;
-            buf += slice_hdr_size - 1;
-            init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
-            encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
-
-            bytestream_put_byte(&slice_hdr, q);
-            slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
-            for (i = 0; i < ctx->num_planes - 1; i++) {
-                bytestream_put_be16(&slice_hdr, sizes[i]);
-                slice_size += sizes[i];
+        for (y = 0; y < ctx->mb_height; y++) {
+            int mbs_per_slice = ctx->mbs_per_slice;
+            for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
+                q = ctx->force_quant ? ctx->force_quant
+                                     : ctx->slice_q[mb + y * ctx->slices_width];
+
+                while (ctx->mb_width - x < mbs_per_slice)
+                    mbs_per_slice >>= 1;
+
+                bytestream_put_byte(&buf, slice_hdr_size << 3);
+                slice_hdr = buf;
+                buf += slice_hdr_size - 1;
+                init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
+                encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
+
+                bytestream_put_byte(&slice_hdr, q);
+                slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
+                for (i = 0; i < ctx->num_planes - 1; i++) {
+                    bytestream_put_be16(&slice_hdr, sizes[i]);
+                    slice_size += sizes[i];
+                }
+                bytestream_put_be16(&slice_sizes, slice_size);
+                buf += slice_size - slice_hdr_size;
             }
-            bytestream_put_be16(&slice_sizes, slice_size);
-            buf += slice_size - slice_hdr_size;
         }
+
+        if (ctx->pictures_per_frame == 1)
+            picture_size = buf - picture_size_pos - 6;
+        else
+            picture_size = buf - picture_size_pos + 1;
+        bytestream_put_be32(&picture_size_pos, picture_size);
     }
 
     orig_buf -= 8;
     frame_size = buf - orig_buf;
-    picture_size = buf - picture_size_pos - 6;
     bytestream_put_be32(&orig_buf, frame_size);
-    bytestream_put_be32(&picture_size_pos, picture_size);
 
     pkt->size   = frame_size;
     pkt->flags |= AV_PKT_FLAG_KEY;
@@ -855,6 +889,7 @@ static av_cold int encode_init(AVCodecContext *avctx)
     int mps;
     int i, j;
     int min_quant, max_quant;
+    int interlaced = !!(avctx->flags & CODEC_FLAG_INTERLACED_DCT);
 
     avctx->bits_per_raw_sample = 10;
     avctx->coded_frame = avcodec_alloc_frame();
@@ -863,7 +898,8 @@ static av_cold int encode_init(AVCodecContext *avctx)
 
     ff_proresdsp_init(&ctx->dsp);
     ff_init_scantable(ctx->dsp.dct_permutation, &ctx->scantable,
-                      ff_prores_progressive_scan);
+                      interlaced ? ff_prores_interlaced_scan
+                                 : ff_prores_progressive_scan);
 
     mps = ctx->mbs_per_slice;
     if (mps & (mps - 1)) {
@@ -879,10 +915,16 @@ static av_cold int encode_init(AVCodecContext *avctx)
     ctx->num_planes    = 3;
 
     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
-    ctx->mb_height     = FFALIGN(avctx->height, 16) >> 4;
+
+    if (interlaced)
+        ctx->mb_height = FFALIGN(avctx->height, 32) >> 5;
+    else
+        ctx->mb_height = FFALIGN(avctx->height, 16) >> 4;
+
     ctx->slices_width  = ctx->mb_width / mps;
     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
-    ctx->num_slices    = ctx->mb_height * ctx->slices_width;
+    ctx->slices_per_picture = ctx->mb_height * ctx->slices_width;
+    ctx->pictures_per_frame = 1 + interlaced;
 
     if (ctx->quant_sel == -1)
         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
@@ -898,7 +940,8 @@ static av_cold int encode_init(AVCodecContext *avctx)
     if (!ctx->force_quant) {
         if (!ctx->bits_per_mb) {
             for (i = 0; i < NUM_MB_LIMITS - 1; i++)
-                if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height)
+                if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height *
+                                           ctx->pictures_per_frame)
                     break;
             ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
         } else if (ctx->bits_per_mb < 128) {
@@ -913,7 +956,7 @@ static av_cold int encode_init(AVCodecContext *avctx)
                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
         }
 
-        ctx->slice_q = av_malloc(ctx->num_slices * sizeof(*ctx->slice_q));
+        ctx->slice_q = av_malloc(ctx->slices_per_picture * sizeof(*ctx->slice_q));
         if (!ctx->slice_q) {
             encode_close(avctx);
             return AVERROR(ENOMEM);
@@ -959,16 +1002,20 @@ static av_cold int encode_init(AVCodecContext *avctx)
             ctx->bits_per_mb += ls * 4;
     }
 
-    ctx->frame_size = ctx->num_slices * (2 + 2 * ctx->num_planes
-                                         + (2 * mps * ctx->bits_per_mb) / 8)
-                      + 200;
+    ctx->frame_size_upper_bound = ctx->pictures_per_frame *
+                                  ctx->slices_per_picture *
+                                  (2 + 2 * ctx->num_planes +
+                                   (mps * ctx->bits_per_mb) / 8)
+                                  + 200;
 
     avctx->codec_tag   = ctx->profile_info->tag;
 
-    av_log(avctx, AV_LOG_DEBUG, "profile %d, %d slices, %d bits per MB\n",
-           ctx->profile, ctx->num_slices, ctx->bits_per_mb);
-    av_log(avctx, AV_LOG_DEBUG, "estimated frame size %d\n",
-           ctx->frame_size);
+    av_log(avctx, AV_LOG_DEBUG,
+           "profile %d, %d slices, interlacing: %s, %d bits per MB\n",
+           ctx->profile, ctx->slices_per_picture * ctx->pictures_per_frame,
+           interlaced ? "yes" : "no", ctx->bits_per_mb);
+    av_log(avctx, AV_LOG_DEBUG, "frame size upper bound: %d\n",
+           ctx->frame_size_upper_bound);
 
     return 0;
 }