YUV视频PSNR计算及代码实现

最近在做视频质量评估相关的东西，其中涉及了一个离线的PSNR计算，练个手写了个YUV视频的PSNR计算程序。

对于图像PSNR计算公式为：

PSNR=10log10(255)2MN∑M−1x=0∑N−1y=0|f(x,y)−f′(x,y)|2

其中f(x,y)和f’(x,y)为编码前后图像(x,y)处的像素值，图像分辨率为MxN。

对于视频，就是将每一帧的失真累加，然后在计算PSNR：

PSNR=10log10(255)2MN∗F∑F−1f=0∑M−1x=0∑N−1y=0|f(x,y)−f′(x,y)|2

其中视频共有F帧。

下面来看代码实现，其中需要注意溢出问题。视频的分辨率可能很大，导致计算得到的SSD很大，在后边的log计算中就会出现问题，因此先对SSD除以MN（视频分辨率），可以放置溢出。

main函数main.cpp：

#include "PSNR.h"

#define _CRT_SECURE_NO_WARNINGS


int main(int argc, char *argv[])
{
    if( argc != 5 ) 
    {
        printf ("The format of input parameter is error!\n");
        exit(-1);
    }

    FILE * org_yuv;    
    FILE * out_yuv;     

    if(NULL == (org_yuv =  fopen(argv[1], "rb")))
    {
        printf("File input is can't open!\n");
        return -1;
    }
    if(NULL == (out_yuv = fopen(argv[2], "rb")))
    {
        printf("File output is can't open!\n");
        return -1;
    }
    int IMAGEWIDTH = 1024;
    int IMAGEHEIGHT = 768;
    IMAGEWIDTH = atoi(argv[3]);
    IMAGEHEIGHT = atoi(argv[4]);


    int readsize;
    unsigned char *org_buff;
    org_buff =  (unsigned char *)malloc(IMAGEWIDTH*IMAGEHEIGHT*sizeof(unsigned char)*3/2); 
    memset (org_buff, 0, IMAGEWIDTH*IMAGEHEIGHT*sizeof(unsigned char)*3/2);

    unsigned char *out_buff;
    out_buff=(unsigned char *)malloc(IMAGEWIDTH*IMAGEHEIGHT*sizeof(unsigned char)*3/2);
    memset (out_buff, 0, IMAGEWIDTH*IMAGEHEIGHT*sizeof(unsigned char)*3/2);


    int Frame =24;
    int i = Frame;
    double videoSSD[3] = {0};
    double videoPSNR[3] = {0};

    while(i)
    {
        readsize=fread(org_buff, 1 , IMAGEWIDTH*IMAGEHEIGHT*3/2, org_yuv);
        readsize=fread(out_buff, 1 , IMAGEWIDTH*IMAGEHEIGHT*3/2, out_yuv);
        if(readsize<IMAGEWIDTH*IMAGEHEIGHT*3/2)
            break;

        printf("Frame %d\n",Frame-i);

        calPSNR(org_buff, out_buff, IMAGEWIDTH,IMAGEHEIGHT,videoSSD);

        i--;
    }

    for(int ch=0;ch<3;ch++)
    {
        videoPSNR[ch] = ( videoSSD[ch] ? 10.0 * log10((255 * 255) *Frame/ videoSSD[ch] ) : 999.99 );
    }
    printf("\n\nSUMMARY --------------------------------------------------------\n");
    printf("Y-PSNR : %.4f, U-PSNR : %.4f, V-PSNR : %.4f\n",videoPSNR[0],videoPSNR[1],videoPSNR[2]);

    free(org_buff);
    org_buff=NULL;
    free(out_buff);
    out_buff=NULL;

    fclose(org_yuv);
    fclose(out_yuv);

}

PSNR计算函数calPSNR.cpp：

#include "PSNR.h"
#include<malloc.h>
#include<memory.h>

void calPSNR(unsigned char *yuv_buff,unsigned char *yuv2_buff, int PIC_W, int PIC_H, double *SSD_SUM)
{
    unsigned char *orgY_buff, *outY_buff;
    orgY_buff = yuv_buff;
    outY_buff = yuv2_buff;
    double *mSSD_SUM = SSD_SUM; 
    double SSD[3] = {0};
    double PSNR[3] = {0};
    int h,v;

    for(int ch = 0 ;ch < 3; ch++)
    {
        int shift = 0;
        if(ch)
            shift = 1;
        int height = PIC_H>>shift;
        int width = PIC_W>>shift;

        for(v=0; v<height; v++)
        {
            for(h=0; h<width; h++)
            {
                int iDiff = orgY_buff[h] - outY_buff[h];
                if(iDiff>255)
                    iDiff = 255;
                SSD[ch] += iDiff * iDiff;
            }

            orgY_buff += width;
            outY_buff += width;
        }

        SSD[ch] /=(width*height);    //防止溢出
        mSSD_SUM[ch] += SSD[ch];
        PSNR[ch] = ( SSD[ch] ? 10.0 * log10((255 * 255)/ SSD[ch] ) : 999.99 );
    }

    printf("Y-PSNR : %f, U-PSNR : %f, V-PSNR : %f\n",PSNR[0],PSNR[1],PSNR[2]);
}

头文件PSNR.h：

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

void calPSNR(unsigned char *rgb_buff, unsigned char *yuv_buff, int ImageWidth, int ImageHeight, double *SSD_SUM);