various security fixes and precautionary checks
[libav.git] / libavcodec / ratecontrol.c
CommitLineData
8b4c7dbc 1/*
2ef0f2b2
FB
2 * Rate control for video encoders
3 *
8f2ab833 4 * Copyright (c) 2002-2004 Michael Niedermayer <michaelni@gmx.at>
2ef0f2b2
FB
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 */
983e3246
MN
20
21/**
22 * @file ratecontrol.c
23 * Rate control for video encoders.
24 */
25
8b4c7dbc 26#include "avcodec.h"
2ef0f2b2 27#include "dsputil.h"
8b4c7dbc
MN
28#include "mpegvideo.h"
29
3aa102be
MN
30#undef NDEBUG // allways check asserts, the speed effect is far too small to disable them
31#include <assert.h>
8b4c7dbc 32
471d7dc3
MN
33#ifndef M_E
34#define M_E 2.718281828
35#endif
36
8b4c7dbc 37static int init_pass2(MpegEncContext *s);
3aa102be 38static double get_qscale(MpegEncContext *s, RateControlEntry *rce, double rate_factor, int frame_num);
8b4c7dbc
MN
39
40void ff_write_pass1_stats(MpegEncContext *s){
158c7f05 41 sprintf(s->avctx->stats_out, "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",
1031fabd 42 s->current_picture_ptr->display_picture_number, s->current_picture_ptr->coded_picture_number, s->pict_type,
158c7f05 43 s->current_picture.quality, s->i_tex_bits, s->p_tex_bits, s->mv_bits, s->misc_bits,
1e491e29 44 s->f_code, s->b_code, s->current_picture.mc_mb_var_sum, s->current_picture.mb_var_sum, s->i_count);
8b4c7dbc
MN
45}
46
47int ff_rate_control_init(MpegEncContext *s)
48{
49 RateControlContext *rcc= &s->rc_context;
3aa102be 50 int i;
8b4c7dbc
MN
51 emms_c();
52
3aa102be 53 for(i=0; i<5; i++){
158c7f05 54 rcc->pred[i].coeff= FF_QP2LAMBDA * 7.0;
3aa102be
MN
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
158c7f05 63 rcc->last_qscale_for[i]=FF_QP2LAMBDA * 5;
3aa102be 64 }
fb066639 65 rcc->buffer_index= s->avctx->rc_initial_buffer_occupancy;
3aa102be 66
3aa102be 67 if(s->flags&CODEC_FLAG_PASS2){
8b4c7dbc 68 int i;
3aa102be 69 char *p;
8b4c7dbc 70
3aa102be
MN
71 /* find number of pics */
72 p= s->avctx->stats_in;
73 for(i=-1; p; i++){
74 p= strchr(p+1, ';');
8b4c7dbc 75 }
3aa102be 76 i+= s->max_b_frames;
0ecca7a4
MN
77 if(i<=0 || i>=INT_MAX / sizeof(RateControlEntry))
78 return -1;
3aa102be
MN
79 rcc->entry = (RateControlEntry*)av_mallocz(i*sizeof(RateControlEntry));
80 rcc->num_entries= i;
81
82 /* init all to skiped 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;
158c7f05 86 rce->qscale= rce->new_qscale=FF_QP2LAMBDA * 2;
3aa102be
MN
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++){
8b4c7dbc
MN
94 RateControlEntry *rce;
95 int picture_number;
96 int e;
3aa102be
MN
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);
8b4c7dbc 108 rce= &rcc->entry[picture_number];
3aa102be 109
c5d309f2 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",
3aa102be
MN
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){
9b879566 114 av_log(s->avctx, AV_LOG_ERROR, "statistics are damaged at line %d, parser out=%d\n", i, e);
8b4c7dbc
MN
115 return -1;
116 }
3aa102be 117 p= next;
8b4c7dbc 118 }
8b4c7dbc
MN
119
120 if(init_pass2(s) < 0) return -1;
121 }
122
3aa102be
MN
123 if(!(s->flags&CODEC_FLAG_PASS2)){
124
125 rcc->short_term_qsum=0.001;
126 rcc->short_term_qcount=0.001;
8b4c7dbc 127
23e54f69 128 rcc->pass1_rc_eq_output_sum= 0.001;
3aa102be
MN
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;
158c7f05 145 rce.qscale = FF_QP2LAMBDA * 2;
3aa102be
MN
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] ++;
8b4c7dbc 165
3aa102be 166 bits= rce.i_tex_bits + rce.p_tex_bits;
8b4c7dbc 167
23e54f69 168 q= get_qscale(s, &rce, rcc->pass1_wanted_bits/rcc->pass1_rc_eq_output_sum, i);
14bea432 169 rcc->pass1_wanted_bits+= s->bit_rate/(s->avctx->frame_rate / (double)s->avctx->frame_rate_base);
3aa102be
MN
170 }
171 }
172
173 }
174
8b4c7dbc
MN
175 return 0;
176}
177
178void ff_rate_control_uninit(MpegEncContext *s)
179{
180 RateControlContext *rcc= &s->rc_context;
181 emms_c();
182
6000abfa 183 av_freep(&rcc->entry);
8b4c7dbc
MN
184}
185
3aa102be
MN
186static inline double qp2bits(RateControlEntry *rce, double qp){
187 if(qp<=0.0){
9b879566 188 av_log(NULL, AV_LOG_ERROR, "qp<=0.0\n");
3aa102be
MN
189 }
190 return rce->qscale * (double)(rce->i_tex_bits + rce->p_tex_bits+1)/ qp;
191}
192
193static inline double bits2qp(RateControlEntry *rce, double bits){
194 if(bits<0.9){
9b879566 195 av_log(NULL, AV_LOG_ERROR, "bits<0.9\n");
3aa102be
MN
196 }
197 return rce->qscale * (double)(rce->i_tex_bits + rce->p_tex_bits+1)/ bits;
198}
199
11dffbe1 200int ff_vbv_update(MpegEncContext *s, int frame_size){
3aa102be 201 RateControlContext *rcc= &s->rc_context;
14bea432 202 const double fps= (double)s->avctx->frame_rate / (double)s->avctx->frame_rate_base;
d60a8f85 203 const int buffer_size= s->avctx->rc_buffer_size;
3aa102be
MN
204 const double min_rate= s->avctx->rc_min_rate/fps;
205 const double max_rate= s->avctx->rc_max_rate/fps;
d60a8f85
MN
206
207//printf("%d %f %d %f %f\n", buffer_size, rcc->buffer_index, frame_size, min_rate, max_rate);
3aa102be 208 if(buffer_size){
fb066639
MN
209 int left;
210
3aa102be 211 rcc->buffer_index-= frame_size;
fb066639 212 if(rcc->buffer_index < 0){
9b879566 213 av_log(s->avctx, AV_LOG_ERROR, "rc buffer underflow\n");
fb066639
MN
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
d60a8f85
MN
220 if(rcc->buffer_index > buffer_size){
221 int stuffing= ceil((rcc->buffer_index - buffer_size)/8);
11dffbe1
MN
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;
fb066639 231 }
3aa102be 232 }
11dffbe1 233 return 0;
3aa102be
MN
234}
235
236/**
237 * modifies the bitrate curve from pass1 for one frame
238 */
239static double get_qscale(MpegEncContext *s, RateControlEntry *rce, double rate_factor, int frame_num){
240 RateControlContext *rcc= &s->rc_context;
d55f7b65 241 AVCodecContext *a= s->avctx;
3aa102be
MN
242 double q, bits;
243 const int pict_type= rce->new_pict_type;
244 const double mb_num= s->mb_num;
245 int i;
3aa102be
MN
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],
d55f7b65 262 a->qcompress,
3aa102be
MN
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 };
5c91a675 274 static const char *const_names[]={
3aa102be
MN
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 };
07787186 301 static double (*func1[])(void *, double)={
ec6a3752
FB
302 (void *)bits2qp,
303 (void *)qp2bits,
3aa102be
MN
304 NULL
305 };
5c91a675 306 static const char *func1_names[]={
3aa102be
MN
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
23e54f69 314 rcc->pass1_rc_eq_output_sum+= bits;
3aa102be
MN
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;
6a1f7e7b
MN
338
339 return q;
340}
341
342static 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];
59b571c1 348
6a1f7e7b
MN
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
3aa102be 354 /* last qscale / qdiff stuff */
6a1f7e7b
MN
355 if(rcc->last_non_b_pict_type==pict_type || pict_type!=I_TYPE){
356 double last_q= rcc->last_qscale_for[pict_type];
158c7f05 357 const int maxdiff= FF_QP2LAMBDA * a->max_qdiff;
59b571c1 358
158c7f05
MN
359 if (q > last_q + maxdiff) q= last_q + maxdiff;
360 else if(q < last_q - maxdiff) q= last_q - maxdiff;
6a1f7e7b 361 }
3aa102be
MN
362
363 rcc->last_qscale_for[pict_type]= q; //Note we cant do that after blurring
364
6a1f7e7b
MN
365 if(pict_type!=B_TYPE)
366 rcc->last_non_b_pict_type= pict_type;
367
3aa102be
MN
368 return q;
369}
370
371/**
372 * gets the qmin & qmax for pict_type
373 */
374static void get_qminmax(int *qmin_ret, int *qmax_ret, MpegEncContext *s, int pict_type){
158c7f05
MN
375 int qmin= s->avctx->lmin;
376 int qmax= s->avctx->lmax;
fa12b546
MN
377
378 assert(qmin <= qmax);
3aa102be
MN
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
158c7f05
MN
388 qmin= clip(qmin, 1, FF_LAMBDA_MAX);
389 qmax= clip(qmax, 1, FF_LAMBDA_MAX);
3aa102be 390
fa12b546 391 if(qmax<qmin) qmax= qmin;
3aa102be
MN
392
393 *qmin_ret= qmin;
394 *qmax_ret= qmax;
395}
396
397static 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;
fb066639
MN
403 const double fps= (double)s->avctx->frame_rate / (double)s->avctx->frame_rate_base;
404 const double min_rate= s->avctx->rc_min_rate / fps;
405 const double max_rate= s->avctx->rc_max_rate / fps;
3aa102be
MN
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);
946c8a12 414//printf("q:%f\n", q);
3aa102be
MN
415 /* buffer overflow/underflow protection */
416 if(buffer_size){
6a1f7e7b 417 double expected_size= rcc->buffer_index;
d60a8f85 418 double q_limit;
3aa102be
MN
419
420 if(min_rate){
6a1f7e7b 421 double d= 2*(buffer_size - expected_size)/buffer_size;
3aa102be 422 if(d>1.0) d=1.0;
946c8a12
MN
423 else if(d<0.0001) d=0.0001;
424 q*= pow(d, 1.0/s->avctx->rc_buffer_aggressivity);
6a1f7e7b 425
d60a8f85
MN
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 }
3aa102be
MN
433 }
434
435 if(max_rate){
436 double d= 2*expected_size/buffer_size;
437 if(d>1.0) d=1.0;
946c8a12
MN
438 else if(d<0.0001) d=0.0001;
439 q/= pow(d, 1.0/s->avctx->rc_buffer_aggressivity);
6a1f7e7b 440
d60a8f85
MN
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 }
3aa102be
MN
448 }
449 }
946c8a12 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);
c695ca3b 451 if(s->avctx->rc_qsquish==0.0 || qmin==qmax){
3aa102be
MN
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 }
c5d309f2 466
3aa102be
MN
467 return q;
468}
469
8b4c7dbc
MN
470//----------------------------------
471// 1 Pass Code
472
3aa102be 473static double predict_size(Predictor *p, double q, double var)
8b4c7dbc
MN
474{
475 return p->coeff*var / (q*p->count);
476}
477
ccfddafb 478/*
3aa102be
MN
479static 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}
ccfddafb 484*/
3aa102be 485
8b4c7dbc
MN
486static void update_predictor(Predictor *p, double q, double var, double size)
487{
488 double new_coeff= size*q / (var + 1);
3aa102be 489 if(var<10) return;
8b4c7dbc
MN
490
491 p->count*= p->decay;
492 p->coeff*= p->decay;
493 p->count++;
494 p->coeff+= new_coeff;
495}
496
c5d309f2
MN
497static void adaptive_quantization(MpegEncContext *s, double q){
498 int i;
499 const float lumi_masking= s->avctx->lumi_masking / (128.0*128.0);
5e746b99 500 const float dark_masking= s->avctx->dark_masking / (128.0*128.0);
c5d309f2
MN
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 float bits_sum= 0.0;
505 float cplx_sum= 0.0;
506 float cplx_tab[s->mb_num];
507 float bits_tab[s->mb_num];
158c7f05
MN
508 const int qmin= s->avctx->lmin;
509 const int qmax= s->avctx->lmax;
1e491e29 510 Picture * const pic= &s->current_picture;
c5d309f2
MN
511
512 for(i=0; i<s->mb_num; i++){
2f5feea4 513 const int mb_xy= s->mb_index2xy[i];
158c7f05 514 float temp_cplx= sqrt(pic->mc_mb_var[mb_xy]); //FIXME merge in pow()
2f5feea4
MN
515 float spat_cplx= sqrt(pic->mb_var[mb_xy]);
516 const int lumi= pic->mb_mean[mb_xy];
c5d309f2 517 float bits, cplx, factor;
158c7f05 518#if 0
c5d309f2
MN
519 if(spat_cplx < q/3) spat_cplx= q/3; //FIXME finetune
520 if(temp_cplx < q/3) temp_cplx= q/3; //FIXME finetune
158c7f05
MN
521#endif
522 if(spat_cplx < 4) spat_cplx= 4; //FIXME finetune
523 if(temp_cplx < 4) temp_cplx= 4; //FIXME finetune
524
bb198e19 525 if((s->mb_type[mb_xy]&CANDIDATE_MB_TYPE_INTRA)){//FIXME hq mode
c5d309f2
MN
526 cplx= spat_cplx;
527 factor= 1.0 + p_masking;
528 }else{
529 cplx= temp_cplx;
530 factor= pow(temp_cplx, - temp_cplx_masking);
531 }
532 factor*=pow(spat_cplx, - spatial_cplx_masking);
5e746b99
MN
533
534 if(lumi>127)
535 factor*= (1.0 - (lumi-128)*(lumi-128)*lumi_masking);
536 else
537 factor*= (1.0 - (lumi-128)*(lumi-128)*dark_masking);
c5d309f2
MN
538
539 if(factor<0.00001) factor= 0.00001;
540
541 bits= cplx*factor;
542 cplx_sum+= cplx;
543 bits_sum+= bits;
544 cplx_tab[i]= cplx;
545 bits_tab[i]= bits;
546 }
547
548 /* handle qmin/qmax cliping */
549 if(s->flags&CODEC_FLAG_NORMALIZE_AQP){
e9a4834a 550 float factor= bits_sum/cplx_sum;
c5d309f2
MN
551 for(i=0; i<s->mb_num; i++){
552 float newq= q*cplx_tab[i]/bits_tab[i];
e9a4834a 553 newq*= factor;
c5d309f2
MN
554
555 if (newq > qmax){
556 bits_sum -= bits_tab[i];
557 cplx_sum -= cplx_tab[i]*q/qmax;
558 }
559 else if(newq < qmin){
560 bits_sum -= bits_tab[i];
561 cplx_sum -= cplx_tab[i]*q/qmin;
562 }
563 }
e9a4834a
MN
564 if(bits_sum < 0.001) bits_sum= 0.001;
565 if(cplx_sum < 0.001) cplx_sum= 0.001;
c5d309f2
MN
566 }
567
568 for(i=0; i<s->mb_num; i++){
2f5feea4 569 const int mb_xy= s->mb_index2xy[i];
c5d309f2
MN
570 float newq= q*cplx_tab[i]/bits_tab[i];
571 int intq;
572
573 if(s->flags&CODEC_FLAG_NORMALIZE_AQP){
574 newq*= bits_sum/cplx_sum;
575 }
576
158c7f05 577 intq= (int)(newq + 0.5);
c5d309f2
MN
578
579 if (intq > qmax) intq= qmax;
580 else if(intq < qmin) intq= qmin;
581//if(i%s->mb_width==0) printf("\n");
582//printf("%2d%3d ", intq, ff_sqrt(s->mc_mb_var[i]));
158c7f05 583 s->lambda_table[mb_xy]= intq;
c5d309f2
MN
584 }
585}
158c7f05 586//FIXME rd or at least approx for dquant
c5d309f2
MN
587
588float ff_rate_estimate_qscale(MpegEncContext *s)
8b4c7dbc 589{
8b4c7dbc 590 float q;
c5d309f2 591 int qmin, qmax;
8b4c7dbc
MN
592 float br_compensation;
593 double diff;
594 double short_term_q;
8b4c7dbc 595 double fps;
3aa102be 596 int picture_number= s->picture_number;
8b4c7dbc 597 int64_t wanted_bits;
3aa102be 598 RateControlContext *rcc= &s->rc_context;
d55f7b65 599 AVCodecContext *a= s->avctx;
3aa102be
MN
600 RateControlEntry local_rce, *rce;
601 double bits;
602 double rate_factor;
603 int var;
604 const int pict_type= s->pict_type;
1e491e29 605 Picture * const pic= &s->current_picture;
8b4c7dbc
MN
606 emms_c();
607
3aa102be 608 get_qminmax(&qmin, &qmax, s, pict_type);
8b4c7dbc 609
14bea432 610 fps= (double)s->avctx->frame_rate / (double)s->avctx->frame_rate_base;
946c8a12 611//printf("input_pic_num:%d pic_num:%d frame_rate:%d\n", s->input_picture_number, s->picture_number, s->frame_rate);
8b4c7dbc
MN
612 /* update predictors */
613 if(picture_number>2){
3aa102be
MN
614 const int last_var= s->last_pict_type == I_TYPE ? rcc->last_mb_var_sum : rcc->last_mc_mb_var_sum;
615 update_predictor(&rcc->pred[s->last_pict_type], rcc->last_qscale, sqrt(last_var), s->frame_bits);
8b4c7dbc
MN
616 }
617
3aa102be
MN
618 if(s->flags&CODEC_FLAG_PASS2){
619 assert(picture_number>=0);
620 assert(picture_number<rcc->num_entries);
621 rce= &rcc->entry[picture_number];
622 wanted_bits= rce->expected_bits;
623 }else{
624 rce= &local_rce;
625 wanted_bits= (uint64_t)(s->bit_rate*(double)picture_number/fps);
626 }
8b4c7dbc 627
3aa102be 628 diff= s->total_bits - wanted_bits;
d55f7b65 629 br_compensation= (a->bit_rate_tolerance - diff)/a->bit_rate_tolerance;
3aa102be
MN
630 if(br_compensation<=0.0) br_compensation=0.001;
631
1e491e29 632 var= pict_type == I_TYPE ? pic->mb_var_sum : pic->mc_mb_var_sum;
3aa102be 633
b1609412 634 short_term_q = 0; /* avoid warning */
3aa102be
MN
635 if(s->flags&CODEC_FLAG_PASS2){
636 if(pict_type!=I_TYPE)
637 assert(pict_type == rce->new_pict_type);
638
639 q= rce->new_qscale / br_compensation;
640//printf("%f %f %f last:%d var:%d type:%d//\n", q, rce->new_qscale, br_compensation, s->frame_bits, var, pict_type);
641 }else{
642 rce->pict_type=
643 rce->new_pict_type= pict_type;
1e491e29
MN
644 rce->mc_mb_var_sum= pic->mc_mb_var_sum;
645 rce->mb_var_sum = pic-> mb_var_sum;
158c7f05 646 rce->qscale = FF_QP2LAMBDA * 2;
3aa102be
MN
647 rce->f_code = s->f_code;
648 rce->b_code = s->b_code;
649 rce->misc_bits= 1;
650
3aa102be
MN
651 bits= predict_size(&rcc->pred[pict_type], rce->qscale, sqrt(var));
652 if(pict_type== I_TYPE){
653 rce->i_count = s->mb_num;
654 rce->i_tex_bits= bits;
655 rce->p_tex_bits= 0;
656 rce->mv_bits= 0;
657 }else{
658 rce->i_count = 0; //FIXME we do know this approx
659 rce->i_tex_bits= 0;
660 rce->p_tex_bits= bits*0.9;
661
662 rce->mv_bits= bits*0.1;
8b4c7dbc 663 }
3aa102be
MN
664 rcc->i_cplx_sum [pict_type] += rce->i_tex_bits*rce->qscale;
665 rcc->p_cplx_sum [pict_type] += rce->p_tex_bits*rce->qscale;
666 rcc->mv_bits_sum[pict_type] += rce->mv_bits;
667 rcc->frame_count[pict_type] ++;
8b4c7dbc 668
3aa102be 669 bits= rce->i_tex_bits + rce->p_tex_bits;
23e54f69 670 rate_factor= rcc->pass1_wanted_bits/rcc->pass1_rc_eq_output_sum * br_compensation;
8b4c7dbc 671
3aa102be 672 q= get_qscale(s, rce, rate_factor, picture_number);
8b4c7dbc 673
c695ca3b 674 assert(q>0.0);
3aa102be 675//printf("%f ", q);
6a1f7e7b 676 q= get_diff_limited_q(s, rce, q);
3aa102be
MN
677//printf("%f ", q);
678 assert(q>0.0);
679
680 if(pict_type==P_TYPE || s->intra_only){ //FIXME type dependant blur like in 2-pass
d55f7b65
MN
681 rcc->short_term_qsum*=a->qblur;
682 rcc->short_term_qcount*=a->qblur;
3aa102be
MN
683
684 rcc->short_term_qsum+= q;
685 rcc->short_term_qcount++;
686//printf("%f ", q);
687 q= short_term_q= rcc->short_term_qsum/rcc->short_term_qcount;
688//printf("%f ", q);
689 }
946c8a12
MN
690 assert(q>0.0);
691
3aa102be
MN
692 q= modify_qscale(s, rce, q, picture_number);
693
694 rcc->pass1_wanted_bits+= s->bit_rate/fps;
695
c695ca3b 696 assert(q>0.0);
8b4c7dbc 697 }
59b571c1
MN
698
699 if(s->avctx->debug&FF_DEBUG_RC){
9b879566 700 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",
d8085ea7 701 av_get_pict_type_char(pict_type), qmin, q, qmax, picture_number, (int)wanted_bits/1000, (int)s->total_bits/1000,
59b571c1
MN
702 br_compensation, short_term_q, s->frame_bits, pic->mb_var_sum, pic->mc_mb_var_sum, s->bit_rate/1000, (int)fps
703 );
704 }
3aa102be
MN
705
706 if (q<qmin) q=qmin;
707 else if(q>qmax) q=qmax;
3aa102be 708
c5d309f2
MN
709 if(s->adaptive_quant)
710 adaptive_quantization(s, q);
711 else
712 q= (int)(q + 0.5);
713
714 rcc->last_qscale= q;
1e491e29
MN
715 rcc->last_mc_mb_var_sum= pic->mc_mb_var_sum;
716 rcc->last_mb_var_sum= pic->mb_var_sum;
717#if 0
718{
719 static int mvsum=0, texsum=0;
720 mvsum += s->mv_bits;
721 texsum += s->i_tex_bits + s->p_tex_bits;
722 printf("%d %d//\n\n", mvsum, texsum);
723}
724#endif
c5d309f2 725 return q;
8b4c7dbc
MN
726}
727
728//----------------------------------------------
729// 2-Pass code
730
731static int init_pass2(MpegEncContext *s)
732{
733 RateControlContext *rcc= &s->rc_context;
d55f7b65 734 AVCodecContext *a= s->avctx;
8b4c7dbc 735 int i;
14bea432 736 double fps= (double)s->avctx->frame_rate / (double)s->avctx->frame_rate_base;
8b4c7dbc
MN
737 double complexity[5]={0,0,0,0,0}; // aproximate bits at quant=1
738 double avg_quantizer[5];
739 uint64_t const_bits[5]={0,0,0,0,0}; // quantizer idependant bits
740 uint64_t available_bits[5];
741 uint64_t all_const_bits;
742 uint64_t all_available_bits= (uint64_t)(s->bit_rate*(double)rcc->num_entries/fps);
8b4c7dbc
MN
743 double rate_factor=0;
744 double step;
ccfddafb 745 //int last_i_frame=-10000000;
d55f7b65 746 const int filter_size= (int)(a->qblur*4) | 1;
3aa102be
MN
747 double expected_bits;
748 double *qscale, *blured_qscale;
8b4c7dbc
MN
749
750 /* find complexity & const_bits & decide the pict_types */
751 for(i=0; i<rcc->num_entries; i++){
752 RateControlEntry *rce= &rcc->entry[i];
753
0d1e9246 754 rce->new_pict_type= rce->pict_type;
3aa102be
MN
755 rcc->i_cplx_sum [rce->pict_type] += rce->i_tex_bits*rce->qscale;
756 rcc->p_cplx_sum [rce->pict_type] += rce->p_tex_bits*rce->qscale;
757 rcc->mv_bits_sum[rce->pict_type] += rce->mv_bits;
758 rcc->frame_count[rce->pict_type] ++;
8b4c7dbc
MN
759
760 complexity[rce->new_pict_type]+= (rce->i_tex_bits+ rce->p_tex_bits)*(double)rce->qscale;
761 const_bits[rce->new_pict_type]+= rce->mv_bits + rce->misc_bits;
8b4c7dbc
MN
762 }
763 all_const_bits= const_bits[I_TYPE] + const_bits[P_TYPE] + const_bits[B_TYPE];
764
765 if(all_available_bits < all_const_bits){
9b879566 766 av_log(s->avctx, AV_LOG_ERROR, "requested bitrate is to low\n");
8b4c7dbc
MN
767 return -1;
768 }
3aa102be
MN
769
770 /* find average quantizers */
771 avg_quantizer[P_TYPE]=0;
772 for(step=256*256; step>0.0000001; step*=0.5){
773 double expected_bits=0;
774 avg_quantizer[P_TYPE]+= step;
775
776 avg_quantizer[I_TYPE]= avg_quantizer[P_TYPE]*ABS(s->avctx->i_quant_factor) + s->avctx->i_quant_offset;
777 avg_quantizer[B_TYPE]= avg_quantizer[P_TYPE]*ABS(s->avctx->b_quant_factor) + s->avctx->b_quant_offset;
778
779 expected_bits=
780 + all_const_bits
781 + complexity[I_TYPE]/avg_quantizer[I_TYPE]
782 + complexity[P_TYPE]/avg_quantizer[P_TYPE]
783 + complexity[B_TYPE]/avg_quantizer[B_TYPE];
784
785 if(expected_bits < all_available_bits) avg_quantizer[P_TYPE]-= step;
786//printf("%f %lld %f\n", expected_bits, all_available_bits, avg_quantizer[P_TYPE]);
787 }
405469ce 788//printf("qp_i:%f, qp_p:%f, qp_b:%f\n", avg_quantizer[I_TYPE],avg_quantizer[P_TYPE],avg_quantizer[B_TYPE]);
8b4c7dbc
MN
789
790 for(i=0; i<5; i++){
791 available_bits[i]= const_bits[i] + complexity[i]/avg_quantizer[i];
792 }
793//printf("%lld %lld %lld %lld\n", available_bits[I_TYPE], available_bits[P_TYPE], available_bits[B_TYPE], all_available_bits);
3aa102be 794
8e1e6f31
FB
795 qscale= av_malloc(sizeof(double)*rcc->num_entries);
796 blured_qscale= av_malloc(sizeof(double)*rcc->num_entries);
3aa102be 797
8b4c7dbc 798 for(step=256*256; step>0.0000001; step*=0.5){
3aa102be 799 expected_bits=0;
8b4c7dbc 800 rate_factor+= step;
3aa102be
MN
801
802 rcc->buffer_index= s->avctx->rc_buffer_size/2;
803
8b4c7dbc
MN
804 /* find qscale */
805 for(i=0; i<rcc->num_entries; i++){
3aa102be
MN
806 qscale[i]= get_qscale(s, &rcc->entry[i], rate_factor, i);
807 }
808 assert(filter_size%2==1);
809
810 /* fixed I/B QP relative to P mode */
3aa102be 811 for(i=rcc->num_entries-1; i>=0; i--){
8b4c7dbc 812 RateControlEntry *rce= &rcc->entry[i];
6a1f7e7b
MN
813
814 qscale[i]= get_diff_limited_q(s, rce, qscale[i]);
3aa102be 815 }
8b4c7dbc 816
3aa102be
MN
817 /* smooth curve */
818 for(i=0; i<rcc->num_entries; i++){
819 RateControlEntry *rce= &rcc->entry[i];
820 const int pict_type= rce->new_pict_type;
821 int j;
822 double q=0.0, sum=0.0;
823
824 for(j=0; j<filter_size; j++){
825 int index= i+j-filter_size/2;
826 double d= index-i;
d55f7b65 827 double coeff= a->qblur==0 ? 1.0 : exp(-d*d/(a->qblur * a->qblur));
8b4c7dbc 828
3aa102be
MN
829 if(index < 0 || index >= rcc->num_entries) continue;
830 if(pict_type != rcc->entry[index].new_pict_type) continue;
831 q+= qscale[index] * coeff;
832 sum+= coeff;
8b4c7dbc 833 }
3aa102be 834 blured_qscale[i]= q/sum;
8b4c7dbc 835 }
8b4c7dbc
MN
836
837 /* find expected bits */
838 for(i=0; i<rcc->num_entries; i++){
839 RateControlEntry *rce= &rcc->entry[i];
3aa102be
MN
840 double bits;
841 rce->new_qscale= modify_qscale(s, rce, blured_qscale[i], i);
405469ce 842 bits= qp2bits(rce, rce->new_qscale) + rce->mv_bits + rce->misc_bits;
3aa102be 843//printf("%d %f\n", rce->new_bits, blured_qscale[i]);
11dffbe1 844 bits += 8*ff_vbv_update(s, bits);
3aa102be 845
8b4c7dbc 846 rce->expected_bits= expected_bits;
3aa102be 847 expected_bits += bits;
8b4c7dbc
MN
848 }
849
3aa102be 850// printf("%f %d %f\n", expected_bits, (int)all_available_bits, rate_factor);
8b4c7dbc
MN
851 if(expected_bits > all_available_bits) rate_factor-= step;
852 }
8e1e6f31
FB
853 av_free(qscale);
854 av_free(blured_qscale);
8b4c7dbc 855
3aa102be 856 if(abs(expected_bits/all_available_bits - 1.0) > 0.01 ){
9b879566 857 av_log(s->avctx, AV_LOG_ERROR, "Error: 2pass curve failed to converge\n");
3aa102be 858 return -1;
8b4c7dbc 859 }
8b4c7dbc 860
3aa102be 861 return 0;
8b4c7dbc 862}