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