A script evaluating PSNR metric of two YUV420 frames usage: psnr.py filename1.yuv filename2.yuv frame_width frame_height
Alternatively if filename1 contains width and height in the form file-1200x1600.yuv, the script will extract width and height of a frame from the filename. Then usage even simplier: psnr.py filename1-1200x1600.yuv filename2.yuv
Python, 111 lines
Download
Copy to clipboard1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # Denis Gorodetskiy
# linkedin.com/in/gorodetskiy
import array
import sys
import os
import math
import re
def frame_diff(frame1_data, frame2_data, begin_pos, end_pos, w, h):
luma_size = w * h
chroma_size = w * h/4
frame_size = w * h * 3/2
print "data size: %d, w*h=%d, w*h*3/2=%d" % (end_pos-begin_pos,
luma_size, frame_size)
print "evaluating mse.."
def psnr(mse):
log10 = math.log10
return 10.0*log10(float(256*256)/float(mse))
def mean(seq):
if len(seq) == 0: return 0.0
else: return sum(seq)/float(len(seq))
def sum_square_err(data1, data2, beg, end):
return sum( (a-b)*(a-b) for a,b in zip(data1[beg:end],data2[beg:end]))
y = begin_pos
u = y + luma_size
v = u + chroma_size
begin = [y,u,v,y]
end = [u,v,end_pos,end_pos]
size = [luma_size, chroma_size, chroma_size, frame_size]
colorspace_mse = [sum_square_err(frame1_data,frame2_data,
begin[i], end[i])/float(size[i]) for i in range(4)]
colorspace_psnr = [psnr(m) for m in colorspace_mse]
return colorspace_mse, colorspace_psnr, colorspace_psnr[-1]
def width_height_from_str(s):
m = re.search(".*[_-](\d+)x(\d+).*", s)
if not m:
raise RuntimeError()
w = int(m.group(1))
h = int(m.group(2))
return w,h
def usage(me):
print "usage: %s filename1.yuv filename2.yuv [width height]" % me
print "\tif you don't want to specify width height explicitly,"
print "\tscript will try to extract width, height from filenames,"
print "\tfilename_1024x768.yuv or filename-1024x768.yuv yield (width,height)=(1024,768)"
def main(argv):
if len(argv) < 3:
usage(argv[0])
filename1 = argv[1]
filename2 = argv[2]
if filename1 == filename2:
print "warning! do you really mean to compare the file with itself?"
data1 = array.array('B')
data2 = array.array('B')
file1_size = os.path.getsize(filename1)
file2_size = os.path.getsize(filename2)
minsize = min(file1_size, file2_size)
if file1_size != file2_size:
print "warning, file sizes do not match! comparing min size %d bytes" % minsize
if len(argv) >= 5:
w = int(argv[3])
h = int(argv[4])
else:
try:
w,h = width_height_from_str(filename1)
except RuntimeError:
try:
w,h = width_height_from_str(filename2)
except RuntimeError:
print "failed to parse width,height from filename"
usage(argv[0])
return
assert w*h*3/2 <= minsize
data_end = w*h * 3 /2
data1.fromfile(open(filename1,"rb"),minsize)
data2.fromfile(open(filename2,"rb"),minsize)
colorplane_mse, colorplane_psnr, frame_psnr = frame_diff(data1, data2,
0, data_end, w, h)
print "planes: Y, U, V, Whole frame"
print 'colorplane mse: ', colorplane_mse
print 'colorplane psnr: ', colorplane_psnr
print 'frame psnr: ', frame_psnr
if __name__ == '__main__':
main(sys.argv)
|
