Video Evaluation by Python

最新推荐文章于 2024-03-13 17:39:27 发布

Snail_Walker

最新推荐文章于 2024-03-13 17:39:27 发布

阅读量1.1k

点赞数 1

本文链接：https://blog.csdn.net/c602273091/article/details/49888231

版权

Python && SCM 同时被 3 个专栏收录

38 篇文章 0 订阅

订阅专栏

__本科实习__

34 篇文章 0 订阅

订阅专栏

HEVC

29 篇文章 0 订阅

订阅专栏

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: