Advanced Search

Evaluating PSNR metric of two YUV420 sequences in pure python

/ Published in: Python
Save to your folder(s)

Expand | Embed | Plain Text

Copy this code and paste it in your HTML 
  1. import array
  2. import sys
  3. import os
  4. import math
  5. import re
  6.  
  7. if __name__ == '__main__':
  8.     if len(sys.argv) < 3:
  9.         print "usage: %s filename1.yuv filename2.yuv [width height]" % sys.argv[0]
  10.  
  11.  
  12.     filename1 = sys.argv[1]
  13.     filename2 = sys.argv[2]
  14.  
  15.     if filename1 == filename2:
  16.         print "warning! do you really mean to compare the file with itself?"
  17.  
  18.     data1 = array.array('B')
  19.     data2 = array.array('B')
  20.  
  21.     file1_size = os.path.getsize(filename1)
  22.     file2_size = os.path.getsize(filename2)
  23.  
  24.     minsize = min(file1_size, file2_size)
  25.  
  26.     if file1_size != file2_size:
  27.         print "warning, file sizes do not match! comparing min size %d bytes" % minsize
  28.  
  29.     if len(sys.argv) >= 5:
  30.         w = int(sys.argv[3])
  31.         h = int(sys.argv[4])
  32.     else:
  33.         m = re.search(".*[_-](\d+)x(\d+).*", filename1)
  34.         assert m
  35.         w = int(m.group(1))
  36.         h = int(m.group(2))
  37.         assert w*h*3/2 <= minsize
  38.         print "(%dx%d)" %(w,h)
  39.  
  40.  
  41.     data1.fromfile(open(filename1,"rb"),minsize)
  42.     data2.fromfile(open(filename2,"rb"),minsize)
  43.  
  44.     print "data size: %d, w*h=%d, w*h*3/2=%d" % (len(data1), w*h, w*h*3/2)
  45.     print "evaluating mse.."
  46.  
  47.     def psnr(mse):
  48.         log10 = math.log10
  49.         return 10.0*log10(float(256*256)/float(mse))
  50.  
  51.     def mean(seq):
  52.         if len(seq) == 0: return 0.0
  53.         else: return sum(seq)/float(len(seq))
  54.  
  55.     def sum_square_err(data1, data2, beg, end):
  56.         return sum( (a-b)*(a-b) for a,b in zip(data1[beg:end],data2[beg:end]))
  57.  
  58.     y = 0
  59.     u = y + (w * h)
  60.     v = u + (w/2 * h/2)
  61.     data_end = w*h*3/2
  62.  
  63.     begin = [y,u,v,y]
  64.     end = [u,v,data_end,data_end]
  65.     size = [w*h, w*h/4, w*h/4, w*h*3/2]
  66.  
  67.     colorspace_mse = [sum_square_err(data1,data2,
  68.         begin[i], end[i])/float(size[i]) for i in range(4)]
  69.  
  70.     colorspace_psnr = [psnr(m) for m in colorspace_mse]
  71.  
  72.     print "planes: Y, U, V, Whole frame"
  73.     print 'colorplane mse: ', colorspace_mse
  74.     print 'colorplane psnr: ', colorspace_psnr

Report this snippet


Comments

RSS Icon Subscribe to comments

You need to login to post a comment.