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