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