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