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