libavcodec/dnxhdenc: add support for dnxhr encoding

Signed-off-by: Michael Niedermayer <michael@niedermayer.cc>

4 files changed, 103 insertions, 19 deletions

diff --git a/libavcodec/dnxhddata.c b/libavcodec/dnxhddata.c
index 7d935a3f31..480b8069cf 100644
--- a/libavcodec/dnxhddata.c
+++ b/libavcodec/dnxhddata.c
@@ -1047,35 +1047,35 @@ const CIDEntry ff_dnxhd_cid_table[] = {
       dnxhd_1235_dc_codes, dnxhd_1235_dc_bits,
       dnxhd_1235_ac_codes, dnxhd_1235_ac_bits, dnxhd_1235_ac_info,
       dnxhd_1235_run_codes, dnxhd_1235_run_bits, dnxhd_1235_run,
-      { 0 } },
+      { 0 }, { { 0 } }, { 57344, 255} },
     { 1271, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE,
       0, 6, DNXHD_VARIABLE, 4,
       dnxhd_1241_luma_weight, dnxhd_1241_chroma_weight,
       dnxhd_1235_dc_codes, dnxhd_1235_dc_bits,
       dnxhd_1235_ac_codes, dnxhd_1235_ac_bits, dnxhd_1235_ac_info,
       dnxhd_1235_run_codes, dnxhd_1235_run_bits, dnxhd_1235_run,
-      { 0 } },
+      { 0 }, { { 0 } }, { 28672, 255} },
     { 1272, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE,
       0, 4, 8, 4,
       dnxhd_1238_luma_weight, dnxhd_1238_chroma_weight,
       dnxhd_1237_dc_codes, dnxhd_1237_dc_bits,
       dnxhd_1238_ac_codes, dnxhd_1238_ac_bits, dnxhd_1238_ac_info,
       dnxhd_1235_run_codes, dnxhd_1235_run_bits, dnxhd_1238_run,
-      { 0 } },
+      { 0 }, { { 0 } }, { 28672, 255} },
     { 1273, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE,
       0, 4, 8, 3,
       dnxhd_1237_luma_weight, dnxhd_1237_chroma_weight,
       dnxhd_1237_dc_codes, dnxhd_1237_dc_bits,
       dnxhd_1237_ac_codes, dnxhd_1237_ac_bits, dnxhd_1237_ac_info,
       dnxhd_1237_run_codes, dnxhd_1237_run_bits, dnxhd_1237_run,
-      { 0 } },
+      { 0 }, { { 0 } }, { 18944, 255} },
     { 1274, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE, DNXHD_VARIABLE,
       0, 4, 8, 3,
       dnxhd_1237_luma_weight, dnxhd_1237_chroma_weight,
       dnxhd_1237_dc_codes, dnxhd_1237_dc_bits,
       dnxhd_1237_ac_codes, dnxhd_1237_ac_bits, dnxhd_1237_ac_info,
       dnxhd_1237_run_codes, dnxhd_1237_run_bits, dnxhd_1237_run,
-      { 0 } },
+      { 0 }, { { 0 } }, { 5888, 255} },
 };
 
 int ff_dnxhd_get_cid_table(int cid)
@@ -1110,10 +1110,31 @@ uint64_t avpriv_dnxhd_parse_header_prefix(const uint8_t *buf)
     return ff_dnxhd_check_header_prefix(prefix);
 }
 
