Here is the code to calculate for PSNR and SSIM of YUV.
My code has its advantage that it can process the problem by batch processing.
If U have any problem, U can contact me without hesitation.
#-*-coding:utf-8-*-
'''
Create on 2015/11/16 9:05:59
@author: Chen Yu in RockChip Internship
'''
import os
import sys
import re
import string
import math
import win32api
import win32process
import win32event
from subprocess import Popen,PIPE
# processing folder path
dirGood = r"E:\ErrorConcealment\Sequence\Sequence_Bat_Good"
dirCopy = r"E:\ErrorConcealment\Sequence\Sequence_Bat"
dirFirstBatch = r"E:\ErrorConcealment\Sequence\Sequence_Bat\First Batch"
dirSecondBatch = r"E:\ErrorConcealment\Sequence\Sequence_Bat\Second Batch"
def MyGetFileList(dir, filelist):
if os.path.isdir(dir):
for s in os.listdir(dir):
if re.match(r"([\s\S]*)\.yuv", s):
filelist.append(s)
return filelist
def CalculateForTwoFolder(fileList1, fileList2, dir1, dir2, directive):
for f1 in fileList1:
for f2 in fileList2:
if f1 == f2:
if directive == 1:
CalculateSSIMForTwoYUVs(f1, f2, dir1, dir2)
elif directive == 0:
CalculatePSNRForTwoYUVs(f1, f2, dir1, dir2)
else:
print "---------------------------Sorry, Ur directive is wrong------------------------------\n"
print "---------------------------0: Count PSNR for two YUV files---------------------------\n"
print "---------------------------1: Count SSIM for two YUV files---------------------------\n"
def CalculateSSIMForTwoYUVs(f1, f2, dir1, dir2):
d1 = dir1 + '\\' + f1
d2 = dir2 + '\\' + f2
result = re.findall("([\s\S]*)\_(\d+)x(\d+)\_([\s\S]*)",d1)
storagePath = re.findall("([\s\S]*)\.yuv", d2)
fpPath = storagePath[0] + "SSIM.txt"
ssim = 0
ssimList = []
AvaregeSsim = 0
AverageSsimNum = 0
frameNum = 0
for line in result:
width = string.atoi(line[1])
height = string.atoi(line[2])
chunkSize = (width * height * 3) >> 1
Area = width * height
frame0 = [range(height) for i in range(width)]
frame1 = [range(height) for i in range(width)]
if (os.path.isfile(d1) and os.path.isfile(d2)):
file1Object = open(d1, 'rb')
file2Object = open(d2, 'rb')
fP = open(fpPath, "w")
fP.write("%s\n" % d1)
fP.write("%s\n" % d2)
print d1
print d2
SumOfDev = 0
stdDev = 0
while True:
chunk0 = file1Object.read(chunkSize)
chunk1 = file2Object.read(chunkSize)
if not chunk0:
break
if not chunk1:
break
l0 = list(chunk0)
l1 = list(chunk1)
for i in range (Area):
if i:
ck0 = ord(l0[i])
ck1 = ord(l1[i])
frame0[i % width][i / width] = ck0
frame1[i % width][i / width] = ck1
ssim = Ssim(frame0, frame1, width, height)
if ssim < 1:
AvaregeSsim += ssim
ssimList.append(ssim)
AverageSsimNum += 1
fP.write("------------------------The %-3d------------------- frame's SSIM is | %f \n" % (frameNum, ssim))
print ("------------------------The %-3d------------------- frame's SSIM is | %f \n" % (frameNum, ssim))
frameNum += 1
if AverageSsimNum != 0:
AvaregeSsim /= AverageSsimNum
for itemSSIM in ssimList:
SumOfDev += ((itemSSIM - AvaregeSsim) ** 2)
stdDev = math.sqrt(SumOfDev*1.0/AverageSsimNum)
fP.write("\n------------------------The Average SSIM for %-3d frames is %f------------------- \n" % (AverageSsimNum, AvaregeSsim))
print ("\n------------------------The Average SSIM for %-3d frames is %f------------------- \n" % (AverageSsimNum, AvaregeSsim))
fP.write("\n------------------------The Standard Deviation is %f------------------- \n\n\n" % (stdDev))
print ("\n------------------------The Standard Deviation is %f------------------- \n\n\n" % (stdDev))
else:
fP.write("\n---------------------------------------There is NO error in this YUV-------------------------\n\n\n")
print ("\n---------------------------------------There is NO error in this YUV-------------------------\n\n\n")
file1Object.close()
file2Object.close()
fP.close()
def CalculatePSNRForTwoYUVs(f1, f2, dir1, dir2):
d1 = dir1 + '\\' + f1
d2 = dir2 + '\\' + f2
result = re.findall("([\s\S]*)\_(\d+)x(\d+)\_([\s\S]*)",d1)
storagePath = re.findall("([\s\S]*)\.yuv", d2)
fpPath = storagePath[0] + ".txt"
for line in result:
width = string.atoi(line[1])
height = string.atoi(line[2])
Area = width * height
chunkSize = (Area * 3) >> 1 # 420
frameNum = 0
MSE = 0
MAXDivMSE = 0
PSNR = 0
AveragePSNR = 0
CountAveragePSNRFrame = 0
ListPSNR = []
if (os.path.isfile(d1) and os.path.isfile(d2)):
file1Object = open(d1, 'rb')
file2Object = open(d2, 'rb')
fP = open(fpPath, "w")
fP.write("%s\n" % d1)
fP.write("%s\n" % d2)
while True:
chunk1 = file1Object.read(chunkSize)
chunk2 = file2Object.read(chunkSize)
if not chunk1:
break
if not chunk2:
break
sumOfMSE = 0
l1 = list(chunk1)
l2 = list(chunk2)
for i in range (Area):
ck1 = ord(l1[i])
ck2 = ord(l2[i])
sumOfMSE += (abs(ck1 - ck2)**2)
MSE = math.sqrt(sumOfMSE*1.0/Area)
if MSE != 0:
MAXDivMSE = 255 / MSE
else:
MAXDivMSE = 0
if MAXDivMSE > 0:
PSNR = 20 * math.log10(MAXDivMSE)
else:
PSNR = 0
print ("------------------------The %-3d------------------- frame's PSNR is | %f \n" % (frameNum, PSNR))
fP.write("----------------------The %-3d---------------------- frame's PSNR is | %f \n" % (frameNum, PSNR))
frameNum += 1
if PSNR != 0:
CountAveragePSNRFrame += 1
AveragePSNR += PSNR
ListPSNR.append(PSNR)
AveragePSNR = 0 if CountAveragePSNRFrame == 0 else (AveragePSNR*1.0/CountAveragePSNRFrame)
SumOfDev = 0
stdDev = 0
if CountAveragePSNRFrame != 0:
for itemPSNR in ListPSNR:
SumOfDev += ((itemPSNR - AveragePSNR) ** 2)
stdDev = math.sqrt(SumOfDev*1.0/CountAveragePSNRFrame)
print ("------------------------The Average PSNR is %f-------------------\n" % (AveragePSNR))
print("----------------------The Standard Deviation is %f----------------------\n" % (stdDev))
fP.write ("------------------------The Average PSNR is %f-------------------\n" % (AveragePSNR))
fP.write("----------------------The Standard Deviation is %f----------------------\n" % (stdDev))
file1Object.close()
file2Object.close()
fP.close()
print "--------------------------------------------------------------------\n\n"
def BatchCalculate(directive):
FileListOfGood = MyGetFileList(dirGood, [])
FileListOfCopy = MyGetFileList(dirCopy, [])
FileListOfFirstBatch = MyGetFileList(dirFirstBatch, [])
FileListOfSecondBatch = MyGetFileList(dirSecondBatch, [])
CalculateForTwoFolder(FileListOfGood, FileListOfCopy, dirGood, dirCopy, directive)
CalculateForTwoFolder(FileListOfGood, FileListOfFirstBatch, dirGood, dirFirstBatch, directive)
CalculateForTwoFolder(FileListOfGood, FileListOfSecondBatch, dirGood, dirSecondBatch, directive)
def SsimEnd1(s1, s2, ss, s12):
ssim_c1 = (int)(.01*.01*255*255*64 + .5)
ssim_c2 = (int)(.03*.03*255*255*64*63 + .5)
vars = ss*64 - s1*s1 - s2*s2
covar = s12*64 - s1*s2
return (float)(2*s1*s2 + ssim_c1) * (float)(2*covar + ssim_c2)\
/ ((float)(s1*s1 + s2*s2 + ssim_c1) * (float)(vars + ssim_c2));
# Here I make the decision that for a list list[x][y], x stands column(like a picture)
def SsimEnd4(l1, l2, width):
ssim = 0
for i in range(width):
ssim += SsimEnd1(l1[i][0] + l1[i + 1][0] + l2[i][0] + l2[i + 1][0],
l1[i][1] + l1[i + 1][1] + l2[i][1] + l2[i + 1][1],
l1[i][2] + l1[i + 1][2] + l2[i][2] + l2[i + 1][2],
l1[i][3] + l1[i + 1][3] + l2[i][3] + l2[i + 1][3])
return ssim
# Size is 32 for pix1 and pix2, and sums is [2][4]
def ssim_4x4x2_core( pix1, pix2, sums):
for z in range(2):
s1 = 0
s2 = 0
ss = 0
s12 = 0
for x in range(16):
a = pix1[x + (z << 4)]
b = pix2[x + (z << 4)]
s1 += a
s2 += b
ss += a*a
ss += b*b
s12 += a*b
sums[z][0] = s1
sums[z][1] = s2
sums[z][2] = ss
sums[z][3] = s12
# Calculate Two Frame's SSIM
def Ssim(pixel0, pixel1, width, height):
Sums0 = [range(4) for i in range((width >> 2) + 3)] # width x 4( column x row = width x height )
Sums1 = [range(4) for i in range((width >> 2) + 3)]
width = (width >> 2)
height = (height >> 2)
z = 0
ssim = 0.0
pix0Ck = range(32)
pix1Ck = range(32)
sums0 = [range(4) for i in range(2)]
sums1 = [range(4) for i in range(2)]
for y in range(1,height):
while(z <= y):
for x in range(0, width, 2):
for i in range(4):
for j in range(8):
if j >= 4:
pix0Ck[(i << 2) + j + 12] = pixel0[(x << 2) + j][(z << 2) + i]
pix1Ck[(i << 2) + j + 12] = pixel1[(x << 2) + j][(z << 2) + i]
else:
pix0Ck[(i << 2) + j] = pixel0[(x << 2) + j][(z << 2) + i]
pix1Ck[(i << 2) + j] = pixel1[(x << 2) + j][(z << 2) + i]
if z % 2 == 0:
ssim_4x4x2_core(pix0Ck, pix1Ck, sums0)
Sums0[x] = sums0[0]
Sums0[x + 1] = sums0[1]
else:
ssim_4x4x2_core(pix0Ck, pix1Ck, sums1)
Sums1[x] = sums1[0]
Sums1[x + 1] = sums1[1]
z += 1
for x in range(0, width - 1, 4):
ssim += SsimEnd4(Sums0[x : x + 5], Sums1[x : x + 5], min(4, width - 1 - x))
return 0 if (width - 1 == 0) or (height - 1 == 0) else ssim * 1.0/((width - 1.0)*(height - 1.0))
if __name__ == "__main__":
directive = raw_input("Please input Ur directive: ( 0 for PSNR and 1 for SSIM ) ")
BatchCalculate(directive)
Some tips:
To deal with chunk data.
1、thelist = list(thestring)
2、for character in thestring:
do_something_with(c)
3、results = [do_something_with(c) for c in thestring]
4、results = map(do_something, thestring)
5、import sets
magic_chars = sets.Set(‘abracadabra’)
poppins_chars = sets.Set(’supercalifragilisticexpialidocious’)
print ”.join(magic_chars & poppins_chars) # set intersection
6、 Sums1 = [range(4) for i in range((width >> 2) + 3)]
Above method maybe the best.
7、Passing parameter：
r‘“C:\Users\Administrator\Documents\Visual Studio 2008\Projects\videoDetection\Debug\videoDetection.exe” -i %s -w 1 -l 20 -t 30’ % L0
Remember %, this can be used in passing parameter.
8、What’s the difference between pointer array and array pointer:
http://www.cnblogs.com/Romi/archive/2012/01/10/2317898.html
9、Some basic file operation:
http://www.cnblogs.com/allenblogs/archive/2010/09/13/1824842.html
http://www.cnblogs.com/rollenholt/archive/2012/04/23/2466179.html
