pictor: Correctly check frame dimensions
[ffmpeg.git] / libavcodec / ratecontrol.c
1 /*
2  * Rate control for video encoders
3  *
4  * Copyright (c) 2002-2004 Michael Niedermayer <michaelni@gmx.at>
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 /**
24  * @file
25  * Rate control for video encoders.
26  */
27
28 #include "libavutil/attributes.h"
29 #include "libavutil/internal.h"
30
31 #include "avcodec.h"
32 #include "internal.h"
33 #include "ratecontrol.h"
34 #include "mpegutils.h"
35 #include "mpegvideo.h"
36 #include "libavutil/eval.h"
37
38 #undef NDEBUG // Always check asserts, the speed effect is far too small to disable them.
39 #include <assert.h>
40
41 #ifndef M_E
42 #define M_E 2.718281828
43 #endif
44
45 static inline double qp2bits(RateControlEntry *rce, double qp)
46 {
47     if (qp <= 0.0) {
48         av_log(NULL, AV_LOG_ERROR, "qp<=0.0\n");
49     }
50     return rce->qscale * (double)(rce->i_tex_bits + rce->p_tex_bits + 1) / qp;
51 }
52
53 static inline double bits2qp(RateControlEntry *rce, double bits)
54 {
55     if (bits < 0.9) {
56         av_log(NULL, AV_LOG_ERROR, "bits<0.9\n");
57     }
58     return rce->qscale * (double)(rce->i_tex_bits + rce->p_tex_bits + 1) / bits;
59 }
60
61 static double get_diff_limited_q(MpegEncContext *s, RateControlEntry *rce, double q)
62 {
63     RateControlContext *rcc   = &s->rc_context;
64     AVCodecContext *a         = s->avctx;
65     const int pict_type       = rce->new_pict_type;
66     const double last_p_q     = rcc->last_qscale_for[AV_PICTURE_TYPE_P];
67     const double last_non_b_q = rcc->last_qscale_for[rcc->last_non_b_pict_type];
68
69     if (pict_type == AV_PICTURE_TYPE_I &&
70         (a->i_quant_factor > 0.0 || rcc->last_non_b_pict_type == AV_PICTURE_TYPE_P))
71         q = last_p_q * FFABS(a->i_quant_factor) + a->i_quant_offset;
72     else if (pict_type == AV_PICTURE_TYPE_B &&
73              a->b_quant_factor > 0.0)
74         q = last_non_b_q * a->b_quant_factor + a->b_quant_offset;
75     if (q < 1)
76         q = 1;
77
78     /* last qscale / qdiff stuff */
79     if (rcc->last_non_b_pict_type == pict_type || pict_type != AV_PICTURE_TYPE_I) {
80         double last_q     = rcc->last_qscale_for[pict_type];
81         const int maxdiff = FF_QP2LAMBDA * a->max_qdiff;
82
83         if (q > last_q + maxdiff)
84             q = last_q + maxdiff;
85         else if (q < last_q - maxdiff)
86             q = last_q - maxdiff;
87     }
88
89     rcc->last_qscale_for[pict_type] = q; // Note we cannot do that after blurring
90
91     if (pict_type != AV_PICTURE_TYPE_B)
92         rcc->last_non_b_pict_type = pict_type;
93
94     return q;
95 }
96
97 /**
98  * Get the qmin & qmax for pict_type.
99  */
100 static void get_qminmax(int *qmin_ret, int *qmax_ret, MpegEncContext *s, int pict_type)
101 {
102     int qmin = s->lmin;
103     int qmax = s->lmax;
104
105     assert(qmin <= qmax);
106
107     switch (pict_type) {
108     case AV_PICTURE_TYPE_B:
109         qmin = (int)(qmin * FFABS(s->avctx->b_quant_factor) + s->avctx->b_quant_offset + 0.5);
110         qmax = (int)(qmax * FFABS(s->avctx->b_quant_factor) + s->avctx->b_quant_offset + 0.5);
111         break;
112     case AV_PICTURE_TYPE_I:
113         qmin = (int)(qmin * FFABS(s->avctx->i_quant_factor) + s->avctx->i_quant_offset + 0.5);
114         qmax = (int)(qmax * FFABS(s->avctx->i_quant_factor) + s->avctx->i_quant_offset + 0.5);
115         break;
116     }
117
118     qmin = av_clip(qmin, 1, FF_LAMBDA_MAX);
119     qmax = av_clip(qmax, 1, FF_LAMBDA_MAX);
120
121     if (qmax < qmin)
122         qmax = qmin;
123
124     *qmin_ret = qmin;
125     *qmax_ret = qmax;
126 }
127
128 static double modify_qscale(MpegEncContext *s, RateControlEntry *rce,
129                             double q, int frame_num)
130 {
131     RateControlContext *rcc  = &s->rc_context;
132     const double buffer_size = s->avctx->rc_buffer_size;
133     const double fps         = 1 / av_q2d(s->avctx->time_base);
134     const double min_rate    = s->avctx->rc_min_rate / fps;
135     const double max_rate    = s->avctx->rc_max_rate / fps;
136     const int pict_type      = rce->new_pict_type;
137     int qmin, qmax;
138
139     get_qminmax(&qmin, &qmax, s, pict_type);
140
141     /* modulation */
142     if (s->rc_qmod_freq &&
143         frame_num % s->rc_qmod_freq == 0 &&
144         pict_type == AV_PICTURE_TYPE_P)
145         q *= s->rc_qmod_amp;
146
147     /* buffer overflow/underflow protection */
148     if (buffer_size) {
149         double expected_size = rcc->buffer_index;
150         double q_limit;
151
152         if (min_rate) {
153             double d = 2 * (buffer_size - expected_size) / buffer_size;
154             if (d > 1.0)
155                 d = 1.0;
156             else if (d < 0.0001)
157                 d = 0.0001;
158             q *= pow(d, 1.0 / s->rc_buffer_aggressivity);
159
160             q_limit = bits2qp(rce,
161                               FFMAX((min_rate - buffer_size + rcc->buffer_index) *
162                                     s->avctx->rc_min_vbv_overflow_use, 1));
163
164             if (q > q_limit) {
165                 if (s->avctx->debug & FF_DEBUG_RC)
166                     av_log(s->avctx, AV_LOG_DEBUG,
167                            "limiting QP %f -> %f\n", q, q_limit);
168                 q = q_limit;
169             }
170         }
171
172         if (max_rate) {
173             double d = 2 * expected_size / buffer_size;
174             if (d > 1.0)
175                 d = 1.0;
176             else if (d < 0.0001)
177                 d = 0.0001;
178             q /= pow(d, 1.0 / s->rc_buffer_aggressivity);
179
180             q_limit = bits2qp(rce,
181                               FFMAX(rcc->buffer_index *
182                                     s->avctx->rc_max_available_vbv_use,
183                                     1));
184             if (q < q_limit) {
185                 if (s->avctx->debug & FF_DEBUG_RC)
186                     av_log(s->avctx, AV_LOG_DEBUG,
187                            "limiting QP %f -> %f\n", q, q_limit);
188                 q = q_limit;
189             }
190         }
191     }
192     ff_dlog(s, "q:%f max:%f min:%f size:%f index:%f agr:%f\n",
193             q, max_rate, min_rate, buffer_size, rcc->buffer_index,
194             s->rc_buffer_aggressivity);
195     if (s->rc_qsquish == 0.0 || qmin == qmax) {
196         if (q < qmin)
197             q = qmin;
198         else if (q > qmax)
199             q = qmax;
200     } else {
201         double min2 = log(qmin);
202         double max2 = log(qmax);
203
204         q  = log(q);
205         q  = (q - min2) / (max2 - min2) - 0.5;
206         q *= -4.0;
207         q  = 1.0 / (1.0 + exp(q));
208         q  = q * (max2 - min2) + min2;
209
210         q = exp(q);
211     }
212
213     return q;
214 }
215
216 /**
217  * Modify the bitrate curve from pass1 for one frame.
218  */
219 static double get_qscale(MpegEncContext *s, RateControlEntry *rce,
220                          double rate_factor, int frame_num)
221 {
222     RateControlContext *rcc = &s->rc_context;
223     AVCodecContext *a       = s->avctx;
224     const int pict_type     = rce->new_pict_type;
225     const double mb_num     = s->mb_num;
226     double q, bits;
227     int i;
228
229     double const_values[] = {
230         M_PI,
231         M_E,
232         rce->i_tex_bits * rce->qscale,
233         rce->p_tex_bits * rce->qscale,
234         (rce->i_tex_bits + rce->p_tex_bits) * (double)rce->qscale,
235         rce->mv_bits / mb_num,
236         rce->pict_type == AV_PICTURE_TYPE_B ? (rce->f_code + rce->b_code) * 0.5 : rce->f_code,
237         rce->i_count / mb_num,
238         rce->mc_mb_var_sum / mb_num,
239         rce->mb_var_sum / mb_num,
240         rce->pict_type == AV_PICTURE_TYPE_I,
241         rce->pict_type == AV_PICTURE_TYPE_P,
242         rce->pict_type == AV_PICTURE_TYPE_B,
243         rcc->qscale_sum[pict_type] / (double)rcc->frame_count[pict_type],
244         a->qcompress,
245         rcc->i_cplx_sum[AV_PICTURE_TYPE_I] / (double)rcc->frame_count[AV_PICTURE_TYPE_I],
246         rcc->i_cplx_sum[AV_PICTURE_TYPE_P] / (double)rcc->frame_count[AV_PICTURE_TYPE_P],
247         rcc->p_cplx_sum[AV_PICTURE_TYPE_P] / (double)rcc->frame_count[AV_PICTURE_TYPE_P],
248         rcc->p_cplx_sum[AV_PICTURE_TYPE_B] / (double)rcc->frame_count[AV_PICTURE_TYPE_B],
249         (rcc->i_cplx_sum[pict_type] + rcc->p_cplx_sum[pict_type]) / (double)rcc->frame_count[pict_type],
250         0
251     };
252
253     bits = av_expr_eval(rcc->rc_eq_eval, const_values, rce);
254     if (isnan(bits)) {
255         av_log(s->avctx, AV_LOG_ERROR, "Error evaluating rc_eq \"%s\"\n", s->rc_eq);
256         return -1;
257     }
258
259     rcc->pass1_rc_eq_output_sum += bits;
260     bits *= rate_factor;
261     if (bits < 0.0)
262         bits = 0.0;
263     bits += 1.0; // avoid 1/0 issues
264
265     /* user override */
266     for (i = 0; i < s->avctx->rc_override_count; i++) {
267         RcOverride *rco = s->avctx->rc_override;
268         if (rco[i].start_frame > frame_num)
269             continue;
270         if (rco[i].end_frame < frame_num)
271             continue;
272
273         if (rco[i].qscale)
274             bits = qp2bits(rce, rco[i].qscale);  // FIXME move at end to really force it?
275         else
276             bits *= rco[i].quality_factor;
277     }
278
279     q = bits2qp(rce, bits);
280
281     /* I/B difference */
282     if (pict_type == AV_PICTURE_TYPE_I && s->avctx->i_quant_factor < 0.0)
283         q = -q * s->avctx->i_quant_factor + s->avctx->i_quant_offset;
284     else if (pict_type == AV_PICTURE_TYPE_B && s->avctx->b_quant_factor < 0.0)
285         q = -q * s->avctx->b_quant_factor + s->avctx->b_quant_offset;
286     if (q < 1)
287         q = 1;
288
289     return q;
290 }
291
292 static int init_pass2(MpegEncContext *s)
293 {
294     RateControlContext *rcc = &s->rc_context;
295     AVCodecContext *a       = s->avctx;
296     int i, toobig;
297     double fps             = 1 / av_q2d(s->avctx->time_base);
298     double complexity[5]   = { 0 }; // approximate bits at quant=1
299     uint64_t const_bits[5] = { 0 }; // quantizer independent bits
300     uint64_t all_const_bits;
301     uint64_t all_available_bits = (uint64_t)(s->bit_rate *
302                                              (double)rcc->num_entries / fps);
303     double rate_factor          = 0;
304     double step;
305     const int filter_size = (int)(a->qblur * 4) | 1;
306     double expected_bits;
307     double *qscale, *blurred_qscale, qscale_sum;
308
309     /* find complexity & const_bits & decide the pict_types */
310     for (i = 0; i < rcc->num_entries; i++) {
311         RateControlEntry *rce = &rcc->entry[i];
312
313         rce->new_pict_type                = rce->pict_type;
314         rcc->i_cplx_sum[rce->pict_type]  += rce->i_tex_bits * rce->qscale;
315         rcc->p_cplx_sum[rce->pict_type]  += rce->p_tex_bits * rce->qscale;
316         rcc->mv_bits_sum[rce->pict_type] += rce->mv_bits;
317         rcc->frame_count[rce->pict_type]++;
318
319         complexity[rce->new_pict_type] += (rce->i_tex_bits + rce->p_tex_bits) *
320                                           (double)rce->qscale;
321         const_bits[rce->new_pict_type] += rce->mv_bits + rce->misc_bits;
322     }
323
324     all_const_bits = const_bits[AV_PICTURE_TYPE_I] +
325                      const_bits[AV_PICTURE_TYPE_P] +
326                      const_bits[AV_PICTURE_TYPE_B];
327
328     if (all_available_bits < all_const_bits) {
329         av_log(s->avctx, AV_LOG_ERROR, "requested bitrate is too low\n");
330         return -1;
331     }
332
333     qscale         = av_malloc(sizeof(double) * rcc->num_entries);
334     blurred_qscale = av_malloc(sizeof(double) * rcc->num_entries);
335     if (!qscale || !blurred_qscale) {
336         av_free(qscale);
337         av_free(blurred_qscale);
338         return AVERROR(ENOMEM);
339     }
340     toobig = 0;
341
342     for (step = 256 * 256; step > 0.0000001; step *= 0.5) {
343         expected_bits = 0;
344         rate_factor  += step;
345
346         rcc->buffer_index = s->avctx->rc_buffer_size / 2;
347
348         /* find qscale */
349         for (i = 0; i < rcc->num_entries; i++) {
350             RateControlEntry *rce = &rcc->entry[i];
351
352             qscale[i] = get_qscale(s, &rcc->entry[i], rate_factor, i);
353             rcc->last_qscale_for[rce->pict_type] = qscale[i];
354         }
355         assert(filter_size % 2 == 1);
356
357         /* fixed I/B QP relative to P mode */
358         for (i = rcc->num_entries - 1; i >= 0; i--) {
359             RateControlEntry *rce = &rcc->entry[i];
360
361             qscale[i] = get_diff_limited_q(s, rce, qscale[i]);
362         }
363
364         /* smooth curve */
365         for (i = 0; i < rcc->num_entries; i++) {
366             RateControlEntry *rce = &rcc->entry[i];
367             const int pict_type   = rce->new_pict_type;
368             int j;
369             double q = 0.0, sum = 0.0;
370
371             for (j = 0; j < filter_size; j++) {
372                 int index    = i + j - filter_size / 2;
373                 double d     = index - i;
374                 double coeff = a->qblur == 0 ? 1.0 : exp(-d * d / (a->qblur * a->qblur));
375
376                 if (index < 0 || index >= rcc->num_entries)
377                     continue;
378                 if (pict_type != rcc->entry[index].new_pict_type)
379                     continue;
380                 q   += qscale[index] * coeff;
381                 sum += coeff;
382             }
383             blurred_qscale[i] = q / sum;
384         }
385
386         /* find expected bits */
387         for (i = 0; i < rcc->num_entries; i++) {
388             RateControlEntry *rce = &rcc->entry[i];
389             double bits;
390
391             rce->new_qscale = modify_qscale(s, rce, blurred_qscale[i], i);
392
393             bits  = qp2bits(rce, rce->new_qscale) + rce->mv_bits + rce->misc_bits;
394             bits += 8 * ff_vbv_update(s, bits);
395
396             rce->expected_bits = expected_bits;
397             expected_bits     += bits;
398         }
399
400         ff_dlog(s->avctx,
401                 "expected_bits: %f all_available_bits: %d rate_factor: %f\n",
402                 expected_bits, (int)all_available_bits, rate_factor);
403         if (expected_bits > all_available_bits) {
404             rate_factor -= step;
405             ++toobig;
406         }
407     }
408     av_free(qscale);
409     av_free(blurred_qscale);
410
411     /* check bitrate calculations and print info */
412     qscale_sum = 0.0;
413     for (i = 0; i < rcc->num_entries; i++) {
414         ff_dlog(s, "[lavc rc] entry[%d].new_qscale = %.3f  qp = %.3f\n",
415                 i,
416                 rcc->entry[i].new_qscale,
417                 rcc->entry[i].new_qscale / FF_QP2LAMBDA);
418         qscale_sum += av_clip(rcc->entry[i].new_qscale / FF_QP2LAMBDA,
419                               s->avctx->qmin, s->avctx->qmax);
420     }
421     assert(toobig <= 40);
422     av_log(s->avctx, AV_LOG_DEBUG,
423            "[lavc rc] requested bitrate: %d bps  expected bitrate: %d bps\n",
424            s->bit_rate,
425            (int)(expected_bits / ((double)all_available_bits / s->bit_rate)));
426     av_log(s->avctx, AV_LOG_DEBUG,
427            "[lavc rc] estimated target average qp: %.3f\n",
428            (float)qscale_sum / rcc->num_entries);
429     if (toobig == 0) {
430         av_log(s->avctx, AV_LOG_INFO,
431                "[lavc rc] Using all of requested bitrate is not "
432                "necessary for this video with these parameters.\n");
433     } else if (toobig == 40) {
434         av_log(s->avctx, AV_LOG_ERROR,
435                "[lavc rc] Error: bitrate too low for this video "
436                "with these parameters.\n");
437         return -1;
438     } else if (fabs(expected_bits / all_available_bits - 1.0) > 0.01) {
439         av_log(s->avctx, AV_LOG_ERROR,
440                "[lavc rc] Error: 2pass curve failed to converge\n");
441         return -1;
442     }
443
444     return 0;
445 }
446
447 av_cold int ff_rate_control_init(MpegEncContext *s)
448 {
449     RateControlContext *rcc = &s->rc_context;
450     int i, res;
451     static const char * const const_names[] = {
452         "PI",
453         "E",
454         "iTex",
455         "pTex",
456         "tex",
457         "mv",
458         "fCode",
459         "iCount",
460         "mcVar",
461         "var",
462         "isI",
463         "isP",
464         "isB",
465         "avgQP",
466         "qComp",
467         "avgIITex",
468         "avgPITex",
469         "avgPPTex",
470         "avgBPTex",
471         "avgTex",
472         NULL
473     };
474     static double (* const func1[])(void *, double) = {
475         (double (*)(void *, double)) bits2qp,
476         (double (*)(void *, double)) qp2bits,
477         NULL
478     };
479     static const char * const func1_names[] = {
480         "bits2qp",
481         "qp2bits",
482         NULL
483     };
484     emms_c();
485
486     res = av_expr_parse(&rcc->rc_eq_eval,
487                         s->rc_eq ? s->rc_eq : "tex^qComp",
488                         const_names, func1_names, func1,
489                         NULL, NULL, 0, s->avctx);
490     if (res < 0) {
491         av_log(s->avctx, AV_LOG_ERROR, "Error parsing rc_eq \"%s\"\n", s->rc_eq);
492         return res;
493     }
494
495     for (i = 0; i < 5; i++) {
496         rcc->pred[i].coeff = FF_QP2LAMBDA * 7.0;
497         rcc->pred[i].count = 1.0;
498         rcc->pred[i].decay = 0.4;
499
500         rcc->i_cplx_sum [i] =
501         rcc->p_cplx_sum [i] =
502         rcc->mv_bits_sum[i] =
503         rcc->qscale_sum [i] =
504         rcc->frame_count[i] = 1; // 1 is better because of 1/0 and such
505
506         rcc->last_qscale_for[i] = FF_QP2LAMBDA * 5;
507     }
508     rcc->buffer_index = s->avctx->rc_initial_buffer_occupancy;
509
510     if (s->avctx->flags & AV_CODEC_FLAG_PASS2) {
511         int i;
512         char *p;
513
514         /* find number of pics */
515         p = s->avctx->stats_in;
516         for (i = -1; p; i++)
517             p = strchr(p + 1, ';');
518         i += s->max_b_frames;
519         if (i <= 0 || i >= INT_MAX / sizeof(RateControlEntry))
520             return -1;
521         rcc->entry       = av_mallocz(i * sizeof(RateControlEntry));
522         rcc->num_entries = i;
523         if (!rcc->entry)
524             return AVERROR(ENOMEM);
525
526         /* init all to skipped P-frames
527          * (with B-frames we might have a not encoded frame at the end FIXME) */
528         for (i = 0; i < rcc->num_entries; i++) {
529             RateControlEntry *rce = &rcc->entry[i];
530
531             rce->pict_type  = rce->new_pict_type = AV_PICTURE_TYPE_P;
532             rce->qscale     = rce->new_qscale    = FF_QP2LAMBDA * 2;
533             rce->misc_bits  = s->mb_num + 10;
534             rce->mb_var_sum = s->mb_num * 100;
535         }
536
537         /* read stats */
538         p = s->avctx->stats_in;
539         for (i = 0; i < rcc->num_entries - s->max_b_frames; i++) {
540             RateControlEntry *rce;
541             int picture_number;
542             int e;
543             char *next;
544
545             next = strchr(p, ';');
546             if (next) {
547                 (*next) = 0; // sscanf is unbelievably slow on looong strings // FIXME copy / do not write
548                 next++;
549             }
550             e = sscanf(p, " in:%d ", &picture_number);
551
552             assert(picture_number >= 0);
553             assert(picture_number < rcc->num_entries);
554             rce = &rcc->entry[picture_number];
555
556             e += sscanf(p, " in:%*d out:%*d type:%d q:%f itex:%d ptex:%d mv:%d misc:%d fcode:%d bcode:%d mc-var:%d var:%d icount:%d skipcount:%d hbits:%d",
557                         &rce->pict_type, &rce->qscale, &rce->i_tex_bits, &rce->p_tex_bits,
558                         &rce->mv_bits, &rce->misc_bits,
559                         &rce->f_code, &rce->b_code,
560                         &rce->mc_mb_var_sum, &rce->mb_var_sum,
561                         &rce->i_count, &rce->skip_count, &rce->header_bits);
562             if (e != 14) {
563                 av_log(s->avctx, AV_LOG_ERROR,
564                        "statistics are damaged at line %d, parser out=%d\n",
565                        i, e);
566                 return -1;
567             }
568
569             p = next;
570         }
571
572         if (init_pass2(s) < 0) {
573             ff_rate_control_uninit(s);
574             return -1;
575         }
576     }
577
578     if (!(s->avctx->flags & AV_CODEC_FLAG_PASS2)) {
579         rcc->short_term_qsum   = 0.001;
580         rcc->short_term_qcount = 0.001;
581
582         rcc->pass1_rc_eq_output_sum = 0.001;
583         rcc->pass1_wanted_bits      = 0.001;
584
585         if (s->avctx->qblur > 1.0) {
586             av_log(s->avctx, AV_LOG_ERROR, "qblur too large\n");
587             return -1;
588         }
589         /* init stuff with the user specified complexity */
590         if (s->rc_initial_cplx) {
591             for (i = 0; i < 60 * 30; i++) {
592                 double bits = s->rc_initial_cplx * (i / 10000.0 + 1.0) * s->mb_num;
593                 RateControlEntry rce;
594
595                 if (i % ((s->gop_size + 3) / 4) == 0)
596                     rce.pict_type = AV_PICTURE_TYPE_I;
597                 else if (i % (s->max_b_frames + 1))
598                     rce.pict_type = AV_PICTURE_TYPE_B;
599                 else
600                     rce.pict_type = AV_PICTURE_TYPE_P;
601
602                 rce.new_pict_type = rce.pict_type;
603                 rce.mc_mb_var_sum = bits * s->mb_num / 100000;
604                 rce.mb_var_sum    = s->mb_num;
605
606                 rce.qscale    = FF_QP2LAMBDA * 2;
607                 rce.f_code    = 2;
608                 rce.b_code    = 1;
609                 rce.misc_bits = 1;
610
611                 if (s->pict_type == AV_PICTURE_TYPE_I) {
612                     rce.i_count    = s->mb_num;
613                     rce.i_tex_bits = bits;
614                     rce.p_tex_bits = 0;
615                     rce.mv_bits    = 0;
616                 } else {
617                     rce.i_count    = 0; // FIXME we do know this approx
618                     rce.i_tex_bits = 0;
619                     rce.p_tex_bits = bits * 0.9;
620                     rce.mv_bits    = bits * 0.1;
621                 }
622                 rcc->i_cplx_sum[rce.pict_type]  += rce.i_tex_bits * rce.qscale;
623                 rcc->p_cplx_sum[rce.pict_type]  += rce.p_tex_bits * rce.qscale;
624                 rcc->mv_bits_sum[rce.pict_type] += rce.mv_bits;
625                 rcc->frame_count[rce.pict_type]++;
626
627                 get_qscale(s, &rce, rcc->pass1_wanted_bits / rcc->pass1_rc_eq_output_sum, i);
628
629                 // FIXME misbehaves a little for variable fps
630                 rcc->pass1_wanted_bits += s->bit_rate / (1 / av_q2d(s->avctx->time_base));
631             }
632         }
633     }
634
635     return 0;
636 }
637
638 av_cold void ff_rate_control_uninit(MpegEncContext *s)
639 {
640     RateControlContext *rcc = &s->rc_context;
641     emms_c();
642
643     av_expr_free(rcc->rc_eq_eval);
644     av_freep(&rcc->entry);
645 }
646
647 int ff_vbv_update(MpegEncContext *s, int frame_size)
648 {
649     RateControlContext *rcc = &s->rc_context;
650     const double fps        = 1 / av_q2d(s->avctx->time_base);
651     const int buffer_size   = s->avctx->rc_buffer_size;
652     const double min_rate   = s->avctx->rc_min_rate / fps;
653     const double max_rate   = s->avctx->rc_max_rate / fps;
654
655     ff_dlog(s, "%d %f %d %f %f\n",
656             buffer_size, rcc->buffer_index, frame_size, min_rate, max_rate);
657
658     if (buffer_size) {
659         int left;
660
661         rcc->buffer_index -= frame_size;
662         if (rcc->buffer_index < 0) {
663             av_log(s->avctx, AV_LOG_ERROR, "rc buffer underflow\n");
664             rcc->buffer_index = 0;
665         }
666
667         left = buffer_size - rcc->buffer_index - 1;
668         rcc->buffer_index += av_clip(left, min_rate, max_rate);
669
670         if (rcc->buffer_index > buffer_size) {
671             int stuffing = ceil((rcc->buffer_index - buffer_size) / 8);
672
673             if (stuffing < 4 && s->codec_id == AV_CODEC_ID_MPEG4)
674                 stuffing = 4;
675             rcc->buffer_index -= 8 * stuffing;
676
677             if (s->avctx->debug & FF_DEBUG_RC)
678                 av_log(s->avctx, AV_LOG_DEBUG, "stuffing %d bytes\n", stuffing);
679
680             return stuffing;
681         }
682     }
683     return 0;
684 }
685
686 static double predict_size(Predictor *p, double q, double var)
687 {
688     return p->coeff * var / (q * p->count);
689 }
690
691 static void update_predictor(Predictor *p, double q, double var, double size)
692 {
693     double new_coeff = size * q / (var + 1);
694     if (var < 10)
695         return;
696
697     p->count *= p->decay;
698     p->coeff *= p->decay;
699     p->count++;
700     p->coeff += new_coeff;
701 }
702
703 static void adaptive_quantization(MpegEncContext *s, double q)
704 {
705     int i;
706     const float lumi_masking         = s->avctx->lumi_masking / (128.0 * 128.0);
707     const float dark_masking         = s->avctx->dark_masking / (128.0 * 128.0);
708     const float temp_cplx_masking    = s->avctx->temporal_cplx_masking;
709     const float spatial_cplx_masking = s->avctx->spatial_cplx_masking;
710     const float p_masking            = s->avctx->p_masking;
711     const float border_masking       = s->border_masking;
712     float bits_sum                   = 0.0;
713     float cplx_sum                   = 0.0;
714     float *cplx_tab                  = s->cplx_tab;
715     float *bits_tab                  = s->bits_tab;
716     const int qmin                   = s->avctx->mb_lmin;
717     const int qmax                   = s->avctx->mb_lmax;
718     Picture *const pic               = &s->current_picture;
719     const int mb_width               = s->mb_width;
720     const int mb_height              = s->mb_height;
721
722     for (i = 0; i < s->mb_num; i++) {
723         const int mb_xy = s->mb_index2xy[i];
724         float temp_cplx = sqrt(pic->mc_mb_var[mb_xy]); // FIXME merge in pow()
725         float spat_cplx = sqrt(pic->mb_var[mb_xy]);
726         const int lumi  = pic->mb_mean[mb_xy];
727         float bits, cplx, factor;
728         int mb_x = mb_xy % s->mb_stride;
729         int mb_y = mb_xy / s->mb_stride;
730         int mb_distance;
731         float mb_factor = 0.0;
732         if (spat_cplx < 4)
733             spat_cplx = 4;              // FIXME fine-tune
734         if (temp_cplx < 4)
735             temp_cplx = 4;              // FIXME fine-tune
736
737         if ((s->mb_type[mb_xy] & CANDIDATE_MB_TYPE_INTRA)) { // FIXME hq mode
738             cplx   = spat_cplx;
739             factor = 1.0 + p_masking;
740         } else {
741             cplx   = temp_cplx;
742             factor = pow(temp_cplx, -temp_cplx_masking);
743         }
744         factor *= pow(spat_cplx, -spatial_cplx_masking);
745
746         if (lumi > 127)
747             factor *= (1.0 - (lumi - 128) * (lumi - 128) * lumi_masking);
748         else
749             factor *= (1.0 - (lumi - 128) * (lumi - 128) * dark_masking);
750
751         if (mb_x < mb_width / 5) {
752             mb_distance = mb_width / 5 - mb_x;
753             mb_factor   = (float)mb_distance / (float)(mb_width / 5);
754         } else if (mb_x > 4 * mb_width / 5) {
755             mb_distance = mb_x - 4 * mb_width / 5;
756             mb_factor   = (float)mb_distance / (float)(mb_width / 5);
757         }
758         if (mb_y < mb_height / 5) {
759             mb_distance = mb_height / 5 - mb_y;
760             mb_factor   = FFMAX(mb_factor,
761                                 (float)mb_distance / (float)(mb_height / 5));
762         } else if (mb_y > 4 * mb_height / 5) {
763             mb_distance = mb_y - 4 * mb_height / 5;
764             mb_factor   = FFMAX(mb_factor,
765                                 (float)mb_distance / (float)(mb_height / 5));
766         }
767
768         factor *= 1.0 - border_masking * mb_factor;
769
770         if (factor < 0.00001)
771             factor = 0.00001;
772
773         bits        = cplx * factor;
774         cplx_sum   += cplx;
775         bits_sum   += bits;
776         cplx_tab[i] = cplx;
777         bits_tab[i] = bits;
778     }
779
780     /* handle qmin/qmax clipping */
781     if (s->mpv_flags & FF_MPV_FLAG_NAQ) {
782         float factor = bits_sum / cplx_sum;
783         for (i = 0; i < s->mb_num; i++) {
784             float newq = q * cplx_tab[i] / bits_tab[i];
785             newq *= factor;
786
787             if (newq > qmax) {
788                 bits_sum -= bits_tab[i];
789                 cplx_sum -= cplx_tab[i] * q / qmax;
790             } else if (newq < qmin) {
791                 bits_sum -= bits_tab[i];
792                 cplx_sum -= cplx_tab[i] * q / qmin;
793             }
794         }
795         if (bits_sum < 0.001)
796             bits_sum = 0.001;
797         if (cplx_sum < 0.001)
798             cplx_sum = 0.001;
799     }
800
801     for (i = 0; i < s->mb_num; i++) {
802         const int mb_xy = s->mb_index2xy[i];
803         float newq      = q * cplx_tab[i] / bits_tab[i];
804         int intq;
805
806         if (s->mpv_flags & FF_MPV_FLAG_NAQ) {
807             newq *= bits_sum / cplx_sum;
808         }
809
810         intq = (int)(newq + 0.5);
811
812         if (intq > qmax)
813             intq = qmax;
814         else if (intq < qmin)
815             intq = qmin;
816         s->lambda_table[mb_xy] = intq;
817     }
818 }
819
820 void ff_get_2pass_fcode(MpegEncContext *s)
821 {
822     RateControlContext *rcc = &s->rc_context;
823     RateControlEntry *rce   = &rcc->entry[s->picture_number];
824
825     s->f_code = rce->f_code;
826     s->b_code = rce->b_code;
827 }
828
829 // FIXME rd or at least approx for dquant
830
831 float ff_rate_estimate_qscale(MpegEncContext *s, int dry_run)
832 {
833     float q;
834     int qmin, qmax;
835     float br_compensation;
836     double diff;
837     double short_term_q;
838     double fps;
839     int picture_number = s->picture_number;
840     int64_t wanted_bits;
841     RateControlContext *rcc = &s->rc_context;
842     AVCodecContext *a       = s->avctx;
843     RateControlEntry local_rce, *rce;
844     double bits;
845     double rate_factor;
846     int var;
847     const int pict_type = s->pict_type;
848     Picture * const pic = &s->current_picture;
849     emms_c();
850
851     get_qminmax(&qmin, &qmax, s, pict_type);
852
853     fps = 1 / av_q2d(s->avctx->time_base);
854     /* update predictors */
855     if (picture_number > 2 && !dry_run) {
856         const int last_var = s->last_pict_type == AV_PICTURE_TYPE_I ? rcc->last_mb_var_sum
857                                                                     : rcc->last_mc_mb_var_sum;
858         update_predictor(&rcc->pred[s->last_pict_type],
859                          rcc->last_qscale,
860                          sqrt(last_var), s->frame_bits);
861     }
862
863     if (s->avctx->flags & AV_CODEC_FLAG_PASS2) {
864         assert(picture_number >= 0);
865         assert(picture_number < rcc->num_entries);
866         rce         = &rcc->entry[picture_number];
867         wanted_bits = rce->expected_bits;
868     } else {
869         Picture *dts_pic;
870         rce = &local_rce;
871
872         /* FIXME add a dts field to AVFrame and ensure it is set and use it
873          * here instead of reordering but the reordering is simpler for now
874          * until H.264 B-pyramid must be handled. */
875         if (s->pict_type == AV_PICTURE_TYPE_B || s->low_delay)
876             dts_pic = s->current_picture_ptr;
877         else
878             dts_pic = s->last_picture_ptr;
879
880         if (!dts_pic || dts_pic->f->pts == AV_NOPTS_VALUE)
881             wanted_bits = (uint64_t)(s->bit_rate * (double)picture_number / fps);
882         else
883             wanted_bits = (uint64_t)(s->bit_rate * (double)dts_pic->f->pts / fps);
884     }
885
886     diff = s->total_bits - wanted_bits;
887     br_compensation = (a->bit_rate_tolerance - diff) / a->bit_rate_tolerance;
888     if (br_compensation <= 0.0)
889         br_compensation = 0.001;
890
891     var = pict_type == AV_PICTURE_TYPE_I ? pic->mb_var_sum : pic->mc_mb_var_sum;
892
893     short_term_q = 0; /* avoid warning */
894     if (s->avctx->flags & AV_CODEC_FLAG_PASS2) {
895         if (pict_type != AV_PICTURE_TYPE_I)
896             assert(pict_type == rce->new_pict_type);
897
898         q = rce->new_qscale / br_compensation;
899         ff_dlog(s, "%f %f %f last:%d var:%d type:%d//\n", q, rce->new_qscale,
900                 br_compensation, s->frame_bits, var, pict_type);
901     } else {
902         rce->pict_type     =
903         rce->new_pict_type = pict_type;
904         rce->mc_mb_var_sum = pic->mc_mb_var_sum;
905         rce->mb_var_sum    = pic->mb_var_sum;
906         rce->qscale        = FF_QP2LAMBDA * 2;
907         rce->f_code        = s->f_code;
908         rce->b_code        = s->b_code;
909         rce->misc_bits     = 1;
910
911         bits = predict_size(&rcc->pred[pict_type], rce->qscale, sqrt(var));
912         if (pict_type == AV_PICTURE_TYPE_I) {
913             rce->i_count    = s->mb_num;
914             rce->i_tex_bits = bits;
915             rce->p_tex_bits = 0;
916             rce->mv_bits    = 0;
917         } else {
918             rce->i_count    = 0;    // FIXME we do know this approx
919             rce->i_tex_bits = 0;
920             rce->p_tex_bits = bits * 0.9;
921             rce->mv_bits    = bits * 0.1;
922         }
923         rcc->i_cplx_sum[pict_type]  += rce->i_tex_bits * rce->qscale;
924         rcc->p_cplx_sum[pict_type]  += rce->p_tex_bits * rce->qscale;
925         rcc->mv_bits_sum[pict_type] += rce->mv_bits;
926         rcc->frame_count[pict_type]++;
927
928         bits        = rce->i_tex_bits + rce->p_tex_bits;
929         rate_factor = rcc->pass1_wanted_bits /
930                       rcc->pass1_rc_eq_output_sum * br_compensation;
931
932         q = get_qscale(s, rce, rate_factor, picture_number);
933         if (q < 0)
934             return -1;
935
936         assert(q > 0.0);
937         q = get_diff_limited_q(s, rce, q);
938         assert(q > 0.0);
939
940         // FIXME type dependent blur like in 2-pass
941         if (pict_type == AV_PICTURE_TYPE_P || s->intra_only) {
942             rcc->short_term_qsum   *= a->qblur;
943             rcc->short_term_qcount *= a->qblur;
944
945             rcc->short_term_qsum += q;
946             rcc->short_term_qcount++;
947             q = short_term_q = rcc->short_term_qsum / rcc->short_term_qcount;
948         }
949         assert(q > 0.0);
950
951         q = modify_qscale(s, rce, q, picture_number);
952
953         rcc->pass1_wanted_bits += s->bit_rate / fps;
954
955         assert(q > 0.0);
956     }
957
958     if (s->avctx->debug & FF_DEBUG_RC) {
959         av_log(s->avctx, AV_LOG_DEBUG,
960                "%c qp:%d<%2.1f<%d %d want:%d total:%d comp:%f st_q:%2.2f "
961                "size:%d var:%d/%d br:%d fps:%d\n",
962                av_get_picture_type_char(pict_type),
963                qmin, q, qmax, picture_number,
964                (int)wanted_bits / 1000, (int)s->total_bits / 1000,
965                br_compensation, short_term_q, s->frame_bits,
966                pic->mb_var_sum, pic->mc_mb_var_sum,
967                s->bit_rate / 1000, (int)fps);
968     }
969
970     if (q < qmin)
971         q = qmin;
972     else if (q > qmax)
973         q = qmax;
974
975     if (s->adaptive_quant)
976         adaptive_quantization(s, q);
977     else
978         q = (int)(q + 0.5);
979
980     if (!dry_run) {
981         rcc->last_qscale        = q;
982         rcc->last_mc_mb_var_sum = pic->mc_mb_var_sum;
983         rcc->last_mb_var_sum    = pic->mb_var_sum;
984     }
985     return q;
986 }