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