Add ssim, psnr

measure.py

@@ -0,0 +1,78 @@
+"""
+Video Quality Metrics
+Copyright (c) 2014 Alex Izvorski <[email protected]>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import numpy
+import re
+import sys
+import scipy.misc
+
+import vifp
+import ssim
+#import ssim_theano
+import psnr
+
+def img_greyscale(img):
+    return 0.299 * img[:,:,0] + 0.587 * img[:,:,1] + 0.114 * img[:,:,2]
+
+def img_read_yuv(src_file, width, height):
+    y_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=(width * height)).reshape( (height, width) )
+    u_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=((width/2) * (height/2))).reshape( (height/2, width/2) )
+    v_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=((width/2) * (height/2))).reshape( (height/2, width/2) )
+    return (y_img, u_img, v_img)
+
+ref_file = sys.argv[1]
+dist_file = sys.argv[2]
+
+if ".yuv" in ref_file:
+    # Inputs are uncompressed video in YUV420 planar format
+
+    # Get resolution from file name
+    m = re.search(r"(\d+)x(\d+)", ref_file)
+    if not m:
+        print "Could not find resolution in file name: %s" % (ref_file)
+        exit(1)
+
+    width, height = int(m.group(1)), int(m.group(2))
+    print "Comparing %s to %s, resolution %d x %d" % (ref_file, dist_file, width, height)
+
+    ref_fh = open(ref_file, "rb")
+    dist_fh = open(dist_file, "rb")
+
+    frame_num = 0
+    while True:
+        ref, _, _ = img_read_yuv(ref_fh, width, height)
+        dist, _, _ = img_read_yuv(dist_fh, width, height)
+        vifp_value = vifp.vifp_mscale(ref.astype(float), dist.astype(float))
+        ssim_value = ssim.ssim(ref, dist)
+        print "Frame=%d VIFP=%f SSIM=%f" % (frame_num, vifp_value, ssim_value)
+        frame_num += 1
+
+else:
+    # Inputs are image files
+    ref = img_greyscale( scipy.misc.imread(ref_file) ).astype(numpy.float32)
+    dist = img_greyscale( scipy.misc.imread(dist_file) ).astype(numpy.float32)
+    width, height = ref.shape[1], ref.shape[0]
+    print "Comparing %s to %s, resolution %d x %d" % (ref_file, dist_file, width, height)
+    vifp_value = vifp.vifp_mscale(ref, dist)
+    print "VIFP=%f" % (vifp_value)
+    #import timeit
+    #print timeit.timeit('ssim_value = ssim_theano.ssim(ref, dist)', setup='import ssim_theano; from __main__ import ref, dist', number=100)
+    ssim_value = ssim.ssim(ref, dist)
+    print "SSIM=%f" % (ssim_value)
+    psnr_value = psnr.psnr(ref, dist)
+    print "PSNR=%f" % (psnr_value)

psnr.py

@@ -0,0 +1,25 @@
+"""
+Video Quality Metrics
+Copyright (c) 2014 Alex Izvorski <[email protected]>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import numpy
+import math
+
+def psnr(img1, img2):
+	mse = numpy.mean( (img1 - img2) ** 2 )
+	PIXEL_MAX = 255.0
+	return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))

ssim.py

