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