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_packet2(avctx, pkt, pkt_size)) < 0)
 704         return ret;
 705
 706     orig_buf = pkt->data;
 707
 708     // frame atom
 709     orig_buf += 4;                              // frame size
 710     bytestream_put_be32  (&orig_buf, FRAME_ID); // frame container ID
 711     buf = orig_buf;
 712
 713     // frame header
 714     tmp = buf;
 715     buf += 2;                                   // frame header size will be stored here
 716     bytestream_put_be16  (&buf, 0);             // version 1
 717     bytestream_put_buffer(&buf, ctx->vendor, 4);
 718     bytestream_put_be16  (&buf, avctx->width);
 719     bytestream_put_be16  (&buf, avctx->height);
 720     bytestream_put_byte  (&buf, ctx->chroma_factor << 6); // frame flags
 721     bytestream_put_byte  (&buf, 0);             // reserved
 722     bytestream_put_byte  (&buf, avctx->color_primaries);
 723     bytestream_put_byte  (&buf, avctx->color_trc);
 724     bytestream_put_byte  (&buf, avctx->colorspace);
 725     bytestream_put_byte  (&buf, 0x40);          // source format and alpha information
 726     bytestream_put_byte  (&buf, 0);             // reserved
 727     if (ctx->quant_sel != QUANT_MAT_DEFAULT) {
 728         bytestream_put_byte  (&buf, 0x03);      // matrix flags - both matrices are present
 729         // luma quantisation matrix
 730         for (i = 0; i < 64; i++)
 731             bytestream_put_byte(&buf, ctx->quant_mat[i]);
 732         // chroma quantisation matrix
 733         for (i = 0; i < 64; i++)
 734             bytestream_put_byte(&buf, ctx->quant_mat[i]);
 735     } else {
 736         bytestream_put_byte  (&buf, 0x00);      // matrix flags - default matrices are used
 737     }
 738     bytestream_put_be16  (&tmp, buf - orig_buf); // write back frame header size
 739
 740     // picture header
 741     picture_size_pos = buf + 1;
 742     bytestream_put_byte  (&buf, 0x40);          // picture header size (in bits)
 743     buf += 4;                                   // picture data size will be stored here
 744     bytestream_put_be16  (&buf, ctx->num_slices); // total number of slices
 745     bytestream_put_byte  (&buf, av_log2(ctx->mbs_per_slice) << 4); // slice width and height in MBs
 746
 747     // seek table - will be filled during slice encoding
 748     slice_sizes = buf;
 749     buf += ctx->num_slices * 2;
 750
 751     // slices
 752     for (y = 0; y < ctx->mb_height; y++) {
 753         mbs_per_slice = ctx->mbs_per_slice;
 754         if (!ctx->force_quant) {
 755             for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
 756                 while (ctx->mb_width - x < mbs_per_slice)
 757                     mbs_per_slice >>= 1;
 758                 q = find_slice_quant(avctx, pic, (mb + 1) * TRELLIS_WIDTH, x, y,
 759                                      mbs_per_slice);
 760             }
 761
 762             for (x = ctx->slices_width - 1; x >= 0; x--) {
 763                 ctx->slice_q[x] = ctx->nodes[q].quant;
 764                 q = ctx->nodes[q].prev_node;
 765             }
 766         }
 767
 768         mbs_per_slice = ctx->mbs_per_slice;
 769         for (x = mb = 0; x < ctx->mb_width; x += mbs_per_slice, mb++) {
 770             q = ctx->force_quant ? ctx->force_quant : ctx->slice_q[mb];
 771
 772             while (ctx->mb_width - x < mbs_per_slice)
 773                 mbs_per_slice >>= 1;
 774
 775             bytestream_put_byte(&buf, slice_hdr_size << 3);
 776             slice_hdr = buf;
 777             buf += slice_hdr_size - 1;
 778             init_put_bits(&pb, buf, (pkt_size - (buf - orig_buf)) * 8);
 779             encode_slice(avctx, pic, &pb, sizes, x, y, q, mbs_per_slice);
 780
 781             bytestream_put_byte(&slice_hdr, q);
 782             slice_size = slice_hdr_size + sizes[ctx->num_planes - 1];
 783             for (i = 0; i < ctx->num_planes - 1; i++) {
 784                 bytestream_put_be16(&slice_hdr, sizes[i]);
 785                 slice_size += sizes[i];
 786             }
 787             bytestream_put_be16(&slice_sizes, slice_size);
 788             buf += slice_size - slice_hdr_size;
 789         }
 790     }
 791
 792     orig_buf -= 8;
 793     frame_size = buf - orig_buf;
 794     picture_size = buf - picture_size_pos - 6;
 795     bytestream_put_be32(&orig_buf, frame_size);
 796     bytestream_put_be32(&picture_size_pos, picture_size);
 797
 798     pkt->size   = frame_size;
 799     pkt->flags |= AV_PKT_FLAG_KEY;
 800     *got_packet = 1;
 801
 802     return 0;
 803 }
 804
 805 static av_cold int encode_close(AVCodecContext *avctx)
 806 {
 807     ProresContext *ctx = avctx->priv_data;
 808
 809     if (avctx->coded_frame->data[0])
 810         avctx->release_buffer(avctx, avctx->coded_frame);
 811
 812     av_freep(&avctx->coded_frame);
 813
 814     av_freep(&ctx->nodes);
 815     av_freep(&ctx->slice_q);
 816
 817     return 0;
 818 }
 819
 820 static av_cold int encode_init(AVCodecContext *avctx)
 821 {
 822     ProresContext *ctx = avctx->priv_data;
 823     int mps;
 824     int i, j;
 825     int min_quant, max_quant;
 826
 827     avctx->bits_per_raw_sample = 10;
 828     avctx->coded_frame = avcodec_alloc_frame();
 829     if (!avctx->coded_frame)
 830         return AVERROR(ENOMEM);
 831
 832     ff_proresdsp_init(&ctx->dsp, avctx);
 833     ff_init_scantable(ctx->dsp.dct_permutation, &ctx->scantable,
 834                       ff_prores_progressive_scan);
 835
 836     mps = ctx->mbs_per_slice;
 837     if (mps & (mps - 1)) {
 838         av_log(avctx, AV_LOG_ERROR,
 839                "there should be an integer power of two MBs per slice\n");
 840         return AVERROR(EINVAL);
 841     }
 842
 843     ctx->chroma_factor = avctx->pix_fmt == PIX_FMT_YUV422P10
 844                          ? CFACTOR_Y422
 845                          : CFACTOR_Y444;
 846     ctx->profile_info  = prores_profile_info + ctx->profile;
 847     ctx->num_planes    = 3;
 848
 849     ctx->mb_width      = FFALIGN(avctx->width,  16) >> 4;
 850     ctx->mb_height     = FFALIGN(avctx->height, 16) >> 4;
 851     ctx->slices_width  = ctx->mb_width / mps;
 852     ctx->slices_width += av_popcount(ctx->mb_width - ctx->slices_width * mps);
 853     ctx->num_slices    = ctx->mb_height * ctx->slices_width;
 854
 855     if (ctx->quant_sel == -1)
 856         ctx->quant_mat = prores_quant_matrices[ctx->profile_info->quant];
 857     else
 858         ctx->quant_mat = prores_quant_matrices[ctx->quant_sel];
 859
 860     if (strlen(ctx->vendor) != 4) {
 861         av_log(avctx, AV_LOG_ERROR, "vendor ID should be 4 bytes\n");
 862         return AVERROR_INVALIDDATA;
 863     }
 864
 865     ctx->force_quant = avctx->global_quality / FF_QP2LAMBDA;
 866     if (!ctx->force_quant) {
 867         if (!ctx->bits_per_mb) {
 868             for (i = 0; i < NUM_MB_LIMITS - 1; i++)
 869                 if (prores_mb_limits[i] >= ctx->mb_width * ctx->mb_height)
 870                     break;
 871             ctx->bits_per_mb   = ctx->profile_info->br_tab[i];
 872         } else if (ctx->bits_per_mb < 128) {
 873             av_log(avctx, AV_LOG_ERROR, "too few bits per MB, please set at least 128\n");
 874             return AVERROR_INVALIDDATA;
 875         }
 876
 877         min_quant = ctx->profile_info->min_quant;
 878         max_quant = ctx->profile_info->max_quant;
 879         for (i = min_quant; i < MAX_STORED_Q; i++) {
 880             for (j = 0; j < 64; j++)
 881                 ctx->quants[i][j] = ctx->quant_mat[j] * i;
 882         }
 883
 884         ctx->nodes = av_malloc((ctx->slices_width + 1) * TRELLIS_WIDTH
 885                                * sizeof(*ctx->nodes));
 886         if (!ctx->nodes) {
 887             encode_close(avctx);
 888             return AVERROR(ENOMEM);
 889         }
 890         for (i = min_quant; i < max_quant + 2; i++) {
 891             ctx->nodes[i].prev_node = -1;
 892             ctx->nodes[i].bits      = 0;
 893             ctx->nodes[i].score     = 0;
 894         }
 895
 896         ctx->slice_q = av_malloc(ctx->slices_width * sizeof(*ctx->slice_q));
 897         if (!ctx->slice_q) {
 898             encode_close(avctx);
 899             return AVERROR(ENOMEM);
 900         }
 901     } else {
 902         int ls = 0;
 903
 904         if (ctx->force_quant > 64) {
 905             av_log(avctx, AV_LOG_ERROR, "too large quantiser, maximum is 64\n");
 906             return AVERROR_INVALIDDATA;
 907         }
 908
 909         for (j = 0; j < 64; j++) {
 910             ctx->quants[0][j] = ctx->quant_mat[j] * ctx->force_quant;
 911             ls += av_log2((1 << 11)  / ctx->quants[0][j]) * 2 + 1;
 912         }
 913
 914         ctx->bits_per_mb = ls * 8;
 915         if (ctx->chroma_factor == CFACTOR_Y444)
 916             ctx->bits_per_mb += ls * 4;
 917         if (ctx->num_planes == 4)
 918             ctx->bits_per_mb += ls * 4;
 919     }
 920
 921     ctx->frame_size = ctx->num_slices * (2 + 2 * ctx->num_planes
 922                                          + (2 * mps * ctx->bits_per_mb) / 8)
 923                       + 200;
 924
 925     avctx->codec_tag   = ctx->profile_info->tag;
 926
 927     av_log(avctx, AV_LOG_DEBUG, "profile %d, %d slices, %d bits per MB\n",
 928            ctx->profile, ctx->num_slices, ctx->bits_per_mb);
 929     av_log(avctx, AV_LOG_DEBUG, "estimated frame size %d\n",
 930            ctx->frame_size);
 931
 932     return 0;
 933 }
 934
 935 #define OFFSET(x) offsetof(ProresContext, x)
 936 #define VE     AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
 937
 938 static const AVOption options[] = {
 939     { "mbs_per_slice", "macroblocks per slice", OFFSET(mbs_per_slice),
 940         AV_OPT_TYPE_INT, { 8 }, 1, MAX_MBS_PER_SLICE, VE },
 941     { "profile",       NULL, OFFSET(profile), AV_OPT_TYPE_INT,
 942         { PRORES_PROFILE_STANDARD },
 943         PRORES_PROFILE_PROXY, PRORES_PROFILE_HQ, VE, "profile" },
 944     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_PROXY },
 945         0, 0, VE, "profile" },
 946     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_LT },
 947         0, 0, VE, "profile" },
 948     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_STANDARD },
 949         0, 0, VE, "profile" },
 950     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { PRORES_PROFILE_HQ },
 951         0, 0, VE, "profile" },
 952     { "vendor", "vendor ID", OFFSET(vendor),
 953         AV_OPT_TYPE_STRING, { .str = "Lavc" }, CHAR_MIN, CHAR_MAX, VE },
 954     { "bits_per_mb", "desired bits per macroblock", OFFSET(bits_per_mb),
 955         AV_OPT_TYPE_INT, { 0 }, 0, 8192, VE },
 956     { "quant_mat", "quantiser matrix", OFFSET(quant_sel), AV_OPT_TYPE_INT,
 957         { -1 }, -1, QUANT_MAT_DEFAULT, VE, "quant_mat" },
 958     { "auto",          NULL, 0, AV_OPT_TYPE_CONST, { -1 },
 959         0, 0, VE, "quant_mat" },
 960     { "proxy",         NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_PROXY },
 961         0, 0, VE, "quant_mat" },
 962     { "lt",            NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_LT },
 963         0, 0, VE, "quant_mat" },
 964     { "standard",      NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_STANDARD },
 965         0, 0, VE, "quant_mat" },
 966     { "hq",            NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_HQ },
 967         0, 0, VE, "quant_mat" },
 968     { "default",       NULL, 0, AV_OPT_TYPE_CONST, { QUANT_MAT_DEFAULT },
 969         0, 0, VE, "quant_mat" },
 970     { NULL }
 971 };
 972
 973 static const AVClass proresenc_class = {
 974     .class_name = "ProRes encoder",
 975     .item_name  = av_default_item_name,
 976     .option     = options,
 977     .version    = LIBAVUTIL_VERSION_INT,
 978 };
 979
 980 AVCodec ff_prores_kostya_encoder = {
 981     .name           = "prores_kostya",
 982     .type           = AVMEDIA_TYPE_VIDEO,
 983     .id             = CODEC_ID_PRORES,
 984     .priv_data_size = sizeof(ProresContext),
 985     .init           = encode_init,
 986     .close          = encode_close,
 987     .encode2        = encode_frame,
 988     .long_name      = NULL_IF_CONFIG_SMALL("Apple ProRes (iCodec Pro)"),
 989     .pix_fmts       = (const enum PixelFormat[]) {
 990                           PIX_FMT_YUV422P10, PIX_FMT_YUV444P10, PIX_FMT_NONE
 991                       },
 992     .priv_class     = &proresenc_class,
 993 };