feat: ✨ add niqe_matlab and brisque_matlab

chaofengc · chaofengc · commit 065c2b73eecc · 2024-08-21T16:26:35.000+08:00
diff --git a/Makefile b/Makefile
@@ -23,6 +23,9 @@ test:
 	pytest tests/test_metric_general.py::test_forward -v  
 	pytest tests/test_metric_general.py::test_cpu_gpu_consistency -v
 
+test_cal:
+	pytest tests/ -m calibration -v
+
 test_cs:
 	pytest tests/test_metric_general.py::test_cpu_gpu_consistency -v
 
diff --git a/README.md b/README.md
@@ -22,7 +22,7 @@ This is a comprehensive image quality assessment (IQA) toolbox built with **pure
 ---
 
 ### :triangular_flag_on_post: Updates/Changelog
-- ✨**Aug, 2024**. Add `piqe` metric.
+- ✨**Aug, 2024**. Add `piqe` metric, and `niqe_matlab, brisque_matlab` with default matlab parameters (results have been calibrated with MATLAB R2021b).
 - 💥**Aug, 2024**. Add `lpips+` and `lpips-vgg+` proposed in our paper [TOPIQ](https://arxiv.org/abs/2308.03060). 
 - 🔥**June, 2024**. Add `arniqa` and its variances trained on different datasets, refer to official repo [here](https://github.com/miccunifi/ARNIQA). Thanks for the contribution from [Lorenzo Agnolucci](https://github.com/LorenzoAgnolucci) 🤗.
 - **Apr 24, 2024**. Add `inception_score` and console entry point with `pyiqa` command.
diff --git a/ResultsCalibra/calibration_summary.csv b/ResultsCalibra/calibration_summary.csv
@@ -1,6 +1,8 @@
 Method,I03.bmp,I04.bmp,I06.bmp,I08.bmp,I19.bmp
 brisque,94.6421,-0.1076,0.9929,5.3583,72.2617
 brisque(ours),94.6452,-0.1082,1.0759,5.1423,66.8381
+brisque_matlab,70.68,-2.36,15.8,23.28,60.97
+brisque_matlab(ours),70.6809,-2.369,15.8031,23.0785,60.8247
 ckdn,0.2833,0.5767,0.6367,0.658,0.5999
 ckdn(ours),0.284,0.565,0.6264,0.6414,0.5935
 cw_ssim,0.2763,0.9996,1.0,0.9068,0.8658
@@ -35,6 +37,8 @@ musiq-spaq,17.685,70.492,78.74,79.015,49.105
 musiq-spaq(ours),17.6804,70.6531,79.0364,79.3189,50.4526
 niqe,15.7536,3.6549,3.2355,3.184,8.6352
 niqe(ours),15.6538,3.6549,3.2342,3.1921,9.0722
+niqe_matlab,7.2,2.99,3.17,3.71,7.69
+niqe_matlab(ours),7.1707,2.9908,3.1551,3.6781,7.4782
 nlpd,0.5616,0.0195,0.0159,0.3028,0.4326
 nlpd(ours),0.5616,0.0139,0.011,0.3033,0.4335
 nrqm,1.3894,8.9394,8.9735,6.829,6.312
diff --git a/ResultsCalibra/results_official.csv b/ResultsCalibra/results_official.csv
@@ -1,5 +1,6 @@
 Method,I03.bmp,I04.bmp,I06.bmp,I08.bmp,I19.bmp
 brisque,94.642100000000000,-0.107618000000000,0.992889000000000,5.358270000000000,72.261700000000000
+brisque_matlab,70.68,-2.36,15.80,23.28,60.97
 ckdn,0.2833,0.5767,0.6367,0.658,0.5999
 cw_ssim,0.2763,0.9996,1.0,0.9068,0.8658
 dists,0.4742,0.1424,0.06825,0.02867,0.3123
@@ -9,6 +10,7 @@ lpips,0.7237,0.2572,0.05079,0.05205,0.4253
 mad,195.2796,80.8379,30.3918,84.3542,202.2371
 ms_ssim,0.6733,0.9996,0.9998,0.9566,0.8462
 niqe,15.7536293917814,3.65492152353770,3.23547743716998,3.18403333858339,8.63519663862637
+niqe_matlab,7.20,2.99,3.17,3.71,7.69
 piqe,100.00,21.62,35.86,41.15,76.95
 nlpd,0.561610096893874,0.019534798560102,0.015915631543598,0.302802106557736,0.432604962261603
 psnr,21.11,20.99,27.01,23.3,21.62
diff --git a/docs/ModelCard.md b/docs/ModelCard.md
@@ -47,9 +47,9 @@ print(pyiqa.list_models())
 | CNNIQA                       |  `cnniqa` |
 | NRQM(Ma)<sup>[2](#fn2)</sup> |  `nrqm` | No backward                                                                         |
 | PI(Perceptual Index)         |  `pi` | No backward                                                                         |
-| BRISQUE                      | `brisque` | No backward                                                                         |
+| BRISQUE                      | `brisque`, `brisque_matlab` | No backward                                                                         |
 | ILNIQE                       | `ilniqe` | No backward                                                                         |
-| NIQE                         | `niqe` | No backward                                                                         |
+| NIQE                         | `niqe`, `niqe_matlab` | No backward                                                                         |
 | PIQE                         | `piqe` | No backward                                                                         |
 <!-- </tr>
 </table> -->
diff --git a/pyiqa/archs/brisque_arch.py b/pyiqa/archs/brisque_arch.py
@@ -10,15 +10,19 @@
 
 """
 
+import scipy
+import numpy as np
 import torch
 from pyiqa.utils.color_util import to_y_channel
 from pyiqa.matlab_utils import imresize
+from pyiqa.matlab_utils.nss_feature import compute_nss_features
 from .func_util import estimate_ggd_param, estimate_aggd_param, normalize_img_with_guass
 from pyiqa.utils.download_util import load_file_from_url
 from pyiqa.utils.registry import ARCH_REGISTRY
 
 default_model_urls = {
-    'url': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/brisque_svm_weights.pth'
+    'url': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/brisque_svm_weights.pth',
+    'brisque_matlab': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/brisque_matlab.mat',
 }
 
 
@@ -50,29 +54,45 @@ def brisque(x: torch.Tensor,
         x = to_y_channel(x, 255.)
     else:
         x = x * 255
+    
+    use_matlab_version = 'matlab' in pretrained_model_path
 
     features = []
     num_of_scales = 2
     for _ in range(num_of_scales):
-        features.append(natural_scene_statistics(x, kernel_size, kernel_sigma))
+        if use_matlab_version: 
+            xnorm = normalize_img_with_guass(x, kernel_size, kernel_sigma, padding='replicate')
+            features.append(compute_nss_features(xnorm))
+        else:
+            features.append(natural_scene_statistics(x, kernel_size, kernel_sigma))
         x = imresize(x, scale=0.5, antialiasing=True)
-
     features = torch.cat(features, dim=-1)
     scaled_features = scale_features(features)
 
-    if pretrained_model_path:
+    if pretrained_model_path and not use_matlab_version:
+        # gamma and rho are SVM model parameters taken from official implementation of BRISQUE on MATLAB
+        # Source: https://live.ece.utexas.edu/research/Quality/index_algorithms.htm
         sv_coef, sv = torch.load(pretrained_model_path, weights_only=False)
         sv_coef = sv_coef.to(x)
         sv = sv.to(x)
+        gamma = 0.05
+        rho = -153.591
+    elif use_matlab_version:
+        params = scipy.io.loadmat(pretrained_model_path)
+        sv = params['sv']
+        sv_coef = np.ravel(params['sv_coef'])
+        sv = torch.from_numpy(sv).to(features)
+        sv_coef = torch.from_numpy(sv_coef).to(features)
+        scale = 0.3210
+        scaled_features = features / scale
+        sv = sv / scale
+        gamma = 1
+        rho = - 43.4582
 
-    # gamma and rho are SVM model parameters taken from official implementation of BRISQUE on MATLAB
-    # Source: https://live.ece.utexas.edu/research/Quality/index_algorithms.htm
-    gamma = 0.05
-    rho = -153.591
     sv.t_()
     kernel_features = rbf_kernel(features=scaled_features, sv=sv, gamma=gamma)
-    score = kernel_features @ sv_coef
-    return score - rho
+    score = kernel_features @ sv_coef - rho
+    return score 
 
 
 def natural_scene_statistics(luma: torch.Tensor, kernel_size: int = 7, sigma: float = 7. / 6) -> torch.Tensor:
@@ -137,6 +157,7 @@ def __init__(self,
                  kernel_size: int = 7,
                  kernel_sigma: float = 7 / 6,
                  test_y_channel: bool = True,
+                 version: str = 'original',
                  pretrained_model_path: str = None) -> None:
         super().__init__()
         self.kernel_size = kernel_size
@@ -150,8 +171,10 @@ def __init__(self,
         self.test_y_channel = test_y_channel
         if pretrained_model_path is not None:
             self.pretrained_model_path = pretrained_model_path
-        else:
+        elif version == 'original':
             self.pretrained_model_path = load_file_from_url(default_model_urls['url'])
+        elif version == 'matlab':
+            self.pretrained_model_path = load_file_from_url(default_model_urls['brisque_matlab']) 
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         r"""Computation of BRISQUE score as a loss function.
diff --git a/pyiqa/archs/niqe_arch.py b/pyiqa/archs/niqe_arch.py
@@ -28,6 +28,7 @@
 default_model_urls = {
     'url': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/niqe_modelparameters.mat',
     'niqe': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/niqe_modelparameters.mat',
+    'niqe_matlab': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/niqe_matlab_params.mat',
     'ilniqe': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/ILNIQE_templateModel.mat',
 }
 
@@ -427,6 +428,7 @@ def __init__(self,
                  test_y_channel: bool = True,
                  color_space: str = 'yiq',
                  crop_border: int = 0,
+                 version: str = 'original',
                  pretrained_model_path: str = None) -> None:
 
         super(NIQE, self).__init__()
@@ -436,8 +438,10 @@ def __init__(self,
         self.crop_border = crop_border
         if pretrained_model_path is not None:
             self.pretrained_model_path = pretrained_model_path
-        else:
+        elif version == 'original':
             self.pretrained_model_path = load_file_from_url(default_model_urls['url'])
+        elif version == 'matlab':
+            self.pretrained_model_path = load_file_from_url(default_model_urls['niqe_matlab'])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         r"""Computation of NIQE metric.
diff --git a/pyiqa/default_model_configs.py b/pyiqa/default_model_configs.py
@@ -210,6 +210,16 @@
         'lower_better': True,
         'score_range': '~0, ~100',
     },
+    'niqe_matlab': {
+        'metric_opts': {
+            'type': 'NIQE',
+            'test_y_channel': True,
+            'version': 'matlab',
+        },
+        'metric_mode': 'NR',
+        'lower_better': True,
+        'score_range': '~0, ~100',
+    },
     'ilniqe': {
         'metric_opts': {
             'type': 'ILNIQE',
@@ -227,6 +237,16 @@
         'lower_better': True,
         'score_range': '~0, ~150',
     },
+    'brisque_matlab': {
+        'metric_opts': {
+            'type': 'BRISQUE',
+            'test_y_channel': True,
+            'version': 'matlab',
+        },
+        'metric_mode': 'NR',
+        'lower_better': True,
+        'score_range': '~0, ~150',
+    },
     'nrqm': {
         'metric_opts': {
             'type': 'NRQM',
diff --git a/pyiqa/matlab_utils/nss_feature.py b/pyiqa/matlab_utils/nss_feature.py
@@ -0,0 +1,69 @@
+import torch
+from typing import Tuple
+
+
+def estimate_aggd_param(
+    block: torch.Tensor, return_sigma=False
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Estimate AGGD (Asymmetric Generalized Gaussian Distribution) parameters.
+    Args:
+        block (Tensor): Image block with shape (b, 1, h, w).
+    Returns:
+        Tensor: alpha, beta_l and beta_r for the AGGD distribution
+        (Estimating the parames in Equation 7 in the paper).
+    """
+    gam = torch.arange(0.2, 10 + 0.001, 0.001).to(block)
+    r_gam = (
+        2 * torch.lgamma(2.0 / gam)
+        - (torch.lgamma(1.0 / gam) + torch.lgamma(3.0 / gam))
+    ).exp()
+    r_gam = r_gam.repeat(block.shape[0], 1)
+
+    mask_left = block < 0
+    mask_right = block > 0
+    count_left = mask_left.sum(dim=(-1, -2), dtype=torch.float32)
+    count_right = mask_right.sum(dim=(-1, -2), dtype=torch.float32)
+
+    left_std = torch.sqrt((block * mask_left).pow(2).sum(dim=(-1, -2)) / (count_left))
+    right_std = torch.sqrt(
+        (block * mask_right).pow(2).sum(dim=(-1, -2)) / (count_right)
+    )
+
+    gammahat = left_std / right_std
+    rhat = block.abs().mean(dim=(-1, -2)).pow(2) / block.pow(2).mean(dim=(-1, -2))
+    rhatnorm = (rhat * (gammahat.pow(3) + 1) * (gammahat + 1)) / (
+        gammahat.pow(2) + 1
+    ).pow(2)
+    array_position = (r_gam - rhatnorm).abs().argmin(dim=-1)
+
+    alpha = gam[array_position]
+    beta_l = (
+        left_std.squeeze(-1)
+        * (torch.lgamma(1 / alpha) - torch.lgamma(3 / alpha)).exp().sqrt()
+    )
+    beta_r = (
+        right_std.squeeze(-1)
+        * (torch.lgamma(1 / alpha) - torch.lgamma(3 / alpha)).exp().sqrt()
+    )
+
+    if return_sigma:
+        return alpha, left_std.squeeze(-1), right_std.squeeze(-1)
+    else:
+        return alpha, beta_l, beta_r
+
+
+def compute_nss_features(luma_nrmlzd: torch.Tensor) -> torch.Tensor:
+
+    alpha, betal, betar = estimate_aggd_param(luma_nrmlzd, return_sigma=False)
+    features = [alpha, (betal + betar) / 2]
+
+    shifts = [(0, 1), (1, 0), (1, 1), (-1, 1)]
+
+    for shift in shifts:
+        shifted_luma_nrmlzd = torch.roll(luma_nrmlzd, shifts=shift, dims=(-2, -1))
+        alpha, betal, betar = estimate_aggd_param(luma_nrmlzd * shifted_luma_nrmlzd, return_sigma=False)
+        distmean = (betar - betal) * torch.exp(torch.lgamma(2/alpha) - torch.lgamma(1/alpha))
+        features.extend((alpha, distmean, betal, betar))
+
+    return torch.stack(features, dim=-1)
+
diff --git a/tests/test_metric_general.py b/tests/test_metric_general.py
@@ -14,6 +14,7 @@
 TOL_DICT = {
     'brisque': (1e-2, 8e-2),
     'niqe': (1e-2, 6e-2),
+    'niqe_matlab': (1e-2, 6e-2),
     'pi': (1e-2, 3e-2),
     'ilniqe': (1e-2, 4e-2),
     'ckdn': (1e-2, 3e-2),
