git.sesse.net Git - ffmpeg/blob - libavcodec/proresenc_kostya.c

   1 /*
   2  * Apple ProRes encoder
   3  *
   4  * Copyright (c) 2012 Konstantin Shishkov
   5  *
   6  * This file is part of Libav.
   7  *
   8  * Libav is free software; you can redistribute it and/or
   9  * modify it under the terms of the GNU Lesser General Public
  10  * License as published by the Free Software Foundation; either
  11  * version 2.1 of the License, or (at your option) any later version.
  12  *
  13  * Libav is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16  * Lesser General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU Lesser General Public
  19  * License along with Libav; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  21  */
  22
  23 #include "libavutil/opt.h"
  24 #include "avcodec.h"
  25 #include "put_bits.h"
  26 #include "bytestream.h"
  27 #include "internal.h"
  28 #include "proresdsp.h"
  29 #include "proresdata.h"
  30
  31 #define CFACTOR_Y422 2
  32 #define CFACTOR_Y444 3
  33
  34 #define MAX_MBS_PER_SLICE 8
  35
  36 #define MAX_PLANES 3 // should be increased to 4 when there's PIX_FMT_YUV444AP10
  37
  38 enum {
  39     PRORES_PROFILE_PROXY = 0,
  40     PRORES_PROFILE_LT,
  41     PRORES_PROFILE_STANDARD,
  42     PRORES_PROFILE_HQ,
  43 };
  44
  45 enum {
  46     QUANT_MAT_PROXY = 0,
  47     QUANT_MAT_LT,
  48     QUANT_MAT_STANDARD,
  49     QUANT_MAT_HQ,
  50     QUANT_MAT_DEFAULT,
  51 };
  52
  53 static const uint8_t prores_quant_matrices[][64] = {
  54     { // proxy
  55          4,  7,  9, 11, 13, 14, 15, 63,
  56          7,  7, 11, 12, 14, 15, 63, 63,
  57          9, 11, 13, 14, 15, 63, 63, 63,
  58         11, 11, 13, 14, 63, 63, 63, 63,
  59         11, 13, 14, 63, 63, 63, 63, 63,
  60         13, 14, 63, 63, 63, 63, 63, 63,
  61         13, 63, 63, 63, 63, 63, 63, 63,
  62         63, 63, 63, 63, 63, 63, 63, 63,
  63     },
  64     { // LT
  65          4,  5,  6,  7,  9, 11, 13, 15,
  66          5,  5,  7,  8, 11, 13, 15, 17,
  67          6,  7,  9, 11, 13, 15, 15, 17,
  68          7,  7,  9, 11, 13, 15, 17, 19,
  69          7,  9, 11, 13, 14, 16, 19, 23,
  70          9, 11, 13, 14, 16, 19, 23, 29,
  71          9, 11, 13, 15, 17, 21, 28, 35,
  72         11, 13, 16, 17, 21, 28, 35, 41,
  73     },
  74     { // standard
  75          4,  4,  5,  5,  6,  7,  7,  9,
  76          4,  4,  5,  6,  7,  7,  9,  9,
  77          5,  5,  6,  7,  7,  9,  9, 10,
  78          5,  5,  6,  7,  7,  9,  9, 10,
  79          5,  6,  7,  7,  8,  9, 10, 12,
  80          6,  7,  7,  8,  9, 10, 12, 15,
  81          6,  7,  7,  9, 10, 11, 14, 17,
  82          7,  7,  9, 10, 11, 14, 17, 21,
  83     },
  84     { // high quality
  85          4,  4,  4,  4,  4,  4,  4,  4,
  86          4,  4,  4,  4,  4,  4,  4,  4,
  87          4,  4,  4,  4,  4,  4,  4,  4,
  88          4,  4,  4,  4,  4,  4,  4,  5,
  89          4,  4,  4,  4,  4,  4,  5,  5,
  90          4,  4,  4,  4,  4,  5,  5,  6,
  91          4,  4,  4,  4,  5,  5,  6,  7,
  92          4,  4,  4,  4,  5,  6,  7,  7,
  93     },
  94     { // codec default
  95          4,  4,  4,  4,  4,  4,  4,  4,
  96          4,  4,  4,  4,  4,  4,  4,  4,
  97          4,  4,  4,  4,  4,  4,  4,  4,
  98          4,  4,  4,  4,  4,  4,  4,  4,
  99          4,  4,  4,  4,  4,  4,  4,  4,
 100          4,  4,  4,  4,  4,  4,  4,  4,
 101          4,  4,  4,  4,  4,  4,  4,  4,
 102          4,  4,  4,  4,  4,  4,  4,  4,
 103     },
 104 };
 105
 106 #define NUM_MB_LIMITS 4
 107 static const int prores_mb_limits[NUM_MB_LIMITS] = {
 108     1620, // up to 720x576
 109     2700, // up to 960x720
 110     6075, // up to 1440x1080
 111     9216, // up to 2048x1152
 112 };
 113
 114 static const struct prores_profile {
 115     const char *full_name;
 116     uint32_t    tag;
 117     int         min_quant;
 118     int         max_quant;
 119     int         br_tab[NUM_MB_LIMITS];
 120     int         quant;
 121 } prores_profile_info[4] = {
 122     {
 123         .full_name = "proxy",
 124         .tag       = MKTAG('a', 'p', 'c', 'o'),
 125         .min_quant = 4,
 126         .max_quant = 8,
 127         .br_tab    = { 300, 242, 220, 194 },
 128         .quant     = QUANT_MAT_PROXY,
 129     },
 130     {
 131         .full_name = "LT",
 132         .tag       = MKTAG('a', 'p', 'c', 's'),
 133         .min_quant = 1,
 134         .max_quant = 9,
 135         .br_tab    = { 720, 560, 490, 440 },
 136         .quant     = QUANT_MAT_LT,
 137     },
 138     {
 139         .full_name = "standard",
 140         .tag       = MKTAG('a', 'p', 'c', 'n'),
 141         .min_quant = 1,
 142         .max_quant = 6,
 143         .br_tab    = { 1050, 808, 710, 632 },
 144         .quant     = QUANT_MAT_STANDARD,
 145     },
 146     {
 147         .full_name = "high quality",
 148         .tag       = MKTAG('a', 'p', 'c', 'h'),
 149         .min_quant = 1,
 150         .max_quant = 6,
 151         .br_tab    = { 1566, 1216, 1070, 950 },
 152         .quant     = QUANT_MAT_HQ,
 153     }
 154 // for 4444 profile bitrate numbers are { 2350, 1828, 1600, 1425 }
 155 };
 156
 157 #define TRELLIS_WIDTH 16
 158 #define SCORE_LIMIT   INT_MAX / 2
 159
 160 struct TrellisNode {
 161     int prev_node;
 162     int quant;
 163     int bits;
 164     int score;
 165 };
 166
 167 #define MAX_STORED_Q 16
 168
 169 typedef struct ProresContext {
 170     AVClass *class;
 171     DECLARE_ALIGNED(16, DCTELEM, blocks)[MAX_PLANES][64 * 4 * MAX_MBS_PER_SLICE];
 172     DECLARE_ALIGNED(16, uint16_t, emu_buf)[16*16];
 173     int16_t quants[MAX_STORED_Q][64];
 174     int16_t custom_q[64];
 175     const uint8_t *quant_mat;
 176
 177     ProresDSPContext dsp;
 178     ScanTable  scantable;
 179
 180     int mb_width, mb_height;
 181     int mbs_per_slice;
 182     int num_chroma_blocks, chroma_factor;
 183     int slices_width;
 184     int num_slices;
 185     int num_planes;
 186     int bits_per_mb;
 187     int force_quant;
 188
 189     char *vendor;
 190     int quant_sel;
 191
 192     int frame_size;
 193
 194     int profile;
 195     const struct prores_profile *profile_info;
 196
 197     struct TrellisNode *nodes;
 198     int *slice_q;
 199 } ProresContext;
 200
 201 static void get_slice_data(ProresContext *ctx, const uint16_t *src,
 202                            int linesize, int x, int y, int w, int h,
 203                            DCTELEM *blocks,
 204                            int mbs_per_slice, int blocks_per_mb, int is_chroma)
 205 {
 206     const uint16_t *esrc;
 207     const int mb_width = 4 * blocks_per_mb;
 208     int elinesize;
 209     int i, j, k;
 210
 211     for (i = 0; i < mbs_per_slice; i++, src += mb_width) {
 212         if (x >= w) {
 213             memset(blocks, 0, 64 * (mbs_per_slice - i) * blocks_per_mb
 214                               * sizeof(*blocks));
 215             return;
 216         }
 217         if (x + mb_width <= w && y + 16 <= h) {
 218             esrc      = src;
 219             elinesize = linesize;
 220         } else {
 221             int bw, bh, pix;
 222
 223             esrc      = ctx->emu_buf;
 224             elinesize = 16 * sizeof(*ctx->emu_buf);
 225
 226             bw = FFMIN(w - x, mb_width);
 227             bh = FFMIN(h - y, 16);
 228
 229             for (j = 0; j < bh; j++) {
 230                 memcpy(ctx->emu_buf + j * 16,
 231                        (const uint8_t*)src + j * linesize,
 232                        bw * sizeof(*src));
 233                 pix = ctx->emu_buf[j * 16 + bw - 1];
 234                 for (k = bw; k < mb_width; k++)
 235                     ctx->emu_buf[j * 16 + k] = pix;
 236             }
 237             for (; j < 16; j++)
 238                 memcpy(ctx->emu_buf + j * 16,
 239                        ctx->emu_buf + (bh - 1) * 16,
 240                        mb_width * sizeof(*ctx->emu_buf));
 241         }
 242         if (!is_chroma) {
 243             ctx->dsp.fdct(esrc, elinesize, blocks);
 244             blocks += 64;
 245             if (blocks_per_mb > 2) {
 246                 ctx->dsp.fdct(src + 8, linesize, blocks);
 247                 blocks += 64;
 248             }
 249             ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
 250             blocks += 64;
 251             if (blocks_per_mb > 2) {
 252                 ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
 253                 blocks += 64;
 254             }
 255         } else {
 256             ctx->dsp.fdct(esrc, elinesize, blocks);
 257             blocks += 64;
 258             ctx->dsp.fdct(src + linesize * 4, linesize, blocks);
 259             blocks += 64;
 260             if (blocks_per_mb > 2) {
 261                 ctx->dsp.fdct(src + 8, linesize, blocks);
 262                 blocks += 64;
 263                 ctx->dsp.fdct(src + linesize * 4 + 8, linesize, blocks);
 264                 blocks += 64;
 265             }
 266         }
 267
 268         x += mb_width;
 269     }
 270 }
 271
 272 /**
 273  * Write an unsigned rice/exp golomb codeword.
 274  */
 275 static inline void encode_vlc_codeword(PutBitContext *pb, unsigned codebook, int val)
 276 {
 277     unsigned int rice_order, exp_order, switch_bits, switch_val;
 278     int exponent;
 279
 280     /* number of prefix bits to switch between Rice and expGolomb */
 281     switch_bits = (codebook & 3) + 1;
 282     rice_order  =  codebook >> 5;       /* rice code order */
 283     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
 284
 285     switch_val  = switch_bits << rice_order;
 286
 287     if (val >= switch_val) {
 288         val -= switch_val - (1 << exp_order);
 289         exponent = av_log2(val);
 290
 291         put_bits(pb, exponent - exp_order + switch_bits, 0);
 292         put_bits(pb, 1, 1);
 293         put_bits(pb, exponent, val);
 294     } else {
 295         exponent = val >> rice_order;
 296
 297         if (exponent)
 298             put_bits(pb, exponent, 0);
 299         put_bits(pb, 1, 1);
 300         if (rice_order)
 301             put_sbits(pb, rice_order, val);
 302     }
 303 }
 304
 305 #define GET_SIGN(x)  ((x) >> 31)
 306 #define MAKE_CODE(x) (((x) << 1) ^ GET_SIGN(x))
 307
 308 static void encode_dcs(PutBitContext *pb, DCTELEM *blocks,
 309                        int blocks_per_slice, int scale)
 310 {
 311     int i;
 312     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
 313
 314     prev_dc = (blocks[0] - 0x4000) / scale;
 315     encode_vlc_codeword(pb, FIRST_DC_CB, MAKE_CODE(prev_dc));
 316     sign     = 0;
 317     codebook = 3;
 318     blocks  += 64;
 319
 320     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
 321         dc       = (blocks[0] - 0x4000) / scale;
 322         delta    = dc - prev_dc;
 323         new_sign = GET_SIGN(delta);
 324         delta    = (delta ^ sign) - sign;
 325         code     = MAKE_CODE(delta);
 326         encode_vlc_codeword(pb, ff_prores_dc_codebook[codebook], code);
 327         codebook = (code + (code & 1)) >> 1;
 328         codebook = FFMIN(codebook, 3);
 329         sign     = new_sign;
 330         prev_dc  = dc;
 331     }
 332 }
 333
 334 static void encode_acs(PutBitContext *pb, DCTELEM *blocks,
 335                        int blocks_per_slice,
 336                        int plane_size_factor,
 337                        const uint8_t *scan, const int16_t *qmat)
 338 {
 339     int idx, i;
 340     int run, level, run_cb, lev_cb;
 341     int max_coeffs, abs_level;
 342
 343     max_coeffs = blocks_per_slice << 6;
 344     run_cb     = ff_prores_run_to_cb_index[4];
 345     lev_cb     = ff_prores_lev_to_cb_index[2];
 346     run        = 0;
 347
 348     for (i = 1; i < 64; i++) {
 349         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
 350             level = blocks[idx] / qmat[scan[i]];
 351             if (level) {
 352                 abs_level = FFABS(level);
 353                 encode_vlc_codeword(pb, ff_prores_ac_codebook[run_cb], run);
 354                 encode_vlc_codeword(pb, ff_prores_ac_codebook[lev_cb],
 355                                     abs_level - 1);
 356                 put_sbits(pb, 1, GET_SIGN(level));
 357
 358                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
 359                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
 360                 run    = 0;
 361             } else {
 362                 run++;
 363             }
 364         }
 365     }
 366 }
 367
 368 static int encode_slice_plane(ProresContext *ctx, PutBitContext *pb,
 369                               const uint16_t *src, int linesize,
 370                               int mbs_per_slice, DCTELEM *blocks,
 371                               int blocks_per_mb, int plane_size_factor,
 372                               const int16_t *qmat)
 373 {
 374     int blocks_per_slice, saved_pos;
 375
 376     saved_pos = put_bits_count(pb);
 377     blocks_per_slice = mbs_per_slice * blocks_per_mb;
 378
 379     encode_dcs(pb, blocks, blocks_per_slice, qmat[0]);
 380     encode_acs(pb, blocks, blocks_per_slice, plane_size_factor,
 381                ctx->scantable.permutated, qmat);
 382     flush_put_bits(pb);
 383
 384     return (put_bits_count(pb) - saved_pos) >> 3;
 385 }
 386
 387 static int encode_slice(AVCodecContext *avctx, const AVFrame *pic,
 388                         PutBitContext *pb,
 389                         int sizes[4], int x, int y, int quant,
 390                         int mbs_per_slice)
 391 {
 392     ProresContext *ctx = avctx->priv_data;
 393     int i, xp, yp;
 394     int total_size = 0;
 395     const uint16_t *src;
 396     int slice_width_factor = av_log2(mbs_per_slice);
 397     int num_cblocks, pwidth;
 398     int plane_factor, is_chroma;
 399     uint16_t *qmat;
 400
 401     if (ctx->force_quant) {
 402         qmat = ctx->quants[0];
 403     } else if (quant < MAX_STORED_Q) {
 404         qmat = ctx->quants[quant];
 405     } else {
 406         qmat = ctx->custom_q;
 407         for (i = 0; i < 64; i++)
 408             qmat[i] = ctx->quant_mat[i] * quant;
 409     }
 410
 411     for (i = 0; i < ctx->num_planes; i++) {
 412         is_chroma    = (i == 1 || i == 2);
 413         plane_factor = slice_width_factor + 2;
 414         if (is_chroma)
 415             plane_factor += ctx->chroma_factor - 3;
 416         if (!is_chroma || ctx->chroma_factor == CFACTOR_Y444) {
 417             xp          = x << 4;
 418             yp          = y << 4;
 419             num_cblocks = 4;
 420             pwidth      = avctx->width;
 421         } else {
 422             xp          = x << 3;
 423             yp          = y << 4;
 424             num_cblocks = 2;
 425             pwidth      = avctx->width >> 1;
 426         }
 427         src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
 428
 429         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
 430                        pwidth, avctx->height, ctx->blocks[0],
 431                        mbs_per_slice, num_cblocks, is_chroma);
 432         sizes[i] = encode_slice_plane(ctx, pb, src, pic->linesize[i],
 433                                       mbs_per_slice, ctx->blocks[0],
 434                                       num_cblocks, plane_factor,
 435                                       qmat);
 436         total_size += sizes[i];
 437     }
 438     return total_size;
 439 }
 440
 441 static inline int estimate_vlc(unsigned codebook, int val)
 442 {
 443     unsigned int rice_order, exp_order, switch_bits, switch_val;
 444     int exponent;
 445
 446     /* number of prefix bits to switch between Rice and expGolomb */
 447     switch_bits = (codebook & 3) + 1;
 448     rice_order  =  codebook >> 5;       /* rice code order */
 449     exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */
 450
 451     switch_val  = switch_bits << rice_order;
 452
 453     if (val >= switch_val) {
 454         val -= switch_val - (1 << exp_order);
 455         exponent = av_log2(val);
 456
 457         return exponent * 2 - exp_order + switch_bits + 1;
 458     } else {
 459         return (val >> rice_order) + rice_order + 1;
 460     }
 461 }
 462
 463 static int estimate_dcs(int *error, DCTELEM *blocks, int blocks_per_slice,
 464                         int scale)
 465 {
 466     int i;
 467     int codebook = 3, code, dc, prev_dc, delta, sign, new_sign;
 468     int bits;
 469
 470     prev_dc  = (blocks[0] - 0x4000) / scale;
 471     bits     = estimate_vlc(FIRST_DC_CB, MAKE_CODE(prev_dc));
 472     sign     = 0;
 473     codebook = 3;
 474     blocks  += 64;
 475     *error  += FFABS(blocks[0] - 0x4000) % scale;
 476
 477     for (i = 1; i < blocks_per_slice; i++, blocks += 64) {
 478         dc       = (blocks[0] - 0x4000) / scale;
 479         *error  += FFABS(blocks[0] - 0x4000) % scale;
 480         delta    = dc - prev_dc;
 481         new_sign = GET_SIGN(delta);
 482         delta    = (delta ^ sign) - sign;
 483         code     = MAKE_CODE(delta);
 484         bits    += estimate_vlc(ff_prores_dc_codebook[codebook], code);
 485         codebook = (code + (code & 1)) >> 1;
 486         codebook = FFMIN(codebook, 3);
 487         sign     = new_sign;
 488         prev_dc  = dc;
 489     }
 490
 491     return bits;
 492 }
 493
 494 static int estimate_acs(int *error, DCTELEM *blocks, int blocks_per_slice,
 495                         int plane_size_factor,
 496                         const uint8_t *scan, const int16_t *qmat)
 497 {
 498     int idx, i;
 499     int run, level, run_cb, lev_cb;
 500     int max_coeffs, abs_level;
 501     int bits = 0;
 502
 503     max_coeffs = blocks_per_slice << 6;
 504     run_cb     = ff_prores_run_to_cb_index[4];
 505     lev_cb     = ff_prores_lev_to_cb_index[2];
 506     run        = 0;
 507
 508     for (i = 1; i < 64; i++) {
 509         for (idx = scan[i]; idx < max_coeffs; idx += 64) {
 510             level   = blocks[idx] / qmat[scan[i]];
 511             *error += FFABS(blocks[idx]) % qmat[scan[i]];
 512             if (level) {
 513                 abs_level = FFABS(level);
 514                 bits += estimate_vlc(ff_prores_ac_codebook[run_cb], run);
 515                 bits += estimate_vlc(ff_prores_ac_codebook[lev_cb],
 516                                      abs_level - 1) + 1;
 517
 518                 run_cb = ff_prores_run_to_cb_index[FFMIN(run, 15)];
 519                 lev_cb = ff_prores_lev_to_cb_index[FFMIN(abs_level, 9)];
 520                 run    = 0;
 521             } else {
 522                 run++;
 523             }
 524         }
 525     }
 526
 527     return bits;
 528 }
 529
 530 static int estimate_slice_plane(ProresContext *ctx, int *error, int plane,
 531                                 const uint16_t *src, int linesize,
 532                                 int mbs_per_slice,
 533                                 int blocks_per_mb, int plane_size_factor,
 534                                 const int16_t *qmat)
 535 {
 536     int blocks_per_slice;
 537     int bits;
 538
 539     blocks_per_slice = mbs_per_slice * blocks_per_mb;
 540
 541     bits  = estimate_dcs(error, ctx->blocks[plane], blocks_per_slice, qmat[0]);
 542     bits += estimate_acs(error, ctx->blocks[plane], blocks_per_slice,
 543                          plane_size_factor, ctx->scantable.permutated, qmat);
 544
 545     return FFALIGN(bits, 8);
 546 }
 547
 548 static int find_slice_quant(AVCodecContext *avctx, const AVFrame *pic,
 549                             int trellis_node, int x, int y, int mbs_per_slice)
 550 {
 551     ProresContext *ctx = avctx->priv_data;
 552     int i, q, pq, xp, yp;
 553     const uint16_t *src;
 554     int slice_width_factor = av_log2(mbs_per_slice);
 555     int num_cblocks[MAX_PLANES], pwidth;
 556     int plane_factor[MAX_PLANES], is_chroma[MAX_PLANES];
 557     const int min_quant = ctx->profile_info->min_quant;
 558     const int max_quant = ctx->profile_info->max_quant;
 559     int error, bits, bits_limit;
 560     int mbs, prev, cur, new_score;
 561     int slice_bits[TRELLIS_WIDTH], slice_score[TRELLIS_WIDTH];
 562     int overquant;
 563     uint16_t *qmat;
 564
 565     mbs = x + mbs_per_slice;
 566
 567     for (i = 0; i < ctx->num_planes; i++) {
 568         is_chroma[i]    = (i == 1 || i == 2);
 569         plane_factor[i] = slice_width_factor + 2;
 570         if (is_chroma[i])
 571             plane_factor[i] += ctx->chroma_factor - 3;
 572         if (!is_chroma[i] || ctx->chroma_factor == CFACTOR_Y444) {
 573             xp             = x << 4;
 574             yp             = y << 4;
 575             num_cblocks[i] = 4;
 576             pwidth         = avctx->width;
 577         } else {
 578             xp             = x << 3;
 579             yp             = y << 4;
 580             num_cblocks[i] = 2;
 581             pwidth         = avctx->width >> 1;
 582         }
 583         src = (const uint16_t*)(pic->data[i] + yp * pic->linesize[i]) + xp;
 584
 585         get_slice_data(ctx, src, pic->linesize[i], xp, yp,
 586                        pwidth, avctx->height, ctx->blocks[i],
 587                        mbs_per_slice, num_cblocks[i], is_chroma[i]);
 588     }
 589
 590     for (q = min_quant; q < max_quant + 2; q++) {
 591         ctx->nodes[trellis_node + q].prev_node = -1;
 592         ctx->nodes[trellis_node + q].quant     = q;
 593     }
 594
 595     // todo: maybe perform coarser quantising to fit into frame size when needed
 596     for (q = min_quant; q <= max_quant; q++) {
 597         bits  = 0;
 598         error = 0;
 599         for (i = 0; i < ctx->num_planes; i++) {
 600             bits += estimate_slice_plane(ctx, &error, i,
 601                                          src, pic->linesize[i],
 602                                          mbs_per_slice,
 603                                          num_cblocks[i], plane_factor[i],
 604                                          ctx->quants[q]);
 605         }
 606         if (bits > 65000 * 8) {
 607             error = SCORE_LIMIT;
 608             break;
 609         }
 610         slice_bits[q]  = bits;
 611         slice_score[q] = error;
 612     }
 613     if (slice_bits[max_quant] <= ctx->bits_per_mb * mbs_per_slice) {
 614         slice_bits[max_quant + 1]  = slice_bits[max_quant];
 615         slice_score[max_quant + 1] = slice_score[max_quant] + 1;
 616         overquant = max_quant;
 617     } else {
 618         for (q = max_quant + 1; q < 128; q++) {
 619             bits  = 0;
 620             error = 0;
 621             if (q < MAX_STORED_Q) {
 622                 qmat = ctx->quants[q];
 623             } else {
 624                 qmat = ctx->custom_q;
 625                 for (i = 0; i < 64; i++)
 626                     qmat[i] = ctx->quant_mat[i] * q;
 627             }
 628             for (i = 0; i < ctx->num_planes; i++) {
 629                 bits += estimate_slice_plane(ctx, &error, i,
 630                                              src, pic->linesize[i],
 631                                              mbs_per_slice,
 632                                              num_cblocks[i], plane_factor[i],
 633                                              qmat);
 634             }
 635             if (bits <= ctx->bits_per_mb * mbs_per_slice)
 636                 break;
 637         }
 638
 639         slice_bits[max_quant + 1]  = bits;
 640         slice_score[max_quant + 1] = error;
 641         overquant = q;
 642     }
 643     ctx->nodes[trellis_node + max_quant + 1].quant = overquant;
 644
 645     bits_limit = mbs * ctx->bits_per_mb;
 646     for (pq = min_quant; pq < max_quant + 2; pq++) {
 647         prev = trellis_node - TRELLIS_WIDTH + pq;
 648
 649         for (q = min_quant; q < max_quant + 2; q++) {
 650             cur = trellis_node + q;
 651
 652             bits  = ctx->nodes[prev].bits + slice_bits[q];
 653             error = slice_score[q];
 654             if (bits > bits_limit)
 655                 error = SCORE_LIMIT;
 656
 657             if (ctx->nodes[prev].score < SCORE_LIMIT && error < SCORE_LIMIT)
 658                 new_score = ctx->nodes[prev].score + error;
 659             else
 660                 new_score = SCORE_LIMIT;
 661             if (ctx->nodes[cur].prev_node == -1 ||
 662                 ctx->nodes[cur].score >= new_score) {
 663
 664                 ctx->nodes[cur].bits      = bits;
 665                 ctx->nodes[cur].score     = new_score;
 666                 ctx->nodes[cur].prev_node = prev;
 667             }
 668         }
 669     }
 670
 671     error = ctx->nodes[trellis_node + min_quant].score;
 672     pq    = trellis_node + min_quant;
 673     for (q = min_quant + 1; q < max_quant + 2; q++) {
 674         if (ctx->nodes[trellis_node + q].score <= error) {
 675             error = ctx->nodes[trellis_node + q].score;
 676             pq    = trellis_node + q;
 677         }
 678     }
 679
 680     return pq;
 681 }
 682
 683 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 684                         const AVFrame *pic, int *got_packet)
 685 {
 686     ProresContext *ctx = avctx->priv_data;
 687     uint8_t *orig_buf, *buf, *slice_hdr, *slice_sizes, *tmp;
 688     uint8_t *picture_size_pos;
 689     PutBitContext pb;
 690     int x, y, i, mb, q = 0;
 691     int sizes[4] = { 0 };
 692     int slice_hdr_size = 2 + 2 * (ctx->num_planes - 1);
 693     int frame_size, picture_size, slice_size;
 694     int mbs_per_slice = ctx->mbs_per_slice;
 695     int pkt_size, ret;
 696
 697     *avctx->coded_frame           = *pic;
 698     avctx->coded_frame->pict_type = AV_PICTURE_TYPE_I;
 699     avctx->coded_frame->key_frame = 1;
 700
 701     pkt_size = ctx->frame_size + FF_MIN_BUFFER_SIZE;
 702
 703     if ((ret = ff_alloc_packet(pkt, pkt_size)) < 0) {
 704         av_log(avctx, AV_LOG_ERROR, "Error getting output packet.\n");
 705         return ret;
 706     }
 707
 708     orig_buf = pkt->data;
 709
 710     // frame atom
 711     orig_buf += 4;                              // frame size
 712     bytestream_put_be32  (&orig_buf, FRAME_ID); // frame container ID
 713     buf = orig_buf;
 714
 715     // frame header
 716     tmp = buf;
 717     buf += 2;                                   // frame header size will be stored here
 718     bytestream_put_be16  (&buf, 0);             // version 1
 719     bytestream_put_buffer(&buf, ctx->vendor, 4);
 720     bytestream_put_be16  (&buf, avctx->width);
 721     bytestream_put_be16  (&buf, avctx->height);
 722     bytestream_put_byte  (&buf, ctx->chroma_factor << 6); // frame flags
 723     bytestream_put_byte  (&buf, 0);             // reserved
 724     bytestream_put_byte  (&buf, avctx->color_primaries);
 725     bytestream_put_byte  (&buf, avctx->color_trc);
 726     bytestream_put_byte  (&buf, avctx->colorspace);
 727     bytestream_put_byte  (&buf, 0x40);          // source format and alpha information
 728     bytestream_put_byte  (&buf, 0);             // reserved
 729     if (ctx->quant_sel != QUANT_MAT_DEFAULT) {
 730         bytestream_put_byte  (&buf, 0x03);      // matrix flags - both matrices are present
 731         // luma quantisation matrix
 732         for (i = 0; i < 64; i++)
 733             bytestream_put_byte(&buf, ctx->quant_mat[i]);
 734         // chroma quantisation matrix
 735         for (i = 0; i < 64; i++)
 736             bytestream_put_byte(&buf, ctx->quant_mat[i]);
 737     } else {
 738         bytestream_put_byte  (&buf, 0x00);      // matrix flags - default matrices are used
 739     }
 740     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
 741
 742     // picture header
 743     picture_size_pos = buf + 1;
 744     bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
 745     buf += 4;                                   // picture data size will be stored here
 746     bytestream_put_be16  (&buf, ctx->num_slices); // total number of slices
 747     bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
 748
 749     // seek table - will be filled during slice encoding
 750     slice_sizes = buf;
 751     buf += ctx->num_slices * 2;
 752
 753     // slices
 754     for (y = 0; y < ctx->mb_height; y++) {
 755         mbs_per_slice = ctx->mbs_per_slice;
 756         if (!ctx->force_quant) {
 757             for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
 758                 while (ctx->mb_width - x < mbs_per_slice)
 759                     mbs_per_slice >>= 1;
 760                 q = find_slice_quant(avctx, pic, (mb + 1) * TRELLIS_WIDTH, x, y,
 761                                      mbs_per_slice);
 762             }
 763
 764             for (x = ctx->slices_width - 1; x >= 0; x--) {
 765                 ctx->slice_q[x] = ctx->nodes[q].quant;
 766                 q = ctx->nodes[q].prev_node;
 767             }
 768         }
 769
 770         mbs_per_slice = ctx->mbs_per_slice;
 771         for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
 772             q = ctx->force_quant ? ctx->force_quant : ctx->slice_q[mb];
 773
 774             while (ctx->mb_width - x < mbs_per_slice)
 775                 mbs_per_slice >>= 1;
 776
 777             bytestream_put_byte(&buf, slice_hdr_size << 3);
 778             slice_hdr = buf;
 779             buf += slice_hdr_size - 1;
 780             init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
 781             encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
 782
 783             bytestream_put_byte(&slice_hdr, q);
 784             slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
 785             for (i = 0; i < ctx->num_planes - 1; i++) {
 786                 bytestream_put_be16(&slice_hdr, sizes[i]);
 787                 slice_size += sizes[i];
 788             }
 789             bytestream_put_be16(&slice_sizes, slice_size);
 790             buf += slice_size - slice_hdr_size;
 791         }
 792     }
 793
 794     orig_buf -= 8;
 795     frame_size = buf - orig_buf;
 796     picture_size = buf - picture_size_pos - 6;
 797     bytestream_put_be32(&orig_buf, frame_size);
 798     bytestream_put_be32(&picture_size_pos, picture_size);
 799
 800     pkt->size   = frame_size;
 801     pkt->flags |= AV_PKT_FLAG_KEY;
 802     *got_packet = 1;
 803
 804     return 0;
 805 }
 806
 807 static av_cold int encode_close(AVCodecContext *avctx)
 808 {
 809     ProresContext *ctx = avctx->priv_data;
 810
 811     if (avctx->coded_frame->data[0])
 812         avctx->release_buffer(avctx, avctx->coded_frame);
 813
 814     av_freep(&avctx->coded_frame);
 815
 816     av_freep(&ctx->nodes);
 817     av_freep(&ctx->slice_q);
 818
 819     return 0;
 820 }
 821
 822 static av_cold int encode_init(AVCodecContext *avctx)
 823 {
 824     ProresContext *ctx = avctx->priv_data;
 825     int mps;
 826     int i, j;
 827     int min_quant, max_quant;
 828
 829     avctx->bits_per_raw_sample = 10;
 830     avctx->coded_frame = avcodec_alloc_frame();
 831     if (!avctx->coded_frame)
 832         return AVERROR(ENOMEM);
 833
 834     ff_proresdsp_init(&ctx->dsp, avctx);
 835     ff_init_scantable(ctx->dsp.dct_permutation, &ctx->scantable,
 836                       ff_prores_progressive_scan);
 837
 838     mps = ctx->mbs_per_slice;
 839     if (mps & (mps - 1)) {
 840         av_log(avctx, AV_LOG_ERROR,
 841                "there should be an integer power of two MBs per slice\n");
 842         return AVERROR(EINVAL);
 843     }
 844
 845     ctx->chroma_factor = avctx->pix_fmt == PIX_FMT_YUV422P10
 846                          ? CFACTOR_Y422
 847                          : CFACTOR_Y444;
 848     ctx->profile_info  = prores_profile_info + ctx->profile;
 849     ctx->num_planes    = 3;
 850
 851     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
 852     ctx->mb_height     = FFALIGN(avctx->height, 16) >> 4;
 853     ctx->slices_width  = ctx->mb_width / mps;
 854     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
 855     ctx->num_slices    = ctx->mb_height * ctx->slices_width;
 856
 857     if (ctx->quant_sel == -1)
 858         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
 859     else
 860         ctx->quant_mat = prores_quant_matrices[ctx->quant_sel];
 861
 862     if (strlen(ctx->vendor) != 4) {
 863         av_log(avctx, AV_LOG_ERROR, "vendor ID should be 4 bytes\n");
 864         return AVERROR_INVALIDDATA;
 865     }
 866
 867     ctx->force_quant = avctx->global_quality / FF_QP2LAMBDA;
 868     if (!ctx->force_quant) {
 869         if (!ctx->bits_per_mb) {
 870             for (i = 0; i < NUM_MB_LIMITS - 1; i++)
 871                 if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height)
 872                     break;
 873             ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
 874         } else if (ctx->bits_per_mb < 128) {
 875             av_log(avctx, AV_LOG_ERROR, "too few bits per MB, please set at least 128\n");
 876             return AVERROR_INVALIDDATA;
 877         }
 878
 879         min_quant = ctx->profile_info->min_quant;
 880         max_quant = ctx->profile_info->max_quant;
 881         for (i = min_quant; i < MAX_STORED_Q; i++) {
 882             for (j = 0; j < 64; j++)
 883                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
 884         }
 885
 886         ctx->nodes = av_malloc((ctx->slices_width + 1) * TRELLIS_WIDTH
 887                                * sizeof(*ctx->nodes));
 888         if (!ctx->nodes) {
 889             encode_close(avctx);
 890             return AVERROR(ENOMEM);
 891         }
 892         for (i = min_quant; i < max_quant + 2; i++) {
 893             ctx->nodes[i].prev_node = -1;
 894             ctx->nodes[i].bits      = 0;
 895             ctx->nodes[i].score     = 0;
 896         }
 897
 898         ctx->slice_q = av_malloc(ctx->slices_width * sizeof(*ctx->slice_q));
 899         if (!ctx->slice_q) {
 900             encode_close(avctx);
 901             return AVERROR(ENOMEM);
 902         }
 903     } else {
 904         int ls = 0;
 905
 906         if (ctx->force_quant > 64) {
 907             av_log(avctx, AV_LOG_ERROR, "too large quantiser, maximum is 64\n");
 908             return AVERROR_INVALIDDATA;
 909         }
 910
 911         for (j = 0; j < 64; j++) {
 912             ctx->quants[0][j] = ctx->quant_mat[j] * ctx->force_quant;
 913             ls += av_log2((1 << 11)  / ctx->quants[0][j]) * 2 + 1;
 914         }
 915
 916         ctx->bits_per_mb = ls * 8;
 917         if (ctx->chroma_factor == CFACTOR_Y444)
 918             ctx->bits_per_mb += ls * 4;
 919         if (ctx->num_planes == 4)
 920             ctx->bits_per_mb += ls * 4;
 921     }
 922
 923     ctx->frame_size = ctx->num_slices * (2 + 2 * ctx->num_planes
 924                                          + (2 * mps * ctx->bits_per_mb) / 8)
 925                       + 200;
 926
 927     avctx->codec_tag   = ctx->profile_info->tag;
 928
 929     av_log(avctx, AV_LOG_DEBUG, "profile %d, %d slices, %d bits per MB\n",
 930            ctx->profile, ctx->num_slices, ctx->bits_per_mb);
 931     av_log(avctx, AV_LOG_DEBUG, "estimated frame size %d\n",
 932            ctx->frame_size);
 933
 934     return 0;
 935 }
 936
 937 #define OFFSET(x) offsetof(ProresContext, x)
 938 #define VE     AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
 939
 940 static const AVOption options[] = {
 941     { "mbs_per_slice", "macroblocks per slice", OFFSET(mbs_per_slice),
 942         AV_OPT_TYPE_INT, { 8 }, 1, MAX_MBS_PER_SLICE, VE },
 943     { "profile",       NULL, OFFSET(profile), AV_OPT_TYPE_INT,
 944         { PRORES_PROFILE_STANDARD },
 945         PRORES_PROFILE_PROXY, PRORES_PROFILE_HQ, VE, "profile" },
 946     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_PROXY },
 947         0, 0, VE, "profile" },
 948     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_LT },
 949         0, 0, VE, "profile" },
 950     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_STANDARD },
 951         0, 0, VE, "profile" },
 952     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_HQ },
 953         0, 0, VE, "profile" },
 954     { "vendor", "vendor ID", OFFSET(vendor),
 955         AV_OPT_TYPE_STRING, { .str = "Lavc" }, CHAR_MIN, CHAR_MAX, VE },
 956     { "bits_per_mb", "desired bits per macroblock", OFFSET(bits_per_mb),
 957         AV_OPT_TYPE_INT, { 0 }, 0, 8192, VE },
 958     { "quant_mat", "quantiser matrix", OFFSET(quant_sel), AV_OPT_TYPE_INT,
 959         { -1 }, -1, QUANT_MAT_DEFAULT, VE, "quant_mat" },
 960     { "auto",          NULL, 0, AV_OPT_TYPE_CONST, { -1 },
 961         0, 0, VE, "quant_mat" },
 962     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_PROXY },
 963         0, 0, VE, "quant_mat" },
 964     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_LT },
 965         0, 0, VE, "quant_mat" },
 966     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_STANDARD },
 967         0, 0, VE, "quant_mat" },
 968     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_HQ },
 969         0, 0, VE, "quant_mat" },
 970     { "default",       NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_DEFAULT },
 971         0, 0, VE, "quant_mat" },
 972     { NULL }
 973 };
 974
 975 static const AVClass proresenc_class = {
 976     .class_name = "ProRes encoder",
 977     .item_name  = av_default_item_name,
 978     .option     = options,
 979     .version    = LIBAVUTIL_VERSION_INT,
 980 };
 981
 982 AVCodec ff_prores_kostya_encoder = {
 983     .name           = "prores_kostya",
 984     .type           = AVMEDIA_TYPE_VIDEO,
 985     .id             = CODEC_ID_PRORES,
 986     .priv_data_size = sizeof(ProresContext),
 987     .init           = encode_init,
 988     .close          = encode_close,
 989     .encode2        = encode_frame,
 990     .long_name      = NULL_IF_CONFIG_SMALL("Apple ProRes (iCodec Pro)"),
 991     .pix_fmts       = (const enum PixelFormat[]) {
 992                           PIX_FMT_YUV422P10, PIX_FMT_YUV444P10, PIX_FMT_NONE
 993                       },
 994     .priv_class     = &proresenc_class,
 995 };