new PSNR code (now works with chroma, b frames, ...)

rename *_TYPE to FF_*_TYPE for the external API
allow user specified pict_type

Originally committed as revision 1308 to svn://svn.ffmpeg.org/ffmpeg/trunk

6 files changed, 132 insertions, 101 deletions

diff --git a/ffmpeg.c b/ffmpeg.c
index 109bbf95ad..04fc9c99fd 100644
--- a/ffmpeg.c
+++ b/ffmpeg.c
@@ -650,6 +650,11 @@ static void do_video_out(AVFormatContext *s,
     av_free(buf1);
 }
 
+static double psnr(double d){
+    if(d==0) return INFINITY;
+    return -10.0*log(d)/log(10);
+}
+
 static void do_video_stats(AVFormatContext *os, AVOutputStream *ost, 
                            int frame_size)
 {
@@ -682,8 +687,8 @@ static void do_video_stats(AVFormatContext *os, AVOutputStream *ost,
     if (enc->codec_type == CODEC_TYPE_VIDEO) {
         frame_number = ost->frame_number;
         fprintf(fvstats, "frame= %5d q= %2.1f ", frame_number, enc->coded_picture->quality);
-        if (do_psnr)
-            fprintf(fvstats, "PSNR= %6.2f ", enc->psnr_y);
+        if (enc->flags&CODEC_FLAG_PSNR)
+            fprintf(fvstats, "PSNR= %6.2f ", psnr(enc->coded_picture->error[0]/(enc->width*enc->height*255.0*255.0)));
         
         fprintf(fvstats,"f_size= %6d ", frame_size);
         /* compute pts value */
@@ -745,8 +750,8 @@ void print_report(AVFormatContext **output_files,
             frame_number = ost->frame_number;
             sprintf(buf + strlen(buf), "frame=%5d q=%2.1f ",
                     frame_number, enc->coded_picture ? enc->coded_picture->quality : 0);
-            if (do_psnr)
-                sprintf(buf + strlen(buf), "PSNR=%6.2f ", enc->psnr_y);
+            if (enc->flags&CODEC_FLAG_PSNR)
+                sprintf(buf + strlen(buf), "PSNR= %6.2f ", psnr(enc->coded_picture->error[0]/(enc->width*enc->height*255.0*255.0)));
             vid = 1;
         }
         /* compute min output value */
@@ -2161,9 +2166,7 @@ void opt_output_file(const char *filename)
                 }
             
                 if (do_psnr)
-                    video_enc->get_psnr = 1;
-                else
-                    video_enc->get_psnr = 0;
+                    video_enc->flags|= CODEC_FLAG_PSNR;
             
                 video_enc->me_method = me_method;
 
diff --git a/libavcodec/avcodec.h b/libavcodec/avcodec.h
index 287b17c09e..470fbc3d95 100644
--- a/libavcodec/avcodec.h
+++ b/libavcodec/avcodec.h
@@ -5,8 +5,8 @@
 
 #define LIBAVCODEC_VERSION_INT 0x000406
 #define LIBAVCODEC_VERSION     "0.4.6"
-#define LIBAVCODEC_BUILD       4642
-#define LIBAVCODEC_BUILD_STR   "4642"
+#define LIBAVCODEC_BUILD       4643
+#define LIBAVCODEC_BUILD_STR   "4643"
 
 enum CodecID {
     CODEC_ID_NONE, 
@@ -140,6 +140,7 @@ static const int Motion_Est_QTab[] = { ME_ZERO, ME_PHODS, ME_LOG,
 #define CODEC_FLAG_EXTERN_HUFF 0x1000 /* use external huffman table (for mjpeg) */
 #define CODEC_FLAG_GRAY  0x2000   /* only decode/encode grayscale */
 #define CODEC_FLAG_EMU_EDGE 0x4000/* dont draw edges */
+#define CODEC_FLAG_PSNR           0x8000 /* error[?] variables will be set during encoding */
 #define CODEC_FLAG_TRUNCATED  0x00010000 /* input bitstream might be truncated at a random location instead 
                                             of only at frame boundaries */
 #define CODEC_FLAG_NORMALIZE_AQP  0x00020000 /* normalize adaptive quantization */
@@ -252,12 +253,18 @@ static const int Motion_Est_QTab[] = { ME_ZERO, ME_PHODS, ME_LOG,
      * decoding: set by user\
      */\
     void *opaque;\
+\
+    /**\
+     * error\
+     * encoding: set by lavc if flags&CODEC_FLAG_PSNR\
+     * decoding: unused\
+     */\
+    uint64_t error[4];\
 
-/* FIXME: these should have FF_ */
-#define I_TYPE 1 // Intra
-#define P_TYPE 2 // Predicted
-#define B_TYPE 3 // Bi-dir predicted
-#define S_TYPE 4 // S(GMC)-VOP MPEG4
+#define FF_I_TYPE 1 // Intra
+#define FF_P_TYPE 2 // Predicted
+#define FF_B_TYPE 3 // Bi-dir predicted
+#define FF_S_TYPE 4 // S(GMC)-VOP MPEG4
 
 typedef struct AVVideoFrame {
     FF_COMMON_PICTURE
@@ -464,17 +471,6 @@ typedef struct AVCodecContext {
     /* with a Start Code (it should) H.263 does   */
     void (*rtp_callback)(void *data, int size, int packet_number); 
 
-    /**
-     * if you set get_psnr to 1 then after encoding you will have the 
-     * PSNR on psnr_y/cb/cr
-     * encoding: set by user (1-> on, 0-> off)
-     * decoding: unused
-     */
-    int get_psnr;
-    float psnr_y;
-    float psnr_cb;
-    float psnr_cr;
-    
     /* statistics, used for 2-pass encoding */
     int mv_bits;
     int header_bits;
@@ -826,6 +822,13 @@ typedef struct AVCodecContext {
 #define FF_DEBUG_QP        16
 #define FF_DEBUG_MV        32
 #define FF_DEBUG_VIS_MV    64
+    
+    /**
+     * error
+     * encoding: set by lavc if flags&CODEC_FLAG_PSNR
+     * decoding: unused
+     */
+    uint64_t error[4];
 } AVCodecContext;
 
 typedef struct AVCodec {
diff --git a/libavcodec/dsputil.c b/libavcodec/dsputil.c
index 9039d2625f..1e177116a4 100644
--- a/libavcodec/dsputil.c
+++ b/libavcodec/dsputil.c
@@ -1528,37 +1528,3 @@ void avcodec_set_bit_exact(void)
 //    dsputil_set_bit_exact_mmx();
 #endif
 }
-
-void get_psnr(UINT8 *orig_image[3], UINT8 *coded_image[3],
-              int orig_linesize[3], int coded_linesize,
-              AVCodecContext *avctx)
-{
-    int quad, diff, x, y;
-    UINT8 *orig, *coded;
-    UINT32 *sq = squareTbl + 256;
-    
-    quad = 0;
-    diff = 0;
-    
-    /* Luminance */
-    orig = orig_image[0];
-    coded = coded_image[0];
-    
-    for (y=0;y<avctx->height;y++) {
-        for (x=0;x<avctx->width;x++) {
-            diff = *(orig + x) - *(coded + x);
-            quad += sq[diff];
-        }
-        orig += orig_linesize[0];
-        coded += coded_linesize;
-    }
-   
-    avctx->psnr_y = (float) quad / (float) (avctx->width * avctx->height);
-    
-    if (avctx->psnr_y) {
-        avctx->psnr_y = (float) (255 * 255) / avctx->psnr_y;
-        avctx->psnr_y = 10 * (float) log10 (avctx->psnr_y); 
-    } else
-        avctx->psnr_y = 99.99;
-}
-
diff --git a/libavcodec/huffyuv.c b/libavcodec/huffyuv.c
index 0af05d1500..13627b51e7 100644
--- a/libavcodec/huffyuv.c
+++ b/libavcodec/huffyuv.c
@@ -461,7 +461,7 @@ static int encode_init(AVCodecContext *avctx)
     s->version=2;
     
     avctx->coded_picture= &s->picture;
-    s->picture.pict_type= I_TYPE;
+    s->picture.pict_type= FF_I_TYPE;
     s->picture.key_frame= 1;
     
     switch(avctx->pix_fmt){
diff --git a/libavcodec/mpegvideo.c b/libavcodec/mpegvideo.c
index 6b2a376e95..1c574f7768 100644
--- a/libavcodec/mpegvideo.c
+++ b/libavcodec/mpegvideo.c
@@ -985,22 +985,50 @@ static void select_input_picture(MpegEncContext *s){
 
     /* set next picture types & ordering */
     if(s->reordered_input_picture[0]==NULL && s->input_picture[0]){
-        if(/*s->picture_in_gop_number >= s->gop_size ||*/ s->next_picture.data[0]==NULL || s->intra_only){
-            s->reordered_input_picture[0]= s->input_picture[0];
-            s->reordered_input_picture[0]->pict_type= I_TYPE;
-            s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
+        if(s->input_picture[0]->pict_type){
+            /* user selected pict_type */
+            if(s->input_picture[0]->pict_type == I_TYPE){
+                s->reordered_input_picture[0]= s->input_picture[0];
+                s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
+            }else{
+                int b_frames;
+
+                for(b_frames=0; b_frames<s->max_b_frames+1; b_frames++){
+                    if(s->input_picture[b_frames]->pict_type!=B_TYPE) break;
+                }
+                
+                if(b_frames > s->max_b_frames){
+                    fprintf(stderr, "warning, too many bframes in a row\n");
+                    b_frames = s->max_b_frames;
+                    s->input_picture[b_frames]->pict_type= I_TYPE;
+                }
+                
+                s->reordered_input_picture[0]= s->input_picture[b_frames];
+                s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
+                for(i=0; i<b_frames; i++){
+                    coded_pic_num++;
+                    s->reordered_input_picture[i+1]= s->input_picture[i];
+                    s->reordered_input_picture[i+1]->coded_picture_number= coded_pic_num;
+                }    
+            }
         }else{
-            s->reordered_input_picture[0]= s->input_picture[s->max_b_frames];
-            if(s->picture_in_gop_number + s->max_b_frames >= s->gop_size)
+            if(/*s->picture_in_gop_number >= s->gop_size ||*/ s->next_picture.data[0]==NULL || s->intra_only){
+                s->reordered_input_picture[0]= s->input_picture[0];
                 s->reordered_input_picture[0]->pict_type= I_TYPE;
-            else
-                s->reordered_input_picture[0]->pict_type= P_TYPE;
-            s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
-            for(i=0; i<s->max_b_frames; i++){
-                coded_pic_num++;
-                s->reordered_input_picture[i+1]= s->input_picture[i];
-                s->reordered_input_picture[i+1]->pict_type= B_TYPE;
-                s->reordered_input_picture[i+1]->coded_picture_number= coded_pic_num;
+                s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
+            }else{
+                s->reordered_input_picture[0]= s->input_picture[s->max_b_frames];
+                if(s->picture_in_gop_number + s->max_b_frames >= s->gop_size)
+                    s->reordered_input_picture[0]->pict_type= I_TYPE;
+                else
+                    s->reordered_input_picture[0]->pict_type= P_TYPE;
+                s->reordered_input_picture[0]->coded_picture_number= coded_pic_num;
+                for(i=0; i<s->max_b_frames; i++){
+                    coded_pic_num++;
+                    s->reordered_input_picture[i+1]= s->input_picture[i];
+                    s->reordered_input_picture[i+1]->pict_type= B_TYPE;
+                    s->reordered_input_picture[i+1]->coded_picture_number= coded_pic_num;
+                }    
             }
         }
     }
@@ -1027,6 +1055,7 @@ int MPV_encode_picture(AVCodecContext *avctx,
 {
     MpegEncContext *s = avctx->priv_data;
     AVVideoFrame *pic_arg = data;
+    int i;
 
     init_put_bits(&s->pb, buf, buf_size, NULL, NULL);
 
@@ -1076,28 +1105,11 @@ int MPV_encode_picture(AVCodecContext *avctx,
     
     s->total_bits += s->frame_bits;
     avctx->frame_bits  = s->frame_bits;
-//printf("fcode: %d, type: %d, head: %d, mv: %d, misc: %d, frame: %d, itex: %d, ptex: %d\n", 
-//s->f_code, avctx->key_frame, s->header_bits, s->mv_bits, s->misc_bits, s->frame_bits, s->i_tex_bits, s->p_tex_bits);
-#if 0 //dump some stats to stats.txt for testing/debuging
-if(s->max_b_frames==0)
-{
-    static FILE *f=NULL;
-    if(!f) f= fopen("stats.txt", "wb");
-    get_psnr(pict->data, s->current_picture,
-             pict->linesize, s->linesize, avctx);
-    fprintf(f, "%7d, %7d, %2.4f\n", pbBufPtr(&s->pb) - s->pb.buf, s->qscale, avctx->psnr_y);
-}
-#endif
-#if 0
-    if (avctx->get_psnr) {
-        /* At this point pict->data should have the original frame   */
-        /* an s->current_picture should have the coded/decoded frame */
-        get_psnr(pict->data, s->current_picture.data,
-                 pict->linesize, s->linesize, avctx);
-//        printf("%f\n", avctx->psnr_y);
-    }
-#endif
 
+    for(i=0; i<4; i++){
+        avctx->error[i] += s->current_picture.error[i];
+    }
+    
     return pbBufPtr(&s->pb) - s->pb.buf;
 }
 
@@ -1819,7 +1831,7 @@ void MPV_decode_mb(MpegEncContext *s, DCTELEM block[6][64])
         }
     }
     
-    if (!(s->encoding && (s->intra_only || s->pict_type==B_TYPE))) {
+    if ((s->flags&CODEC_FLAG_PSNR) || !(s->encoding && (s->intra_only || s->pict_type==B_TYPE))) { //FIXME precalc
         UINT8 *dest_y, *dest_cb, *dest_cr;
         int dct_linesize, dct_offset;
         op_pixels_func (*op_pix)[4];
@@ -2544,6 +2556,22 @@ static inline void encode_mb_hq(MpegEncContext *s, MpegEncContext *backup, MpegE
         copy_context_after_encode(best, s, type);
     }
 }
+                
+static inline int sse(MpegEncContext *s, uint8_t *src1, uint8_t *src2, int w, int h, int stride){
+    uint32_t *sq = squareTbl + 256;
+    int acc=0;
+    int x,y;
+    
+    if(w==16 && h==16) 
+        return s->dsp.pix_norm(src1, src2, stride);
+    
+    for(y=0; y<h; y++){
+        for(x=0; x<w; x++){
+            acc+= sq[src1[x + y*stride] - src2[x + y*stride]];
+        } 
+    }
+    return acc;
+}
 
 static void encode_picture(MpegEncContext *s, int picture_number)
 {
@@ -2723,11 +2751,13 @@ static void encode_picture(MpegEncContext *s, int picture_number)
     s->b_count=0;
     s->skip_count=0;
 
-    /* init last dc values */
-    /* note: quant matrix value (8) is implied here */
-    s->last_dc[0] = 128;
-    s->last_dc[1] = 128;
-    s->last_dc[2] = 128;
+    for(i=0; i<3; i++){
+        /* init last dc values */
+        /* note: quant matrix value (8) is implied here */
+        s->last_dc[i] = 128;
+        
+        s->current_picture.error[i] = 0;
+    }
     s->mb_incr = 1;
     s->last_mv[0][0][0] = 0;
     s->last_mv[0][0][1] = 0;
@@ -2992,6 +3022,30 @@ static void encode_picture(MpegEncContext *s, int picture_number)
             }
 
             MPV_decode_mb(s, s->block);
+            
+            if(s->flags&CODEC_FLAG_PSNR){
+                int w= 16;
+                int h= 16;
+
+                if(s->mb_x*16 + 16 > s->width ) w= s->width - s->mb_x*16;
+                if(s->mb_y*16 + 16 > s->height) h= s->height- s->mb_y*16;
+                
+                s->current_picture.error[0] += sse(
+                    s,
+                    s->new_picture    .data[0] + s->mb_x*16 + s->mb_y*s->linesize*16,
+                    s->current_picture.data[0] + s->mb_x*16 + s->mb_y*s->linesize*16,
+                    w, h, s->linesize);
+                s->current_picture.error[1] += sse(
+                    s,
+                    s->new_picture    .data[1] + s->mb_x*8  + s->mb_y*s->uvlinesize*8,
+                    s->current_picture.data[1] + s->mb_x*8  + s->mb_y*s->uvlinesize*8,
+                    w>>1, h>>1, s->uvlinesize);
+                s->current_picture.error[2] += sse(
+                    s,
+                    s->new_picture    .data[2] + s->mb_x*8  + s->mb_y*s->uvlinesize*8,
+                    s->current_picture.data[2] + s->mb_x*8  + s->mb_y*s->uvlinesize*8,
+                    w>>1, h>>1, s->uvlinesize);
+            }
 //printf("MB %d %d bits\n", s->mb_x+s->mb_y*s->mb_width, get_bit_count(&s->pb));
         }
 
diff --git a/libavcodec/mpegvideo.h b/libavcodec/mpegvideo.h
index 7a0682fc5a..46e19d4b17 100644
--- a/libavcodec/mpegvideo.h
+++ b/libavcodec/mpegvideo.h
@@ -48,6 +48,11 @@ enum OutputFormat {
 #define MAX_RUN    64
 #define MAX_LEVEL  64
 
+#define I_TYPE FF_I_TYPE  // Intra
+#define P_TYPE FF_P_TYPE  // Predicted
+#define B_TYPE FF_B_TYPE  // Bi-dir predicted
+#define S_TYPE FF_S_TYPE  // S(GMC)-VOP MPEG4
+
 typedef struct Predictor{
     double coeff;
     double count;