Posted By

runnig on 04/08/11


Tagged

python frame psnr metric yuv differences


Versions (?)

Evaluating PSNR metric of two YUV420 sequences in pure python


 / Published in: Python
 

  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  

You need to login to post a comment.