Restructured code

Author: abhinaukumar

__init__.py

@@ -1,2 +1,4 @@
-from .utils import *
-from .iqa import *
+from .ssim import ssim, ms_ssim
+from .vif import vif_spatial, vif
+from .brisque import brisque
+from .niqe import niqe

brisque.py

@@ -0,0 +1,14 @@
+from .utils import compute_image_mscn_transform, extract_subband_feats
+
+import numpy as np
+import cv2
+
+
+# TODO: Compare with Matlab
+def brisque(image):
+    y_mscn = compute_image_mscn_transform(image)
+    half_scale = cv2.resize(image, dsize=(0, 0), fx=0.5, fy=0.5, interpolation=cv2.INTER_LANCZOS4)
+    y_half_mscn = compute_image_mscn_transform(half_scale)
+    feats_full = extract_subband_feats(y_mscn)
+    feats_half = extract_subband_feats(y_half_mscn)
+    return np.concatenate((feats_full, feats_half))

iqa.py

@@ -0,0 +1,60 @@
+
+from .utils import compute_image_mscn_transform, extract_on_patches
+
+import numpy as np
+import scipy.io
+import cv2
+import os
+
+
+def niqe(img):
+    blocksizerow = 96
+    blocksizecol = 96
+    h, w = img.shape
+
+    module_path = os.path.dirname(__file__)
+    params = scipy.io.loadmat(os.path.join(module_path, 'niqe_nss_parameters.mat'))
+    mu_prisparam = params['mu_prisparam']
+    cov_prisparam = params['cov_prisparam']
+    if (h < blocksizerow) or (w < blocksizecol):
+        print("Input frame is too small")
+        exit(0)
+
+    # ensure that the patch divides evenly into img
+    hoffset = (h % blocksizerow)
+    woffset = (w % blocksizecol)
+
+    if hoffset > 0:
+        img = img[:-hoffset, :]
+    if woffset > 0:
+        img = img[:, :-woffset]
+
+    img = img.astype(np.float32)
+
+    scale_percent = 50  # percent of original size
+    width = int(img.shape[1] * scale_percent / 100)
+    height = int(img.shape[0] * scale_percent / 100)
+
+    img2 = cv2.resize(img, (height, width), interpolation=cv2.INTER_CUBIC)
+
+    mscn1 = compute_image_mscn_transform(img, extend_mode='nearest')
+    mscn1 = mscn1.astype(np.float32)
+
+    mscn2 = compute_image_mscn_transform(img2, extend_mode='nearest')
+    mscn2 = mscn2.astype(np.float32)
+
+    feats_lvl1 = extract_on_patches(mscn1, blocksizerow, blocksizecol)
+    feats_lvl2 = extract_on_patches(mscn2, blocksizerow/2, blocksizecol/2)
+
+    # stack the scale features
+    feats = np.hstack((feats_lvl1, feats_lvl2))
+
+    mu_distparam = np.mean(feats, axis=0)
+    cov_distparam = np.cov(feats.T)
+
+    invcov_param = np.linalg.pinv((cov_prisparam + cov_distparam)/2)
+
+    xd = mu_prisparam - mu_distparam
+    quality = np.sqrt(np.dot(np.dot(xd, invcov_param), xd.T))[0][0]
+
+    return np.hstack((mu_distparam, [quality]))

niqe.py

@@ -0,0 +1,50 @@
+from .utils import moments
+
+import numpy as np
+
+
+def ssim(img_ref, img_dist, k=11, max_val=1, K1=0.01, K2=0.03, no_lum=False, full=False, padding=None, stride=1):
+    x = img_ref.astype('float32')
+    y = img_dist.astype('float32')
+    mu_x, mu_y, var_x, var_y, cov_xy = moments(x, y, k, stride, padding=padding)
+
+    C1 = (max_val*K1)**2
+    C2 = (max_val*K2)**2
+
+    if not no_lum:
+        l = (2*mu_x*mu_y + C1)/(mu_x**2 + mu_y**2 + C1)
+    cs = (2*cov_xy + C2)/(var_x + var_y + C2)
+
+    ssim_map = cs
+    if not no_lum:
+        ssim_map *= l
+
+    if (full):
+        return (np.mean(ssim_map), ssim_map)
+    else:
+        return np.mean(ssim_map)
+
+
+def ms_ssim(img_ref, img_dist, k=11, max_val=1, K1=0.01, K2=0.03, full=False, padding=None, stride=1):
+    x = img_ref.astype('float32')
+    y = img_dist.astype('float32')
+
+    n_levels = 5
+    weights = np.array([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
+    scores = np.ones((n_levels,))
+    for i in range(n_levels-1):
+        if np.min(x.shape) <= k:
+            break
+        scores[i] = ssim(x, y, k, max_val, K1, K2, no_lum=True, padding=padding, stride=stride)
+        x = x[:(x.shape[0]//2)*2, :(x.shape[1]//2)*2]
+        y = y[:(y.shape[0]//2)*2, :(y.shape[1]//2)*2]
+        x = (x[::2, ::2] + x[1::2, ::2] + x[1::2, 1::2] + x[::2, 1::2])/4
+        y = (y[::2, ::2] + y[1::2, ::2] + y[1::2, 1::2] + y[::2, 1::2])/4
+
+    if np.min(x.shape) > k:
+        scores[-1] = ssim(x, y, k, max_val, K1, K2, no_lum=False, padding=padding, stride=stride)
+    msssim = np.prod(np.power(scores, weights))
+    if full:
+        return msssim, scores
+    else:
+        return msssim