计算两段yuv格式视频流中每一帧的psnr值

现在虚拟视点生成了yuv格式的视频流，如何计算每一帧的psnr值，进行了如下步骤：

1.对yuv视频流进行转换为jpg图片；

2.把jpg图片转化为avi视频流，便于批处理；

3.对avi视频中的每一帧计算psnr值。

程序：

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cv.h>
#include <highgui.h>
#define NUM_FRAME 100 //只处理前100帧，根据视频帧数可修改
void CalcPsnr(const char *in1,const char *in2)
{
	cv::VideoCapture vc1;
	cv::VideoCapture vc2;
	bool flag1 = vc1.open(in1);
	bool flag2 = vc2.open(in2);
	if (!flag1||!flag2)
	{
		printf("avi file open error \n");
		system("pause");
		exit(-1);
	}

	int frmCount1 = vc1.get(CV_CAP_PROP_FRAME_COUNT);
	int frmCount2 = vc2.get(CV_CAP_PROP_FRAME_COUNT);
	printf("frmCount: %d \n", frmCount1);
	printf("frmCount: %d \n", frmCount2);
	for (int i = 0; i < frmCount1; i++)
	{
		printf("%d/%d \n", i + 1, frmCount1);
		cv::Mat image_ref;
		vc1 >> image_ref;
		cv::Mat image_obj;
		vc2 >> image_obj;
		double mse = 0;
		double div_r = 0;
		double div_g = 0;
		double div_b = 0;
		int width = image_ref.cols;
		int height = image_ref.rows;
		double psnr = 0;
		for (int v = 0; v < height; v++)
		{
			for (int u = 0; u < width; u++)
			{
				div_r = image_ref.at<cv::Vec3b>(v, u)[0] - image_obj.at<cv::Vec3b>(v, u)[0];
				div_g = image_ref.at<cv::Vec3b>(v, u)[1] - image_obj.at<cv::Vec3b>(v, u)[1];
				div_b = image_ref.at<cv::Vec3b>(v, u)[2] - image_obj.at<cv::Vec3b>(v, u)[2];
				mse += ((div_r*div_r + div_b*div_b + div_g*div_g) / 3);

			}
		}
		mse = mse / (width*height);
		psnr = 10 * log10(255 * 255 / mse);
		printf("%lf\n", mse);
		printf("%lf\n", psnr);
	}
	return;
}
void DisplayYUV2RGB(const char *dir,const char *in,int _w,int _h)
{
	int w = _w;
	int h = _h;
	printf("yuv file w: %d, h: %d \n", w, h);
	FILE* pFileIn = fopen(in, "rb+");
	int bufLen = w*h * 3 / 2;
	unsigned char* pYuvBuf = new unsigned char[bufLen];
	int iCount = 0;


	for (int i = 0; i<NUM_FRAME; i++)
	{
		fread(pYuvBuf, bufLen*sizeof(unsigned char), 1, pFileIn);

		cv::Mat yuvImg;
		yuvImg.create(h * 3 / 2, w, CV_8UC1);
		memcpy(yuvImg.data, pYuvBuf, bufLen*sizeof(unsigned char));
		cv::Mat rgbImg;
		cv::cvtColor(yuvImg, rgbImg, CV_YUV2BGR_I420);

		cv::imshow("img", rgbImg);
		char s[100];
		sprintf(s,"%spic%d%s",dir,i,".jpg");
		cv::imwrite(s, rgbImg);
		cv::waitKey(1);

		printf("%d \n", iCount++);
	}
	delete[] pYuvBuf;
	fclose(pFileIn);
}
void Image_to_video(const char* in, const char* out)
{
	int num = 1;
	CvSize size = cvSize(1024, 768);  //视频帧格式的大小
	double fps = 30; //每秒钟的帧率
	CvVideoWriter *writer = cvCreateVideoWriter(out, CV_FOURCC('D', 'I', 'V', 'X'), fps, size); //创建视频文件
	char cname[100];
	sprintf(cname, in, num); //加载图片的文件夹，图片的名称编号是1开始1，2,3,4,5.。。。
	IplImage *src = cvLoadImage(cname);
	if (!src)
	{
		return;
	}
	IplImage *src_resize = cvCreateImage(size, 8, 3); //创建视频文件格式大小的图片
	cvNamedWindow("avi");
	while (src)
	{
		cvShowImage("avi", src_resize);
		cvWaitKey(1);
		cvResize(src, src_resize); //将读取的图片设置为视频格式大小相同
		cvWriteFrame(writer, src_resize); //保存图片为视频流格式
		cvReleaseImage(&src); //释放空间
		num++;
		sprintf(cname, in, num);
		src = cvLoadImage(cname);       //循环读取数据
	}
	cvReleaseVideoWriter(&writer);
	cvReleaseImage(&src_resize);
}
int main(int argc, char *argv[])
{
	const char *out = "C:/Users/jiang/Desktop/output/book_virtual08.yuv";
	const char *dir = "C:/Users/jiang/Desktop/output/tupian1/";
	DisplayYUV2RGB(dir, out, 1024, 768);
	const char *outImagename = "C:/Users/jiang/Desktop/output/tupian1/pic%d.jpg";
	const char *outVideoname = "C:/Users/jiang/Desktop/output/3outfile.avi";
	Image_to_video(outImagename, outVideoname);

	out = "C:/Users/jiang/Desktop/bookarrival/bookarrival_c_8.yuv";
	dir = "C:/Users/jiang/Desktop/output/tupian1/";
	DisplayYUV2RGB(dir, out, 1024, 768);
	outImagename = "C:/Users/jiang/Desktop/output/tupian1/pic%d.jpg";
	outVideoname = "C:/Users/jiang/Desktop/output/4outfile.avi";
	Image_to_video(outImagename, outVideoname);

	const char *in1 = "C:/Users/jiang/Desktop/output/3outfile.avi";
	const char *in2 = "C:/Users/jiang/Desktop/output/4outfile.avi";
	CalcPsnr(in1, in2);

	getchar();
}