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