+static int dnxhd_find_hr_cid(AVCodecContext *avctx)
+{
+    switch (avctx->profile) {
+    case FF_PROFILE_DNXHR_444:
+        return 1270;
+    case FF_PROFILE_DNXHR_HQX:
+        return 1271;
+    case FF_PROFILE_DNXHR_HQ:
+        return 1272;
+    case FF_PROFILE_DNXHR_SQ:
+        return 1273;
+    case FF_PROFILE_DNXHR_LB:
+        return 1274;
+    }
+    return 0;
+}
+
 int ff_dnxhd_find_cid(AVCodecContext *avctx, int bit_depth)
 {
     int i, j;
     int mbs = avctx->bit_rate / 1000000;
+
+    if (avctx->profile != FF_PROFILE_DNXHD)
+        return dnxhd_find_hr_cid(avctx);
+
     if (!mbs)
         return 0;
     for (i = 0; i < FF_ARRAY_ELEMS(ff_dnxhd_cid_table); i++) {
diff --git a/libavcodec/dnxhddata.h b/libavcodec/dnxhddata.h
index 3806f9b4f1..06e7128d4d 100644
--- a/libavcodec/dnxhddata.h
+++ b/libavcodec/dnxhddata.h
@@ -55,6 +55,7 @@ typedef struct CIDEntry {
     const uint8_t *run_bits, *run;
     int bit_rates[5]; ///< Helper to choose variants, rounded to nearest 5Mb/s
     AVRational frame_rates[5];
+    AVRational packet_scale;
 } CIDEntry;
 
 extern const CIDEntry ff_dnxhd_cid_table[];
diff --git a/libavcodec/dnxhdenc.c b/libavcodec/dnxhdenc.c
index b3721ccf38..b0ee8a2050 100644
--- a/libavcodec/dnxhdenc.c
+++ b/libavcodec/dnxhdenc.c
@@ -49,6 +49,21 @@ static const AVOption options[] = {
     { "ibias", "intra quant bias",
         offsetof(DNXHDEncContext, intra_quant_bias), AV_OPT_TYPE_INT,
         { .i64 = 0 }, INT_MIN, INT_MAX, VE },
+    { "profile",       NULL, offsetof(DNXHDEncContext, profile), AV_OPT_TYPE_INT,
+        { .i64 = FF_PROFILE_DNXHD },
+        FF_PROFILE_DNXHD, FF_PROFILE_DNXHR_444, VE, "profile" },
+    { "dnxhd",     NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHD },
+        0, 0, VE, "profile" },
+    { "dnxhr_444", NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHR_444 },
+        0, 0, VE, "profile" },
+    { "dnxhr_hqx", NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHR_HQX },
+        0, 0, VE, "profile" },
+    { "dnxhr_hq",  NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHR_HQ },
+        0, 0, VE, "profile" },
+    { "dnxhr_sq",  NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHR_SQ },
+        0, 0, VE, "profile" },
+    { "dnxhr_lb",  NULL, 0, AV_OPT_TYPE_CONST, { .i64 = FF_PROFILE_DNXHR_LB },
+        0, 0, VE, "profile" },
     { NULL }
 };
 
@@ -276,8 +291,8 @@ static av_cold int dnxhd_init_rc(DNXHDEncContext *ctx)
     if (ctx->m.avctx->mb_decision != FF_MB_DECISION_RD)
         FF_ALLOCZ_ARRAY_OR_GOTO(ctx->m.avctx, ctx->mb_cmp,
                           ctx->m.mb_num, sizeof(RCCMPEntry), fail);
-    ctx->frame_bits = (ctx->cid_table->coding_unit_size -
-                       640 - 4 - ctx->min_padding) * 8;
+    ctx->frame_bits = (ctx->coding_unit_size -
+                       ctx->data_offset - 4 - ctx->min_padding) * 8;
     ctx->qscale = 1;
     ctx->lambda = 2 << LAMBDA_FRAC_BITS; // qscale 2
     return 0;
@@ -285,6 +300,12 @@ fail:
     return AVERROR(ENOMEM);
 }
 
+static int dnxhd_get_hr_frame_size(const CIDEntry* profile, int mb_num)
+{
+  int result = mb_num * profile->packet_scale.num / profile->packet_scale.den;
+  result = (result + 2048) / 4096 * 4096;
+  return FFMAX(result, 8192);
+}
 static av_cold int dnxhd_encode_init(AVCodecContext *avctx)
 {
     DNXHDEncContext *ctx = avctx->priv_data;
@@ -303,6 +324,14 @@ static av_cold int dnxhd_encode_init(AVCodecContext *avctx)
         return AVERROR(EINVAL);
     }
 
