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