Video Evaluation by Python

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