fix: 🚑 fix ssim calculation in nrqm

Author: chaofengc

pyiqa/archs/func_util.py

@@ -2,12 +2,22 @@
 import torch
 import torch.nn.functional as F
 
+from pyiqa.utils.color_util import to_y_channel
 from pyiqa.matlab_utils import fspecial, imfilter
 from .arch_util import excact_padding_2d
 
 EPS = torch.finfo(torch.float32).eps
 
 
+def preprocess_rgb(x, test_y_channel, data_range=1., color_space='yiq'):
+    if test_y_channel and x.shape[1] == 3:
+        x = to_y_channel(x, data_range, color_space)
+    else:
+        x = x * data_range
+        x = x - x.detach() + x.round()
+    return x
+
+
 def extract_2d_patches(x, kernel, stride=1, dilation=1, padding='same'):
     """
     Ref: https://stackoverflow.com/a/65886666

pyiqa/archs/nrqm_arch.py

@@ -19,7 +19,7 @@
 from pyiqa.utils.download_util import load_file_from_url
 from pyiqa.matlab_utils import imresize, fspecial, SCFpyr_PyTorch, dct2d, im2col
 from pyiqa.archs.func_util import extract_2d_patches
-from pyiqa.archs.ssim_arch import SSIM
+from pyiqa.archs.ssim_arch import ssim as ssim_func
 from pyiqa.archs.arch_util import ExactPadding2d
 from pyiqa.archs.niqe_arch import NIQE
 from warnings import warn
@@ -211,7 +211,7 @@ def norm_sender_normalized(pyr, num_scale=2, num_bands=6, blksz=3, eps=1e-12):
             L, Q = torch.linalg.eigh(C_x)
             L_pos = L * (L > 0)
             L_pos_sum = L_pos.sum(dim=1, keepdim=True)
-            L = L_pos * L.sum(dim=1, keepdim=True) / (L_pos_sum + (L_pos_sum == 0).float())
+            L = L_pos * L.sum(dim=1, keepdim=True) / (L_pos_sum + (L_pos_sum == 0).to(L.dtype))
             C_x = Q @ torch.diag_embed(L) @ Q.transpose(1, 2)
 
             o_c = current_band[:, border:-border, border:-border]
@@ -267,16 +267,16 @@ def global_gsm(img: Tensor):
 
     # structure correlation between scales
     hp_band = pyr[0]
-    ssim_func = SSIM(channels=1, test_y_channel=False)
-    for sb in subbands:
-        sb_tmp = imresize(sb, sizes=hp_band.shape[1:]).unsqueeze(1)
-        tmp_ssim = ssim_func(sb_tmp, hp_band.unsqueeze(1))
-        feat.append(tmp_ssim)
+    for sb in lp_bands:
+        curr_band = imresize(sb, sizes=hp_band.shape[1:]).unsqueeze(1)
+        _, tmpscore = ssim_func(curr_band, hp_band.unsqueeze(1), get_cs=True, data_range=255)
+        feat.append(tmpscore)
 
     # structure correlation between orientations
     for i in range(num_bands):
         for j in range(i + 1, num_bands):
-            feat.append(ssim_func(subbands[i].unsqueeze(1), subbands[j].unsqueeze(1)))
+            _, tmpscore = ssim_func(subbands[i].unsqueeze(1), subbands[j].unsqueeze(1), get_cs=True, data_range=255)
+            feat.append(tmpscore)
 
     feat = torch.stack(feat, dim=1)
     return feat
@@ -316,7 +316,7 @@ def random_forest_regression(feat, ldau, rdau, threshold_value, pred_value, best
                                          best_attri[:, i])
             tmp_pred.append(tmp_result)
         pred.append(tmp_pred)
-    pred = torch.Tensor(pred)
+    pred = torch.tensor(pred)
     return pred.mean(dim=1, keepdim=True)
 
 
@@ -335,7 +335,8 @@ def nrqm(
 
     # crop image
     b, c, h, w = img.shape
-    img_pyr = get_guass_pyramid(img.float() / 255.)
+    img = img.double()
+    img_pyr = get_guass_pyramid(img / 255.)
 
     # DCT features
     f1 = []
@@ -359,7 +360,7 @@ def nrqm(
     for feat, rf in zip([f1, f2, f3], rf_param):
         tmp_pred = random_forest_regression(feat, *rf)
         preds = torch.cat((preds, tmp_pred), dim=1)
-    quality = preds @ torch.Tensor(linear_param)
+    quality = preds @ torch.tensor(linear_param)
 
     return quality.squeeze()
 
@@ -396,7 +397,7 @@ def calculate_nrqm(img: torch.Tensor,
 
     if test_y_channel and img.shape[1] == 3:
         img = to_y_channel(img, 255, color_space)
-
+    
     if crop_border != 0:
         img = img[..., crop_border:-crop_border, crop_border:-crop_border]
 

pyiqa/archs/ssim_arch.py

@@ -24,30 +24,21 @@
 from pyiqa.utils.color_util import to_y_channel
 from pyiqa.matlab_utils import fspecial, SCFpyr_PyTorch, math_util, filter2
 from pyiqa.utils.registry import ARCH_REGISTRY
+from .func_util import preprocess_rgb
 
 
 def ssim(X,
          Y,
-         win,
+         win=None,
          get_ssim_map=False,
          get_cs=False,
          get_weight=False,
          downsample=False,
          data_range=1.,
-         test_y_channel=True,
-         color_space='yiq'):
-
-    data_range = 255
-    # Whether calculate on y channel of ycbcr
-    if test_y_channel and X.shape[1] == 3:
-        X = to_y_channel(X, data_range, color_space)
-        Y = to_y_channel(Y, data_range, color_space)
-    else:
-        X = X * data_range
-        X = X - X.detach() + X.round()
-        Y = Y * data_range
-        Y = Y - Y.detach() + Y.round()
-
+         ):
+    if win is None:
+        win = fspecial(11, 1.5, X.shape[1]).to(X)
+    
     C1 = (0.01 * data_range)**2
     C2 = (0.03 * data_range)**2
 
@@ -58,8 +49,6 @@ def ssim(X,
         X = F.avg_pool2d(X, kernel_size=f)
         Y = F.avg_pool2d(Y, kernel_size=f)
 
-    win = win.to(X.device)
-
     mu1 = filter2(X, win, 'valid')
     mu2 = filter2(Y, win, 'valid')
     mu1_sq = mu1.pow(2)
@@ -98,11 +87,11 @@ class SSIM(torch.nn.Module):
     def __init__(self, channels=3, downsample=False, test_y_channel=True, color_space='yiq', crop_border=0.):
 
         super(SSIM, self).__init__()
-        self.win = fspecial(11, 1.5, channels)
         self.downsample = downsample
         self.test_y_channel = test_y_channel
         self.color_space = color_space
         self.crop_border = crop_border
+        self.data_range = 255
 
     def forward(self, X, Y):
         assert X.shape == Y.shape, f'Input {X.shape} and reference images should have the same shape'
@@ -111,14 +100,11 @@ def forward(self, X, Y):
             crop_border = self.crop_border
             X = X[..., crop_border:-crop_border, crop_border:-crop_border]
             Y = Y[..., crop_border:-crop_border, crop_border:-crop_border]
+        
+        X = preprocess_rgb(X, self.test_y_channel, self.data_range, self.color_space)
+        Y = preprocess_rgb(Y, self.test_y_channel, self.data_range, self.color_space) 
 
-        score = ssim(
-            X,
-            Y,
-            win=self.win,
-            downsample=self.downsample,
-            test_y_channel=self.test_y_channel,
-            color_space=self.color_space)
+        score = ssim(X, Y, data_range=self.data_range, downsample=self.downsample)
         return score
 
 
@@ -185,11 +171,11 @@ class MS_SSIM(torch.nn.Module):
 
     def __init__(self, channels=3, downsample=False, test_y_channel=True, is_prod=True, color_space='yiq'):
         super(MS_SSIM, self).__init__()
-        self.win = fspecial(11, 1.5, channels)
         self.downsample = downsample
         self.test_y_channel = test_y_channel
         self.color_space = color_space
         self.is_prod = is_prod
+        self.data_range = 255
 
     def forward(self, X, Y):
         """Computation of MS-SSIM metric.
