Evaluating PSNR metric of two YUV420 sequences in pure python


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



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.