@@ -0,0 +1,69 @@
+"""
+Video Quality Metrics
+Copyright (c) 2014 Alex Izvorski <[email protected]>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import numpy
+from scipy.ndimage import gaussian_filter
+
+from numpy.lib.stride_tricks import as_strided as ast
+
+# Hat tip: http://stackoverflow.com/a/5078155/1828289
+def block_view(A, block=(3, 3)):
+    """Provide a 2D block view to 2D array. No error checking made.
+    Therefore meaningful (as implemented) only for blocks strictly
+    compatible with the shape of A."""
+    # simple shape and strides computations may seem at first strange
+    # unless one is able to recognize the 'tuple additions' involved ;-)
+    shape= (A.shape[0]/ block[0], A.shape[1]/ block[1])+ block
+    strides= (block[0]* A.strides[0], block[1]* A.strides[1])+ A.strides
+    return ast(A, shape= shape, strides= strides)
+
+
+def ssim(img1, img2, C1=0.01**2, C2=0.03**2):
+
+    bimg1 = block_view(img1, (4,4))
+    bimg2 = block_view(img2, (4,4))
+    s1  = numpy.sum(bimg1, (-1, -2))
+    s2  = numpy.sum(bimg2, (-1, -2))
+    ss  = numpy.sum(bimg1*bimg1, (-1, -2)) + numpy.sum(bimg2*bimg2, (-1, -2))
+    s12 = numpy.sum(bimg1*bimg2, (-1, -2))
+
+    vari = ss - s1*s1 - s2*s2
+    covar = s12 - s1*s2
+
+    ssim_map =  (2*s1*s2 + C1) * (2*covar + C2) / ((s1*s1 + s2*s2 + C1) * (vari + C2))
+    return numpy.mean(ssim_map)
+
+# FIXME there seems to be a problem with this code
+def ssim_exact(img1, img2, sd=1.5, C1=0.01**2, C2=0.03**2):
+
+    mu1 = gaussian_filter(img1, sd)
+    mu2 = gaussian_filter(img2, sd)
+    mu1_sq = mu1 * mu1
+    mu2_sq = mu2 * mu2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = gaussian_filter(img1 * img1, sd) - mu1_sq
+    sigma2_sq = gaussian_filter(img2 * img2, sd) - mu2_sq
+    sigma12 = gaussian_filter(img1 * img2, sd) - mu1_mu2
+
+    ssim_num = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2))
+
+    ssim_den = ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    ssim_map = ssim_num / ssim_den
+    return numpy.mean(ssim_map)
+

ssim_theano.py

@@ -0,0 +1,36 @@
+"""
+Video Quality Metrics
+Copyright (c) 2014 Alex Izvorski <[email protected]>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""
+
+import theano
+from theano import tensor
+
+# declare two symbolic floating-point scalars
+img1 = tensor.fmatrix()
+img2 = tensor.fmatrix()
+
+s1 = img1.sum()
+s2 = img2.sum()
+ss = (img1 * img1).sum() + (img2 * img2).sum()
+s12 = (img1 * img2).sum()
+vari = ss - s1*s1 - s2*s2
+covar = s12 - s1*s2
+ssim_c1 = .01*.01
+ssim_c2 = .03*.03
+ssim_value = (2*s1*s2 + ssim_c1) * (2*covar + ssim_c2) / ((s1*s1 + s2*s2 + ssim_c1) * (vari + ssim_c2))
+
+ssim = theano.function([img1, img2], ssim_value)

vifp.py

@@ -113,54 +113,3 @@ def vifp_mscale(ref, dist):
     return vifp
 
 
-
-
-def img_greyscale(img):
-    return 0.299 * img[:,:,0] + 0.587 * img[:,:,1] + 0.114 * img[:,:,2]
-
-def img_read_yuv(src_file, width, height):
-    y_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=(width * height)).reshape( (height, width) )
-    u_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=((width/2) * (height/2))).reshape( (height/2, width/2) )
-    v_img = numpy.fromfile(src_file, dtype=numpy.uint8, count=((width/2) * (height/2))).reshape( (height/2, width/2) )
-    return (y_img, u_img, v_img)
-
-
-if __name__ == '__main__':
-
-    import re, sys, scipy.misc, pdb
-
-    ref_file = sys.argv[1]
-    dist_file = sys.argv[2]
-
-    if ".yuv" in ref_file:
-        # Inputs are uncompressed video in YUV420 planar format
-
-        # Get resolution from file name
-        m = re.search(r"(\d+)x(\d+)", ref_file)
-        if not m:
-            print "Could not find resolution in file name: %s" % (ref_file)
-            exit(1)
-
-        width, height = int(m.group(1)), int(m.group(2))
-        print "Comparing %s to %s, resolution %d x %d" % (ref_file, dist_file, width, height)
-
-        ref_fh = open(ref_file, "rb")
-        dist_fh = open(dist_file, "rb")
-
-        frame_num = 0
-        while True:
-            ref, _, _ = img_read_yuv(ref_fh, width, height)
-            dist, _, _ = img_read_yuv(dist_fh, width, height)
-            vifp = vifp_mscale(ref.astype(float), dist.astype(float))
-            print "Frame=%d VIFP=%f" % (frame_num, vifp)
-            frame_num += 1
-
-    else:
-        # Inputs are image files
-        ref = img_greyscale( scipy.misc.imread(ref_file) )
-        dist = img_greyscale( scipy.misc.imread(dist_file) )
-        width, height = ref.shape[1], ref.shape[0]
-        print "Comparing %s to %s, resolution %d x %d" % (ref_file, dist_file, width, height)
-        vifp = vifp_mscale(ref, dist)
-        print "VIFP=%f" (vifp)
-