From 5f6d4fbc79d14d56b29c0a0c9c84baa8d33e78c5 Mon Sep 17 00:00:00 2001 From: chaofengc Date: Fri, 16 Jun 2023 15:01:17 +0800 Subject: [PATCH] feat: :triangular_flag_on_post: add entropy for NR IQA --- ResultsCalibra/calibration_summary.csv | 4 +- ResultsCalibra/results_official.csv | 1 + pyiqa/archs/entropy_arch.py | 58 ++++++++++++++++++++++++++ pyiqa/default_model_configs.py | 6 +++ 4 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 pyiqa/archs/entropy_arch.py diff --git a/ResultsCalibra/calibration_summary.csv b/ResultsCalibra/calibration_summary.csv index 66b1efa..72abd18 100644 --- a/ResultsCalibra/calibration_summary.csv +++ b/ResultsCalibra/calibration_summary.csv @@ -7,6 +7,8 @@ cw_ssim,0.2763,0.9996,1.0,0.9068,0.8658 cw_ssim(ours),0.2763,0.9996,1.0,0.9068,0.8658 dists,0.4742,0.1424,0.0683,0.0287,0.3123 dists(ours),0.4742,0.1424,0.0682,0.0287,0.3123 +entropy,6.9511,6.9661,7.5309,7.5566,5.7629 +entropy(ours),6.9511,6.9661,7.5309,7.5565,5.7629 fsim,0.689,0.9702,0.9927,0.9575,0.822 fsim(ours),0.6891,0.9702,0.9927,0.9575,0.822 gmsd,0.2203,0.0005,0.0004,0.1346,0.205 @@ -34,7 +36,7 @@ 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 nrqm(ours),1.3932,8.9419,8.9721,6.8309,6.3031 -paq2piq,44.1340,73.6015,74.3297,76.8748,70.9153 +paq2piq,44.134,73.6015,74.3297,76.8748,70.9153 paq2piq(ours),44.1341,73.6015,74.3297,76.8748,70.9153 pi,11.9235,3.072,2.618,2.8074,6.7713 pi(ours),11.9286,3.073,2.6357,2.7979,6.9546 diff --git a/ResultsCalibra/results_official.csv b/ResultsCalibra/results_official.csv index 3d1b8fa..7103209 100644 --- a/ResultsCalibra/results_official.csv +++ b/ResultsCalibra/results_official.csv @@ -22,3 +22,4 @@ musiq-koniq,12.494,75.332,73.429,75.188,36.938 musiq-paq2piq,46.035,72.660,73.625,74.361,69.006 musiq-spaq,17.685,70.492,78.740,79.015,49.105 paq2piq,44.1340,73.6015,74.3297,76.8748,70.9153 +entropy,6.9511,6.9661,7.5309,7.5566,5.7629 diff --git a/pyiqa/archs/entropy_arch.py b/pyiqa/archs/entropy_arch.py new file mode 100644 index 0000000..3851f3d --- /dev/null +++ b/pyiqa/archs/entropy_arch.py @@ -0,0 +1,58 @@ +r"""Peak signal-to-noise ratio (PSNR) Metric + +Created by: https://github.com/photosynthesis-team/piq + +Modified by: Jiadi Mo (https://github.com/JiadiMo) + +Refer to: + Wikipedia from https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio + QIQA from https://github.com/francois-rozet/piqa/blob/master/piqa/psnr.py + +""" + +import torch +import torch.nn as nn + +from pyiqa.utils.registry import ARCH_REGISTRY +from pyiqa.utils.color_util import to_y_channel + + +def entropy(x, data_range=255., eps=1e-8, color_space='yiq'): + r"""Compute entropy of a gray scale image. + Args: + x: An input tensor. Shape :math:`(N, C, H, W)`. + Returns: + Entropy of the image. + """ + + if (x.shape[1] == 3): + # Convert RGB image to gray scale and use Y-channel + x = to_y_channel(x, data_range, color_space) + + # Compute histogram + hist = nn.functional.one_hot(x.long(), num_classes=int(data_range + 1)).sum(dim=[1, 2, 3]) + hist = hist / hist.sum(dim=1, keepdim=True) + + # Compute entropy + score = -torch.sum(hist * torch.log2(hist + eps), dim=1) + + return score + + +@ARCH_REGISTRY.register() +class Entropy(nn.Module): + r""" + Args: + x (torch.Tensor): image tensor with shape (B, _, H, W), range [0, 1] + Return: + score (torch.Tensor): (B, 1) + """ + + def __init__(self, **kwargs): + super().__init__() + self.kwargs = kwargs + + def forward(self, x): + + score = entropy(x, **self.kwargs) + return score diff --git a/pyiqa/default_model_configs.py b/pyiqa/default_model_configs.py index facdcc4..139ca83 100644 --- a/pyiqa/default_model_configs.py +++ b/pyiqa/default_model_configs.py @@ -344,5 +344,11 @@ 'type': 'CLIPScore', }, 'metric_mode': 'NR', # Caption image similarity + }, + 'entropy': { + 'metric_opts': { + 'type': 'Entropy', + }, + 'metric_mode': 'NR', } })