@@ -201,14 +187,16 @@ def forward(self, X, Y):
         """
         assert X.shape == Y.shape, 'Input and reference images should have the same shape, but got'
         f'{X.shape} and {Y.shape}'
+
+        X = preprocess_rgb(X, self.test_y_channel, self.data_range, self.color_space)
+        Y = preprocess_rgb(Y, self.test_y_channel, self.data_range, self.color_space) 
+
         score = ms_ssim(
-            X,
-            Y,
-            win=self.win,
+            X, Y,
+            data_range=self.data_range,
             downsample=self.downsample,
-            test_y_channel=self.test_y_channel,
-            is_prod=self.is_prod,
-            color_space=self.color_space)
+            is_prod=self.is_prod
+            )
         return score
 
 

pyiqa/matlab_utils/scfpyr_util.py

@@ -63,7 +63,7 @@ def build(self, im_batch):
 
         assert im_batch.device == self.device, 'Devices invalid (pyr = {}, batch = {})'.format(
             self.device, im_batch.device)
-        assert im_batch.dtype == torch.float32, 'Image batch must be torch.float32'
+        # assert im_batch.dtype == torch.float32, 'Image batch must be torch.float32'
         assert im_batch.dim() == 4, 'Image batch must be of shape [N,C,H,W]'
         assert im_batch.shape[1] == 1, 'Second dimension must be 1 encoding grayscale image'