+    if (ctx->profile == FF_PROFILE_DNXHR_HQX ||
+        ctx->profile == FF_PROFILE_DNXHR_HQX) {
+        avpriv_report_missing_feature(avctx,
+               "dnxhr_444 or dnxhr_hqx profile");
+        return AVERROR_PATCHWELCOME;
+    }
+
+    avctx->profile = ctx->profile;
     ctx->cid = ff_dnxhd_find_cid(avctx, bit_depth);
     if (!ctx->cid) {
         av_log(avctx, AV_LOG_ERROR,
@@ -312,6 +341,15 @@ static av_cold int dnxhd_encode_init(AVCodecContext *avctx)
     }
     av_log(avctx, AV_LOG_DEBUG, "cid %d\n", ctx->cid);
 
+    if (ctx->cid >= 1270 && ctx->cid <= 1274)
+        avctx->codec_tag = MKTAG('A','V','d','h');
+
+    if (avctx->width < 256 || avctx->height < 120) {
+        av_log(avctx, AV_LOG_ERROR,
+               "Input dimensions too small, input must be at least 256x120\n");
+        return AVERROR(EINVAL);
+    }
+
     index = ff_dnxhd_get_cid_table(ctx->cid);
     av_assert0(index >= 0);
 
@@ -355,6 +393,20 @@ static av_cold int dnxhd_encode_init(AVCodecContext *avctx)
 
     ctx->m.mb_num = ctx->m.mb_height * ctx->m.mb_width;
 
+    if (ctx->cid_table->frame_size == DNXHD_VARIABLE) {
+        ctx->frame_size = dnxhd_get_hr_frame_size(ctx->cid_table,
+                                                  ctx->m.mb_num);
+        ctx->coding_unit_size = ctx->frame_size;
+    } else {
+        ctx->frame_size = ctx->cid_table->frame_size;
+        ctx->coding_unit_size = ctx->cid_table->coding_unit_size;
+    }
+
+    if (ctx->m.mb_height > 68)
+        ctx->data_offset = 0x170 + (ctx->m.mb_height << 2);
+    else
+        ctx->data_offset = 0x280;
+
 #if FF_API_QUANT_BIAS
 FF_DISABLE_DEPRECATION_WARNINGS
     if (avctx->intra_quant_bias != FF_DEFAULT_QUANT_BIAS)
@@ -415,11 +467,16 @@ fail:  // for FF_ALLOCZ_OR_GOTO
 static int dnxhd_write_header(AVCodecContext *avctx, uint8_t *buf)
 {
     DNXHDEncContext *ctx = avctx->priv_data;
-    static const uint8_t header_prefix[5] = { 0x00, 0x00, 0x02, 0x80, 0x01 };
 
-    memset(buf, 0, 640);
+    memset(buf, 0, ctx->data_offset);
+
+    // * write prefix */
+    AV_WB16(buf + 0x02, ctx->data_offset);
+    if (ctx->cid >= 1270 && ctx->cid <= 1274)
+        buf[4] = 0x03;
+    else
+        buf[4] = 0x01;
 
-    memcpy(buf, header_prefix, 5);
     buf[5] = ctx->interlaced ? ctx->cur_field + 2 : 0x01;
     buf[6] = 0x80; // crc flag off
     buf[7] = 0xa0; // reserved
@@ -573,7 +630,7 @@ void dnxhd_get_blocks(DNXHDEncContext *ctx, int mb_x, int mb_y)
     pdsp->get_pixels(ctx->blocks[2], ptr_u,      ctx->m.uvlinesize);
     pdsp->get_pixels(ctx->blocks[3], ptr_v,      ctx->m.uvlinesize);
 
-    if (mb_y + 1 == ctx->m.mb_height && ctx->m.avctx->height == 1080) {
+    if (mb_y + 1 == ctx->m.mb_height && (ctx->m.avctx->height % 16) != 0) {
         if (ctx->interlaced) {
             ctx->get_pixels_8x4_sym(ctx->blocks[4],
                                     ptr_y + ctx->dct_y_offset,
@@ -674,7 +731,7 @@ static int dnxhd_encode_thread(AVCodecContext *avctx, void *arg,
     DNXHDEncContext *ctx = avctx->priv_data;
     int mb_y = jobnr, mb_x;
     ctx = ctx->thread[threadnr];
-    init_put_bits(&ctx->m.pb, (uint8_t *)arg + 640 + ctx->slice_offs[jobnr],
+    init_put_bits(&ctx->m.pb, (uint8_t *)arg + ctx->data_offset + ctx->slice_offs[jobnr],
                   ctx->slice_size[jobnr]);
 
     ctx->m.last_dc[0] =
@@ -740,7 +797,7 @@ static int dnxhd_mb_var_thread(AVCodecContext *avctx, void *arg,
             int sum;
             int varc;
 
-            if (!partial_last_row && mb_x * 16 <= avctx->width - 16) {
+            if (!partial_last_row && mb_x * 16 <= avctx->width - 16 && (avctx->width % 16) == 0) {
                 sum  = ctx->m.mpvencdsp.pix_sum(pix, ctx->m.linesize);
                 varc = ctx->m.mpvencdsp.pix_norm1(pix, ctx->m.linesize);
             } else {
@@ -1063,7 +1120,7 @@ static int dnxhd_encode_picture(AVCodecContext *avctx, AVPacket *pkt,
     int offset, i, ret;
     uint8_t *buf;
 
-    if ((ret = ff_alloc_packet2(avctx, pkt, ctx->cid_table->frame_size, 0)) < 0)
+    if ((ret = ff_alloc_packet2(avctx, pkt, ctx->frame_size, 0)) < 0)
         return ret;
     buf = pkt->data;
 
@@ -1099,16 +1156,16 @@ encode_coding_unit:
 
     avctx->execute2(avctx, dnxhd_encode_thread, buf, NULL, ctx->m.mb_height);
 
-    av_assert1(640 + offset + 4 <= ctx->cid_table->coding_unit_size);
-    memset(buf + 640 + offset, 0,
-           ctx->cid_table->coding_unit_size - 4 - offset - 640);
+    av_assert1(ctx->data_offset + offset + 4 <= ctx->coding_unit_size);
+    memset(buf + ctx->data_offset + offset, 0,
+           ctx->coding_unit_size - 4 - offset - ctx->data_offset);
 
-    AV_WB32(buf + ctx->cid_table->coding_unit_size - 4, 0x600DC0DE); // EOF
+    AV_WB32(buf + ctx->coding_unit_size - 4, 0x600DC0DE); // EOF
 
     if (ctx->interlaced && first_field) {
         first_field     = 0;
         ctx->cur_field ^= 1;
-        buf            += ctx->cid_table->coding_unit_size;
+        buf            += ctx->coding_unit_size;
         goto encode_coding_unit;
     }
 
diff --git a/libavcodec/dnxhdenc.h b/libavcodec/dnxhdenc.h
index e9d94294dc..7f8e7b35f2 100644
--- a/libavcodec/dnxhdenc.h
+++ b/libavcodec/dnxhdenc.h
@@ -45,6 +45,7 @@ typedef struct DNXHDEncContext {
     MpegEncContext m; ///< Used for quantization dsp functions
 
     int cid;
+    int profile;
     const CIDEntry *cid_table;
     uint8_t *msip; ///< Macroblock Scan Indexes Payload
     uint32_t *slice_size;
@@ -58,6 +59,10 @@ typedef struct DNXHDEncContext {
     unsigned dct_uv_offset;
     unsigned block_width_l2;
 
+    int frame_size;
+    int coding_unit_size;
+    int data_offset;
+
     int interlaced;
     int cur_field;