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 }