Initial commit

.gitattributes

@@ -0,0 +1,42 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/01_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/02_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/03_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/05_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/06_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/08_1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/10_1.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: FR IQA
+emoji: 🌖
+colorFrom: gray
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.42.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: Compute similarity between two images using FR-IQA metrics
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,125 @@
+import os
+import gradio as gr
+from DISTS_pytorch import DISTS
+from torchvision.io import read_image
+import torch
+import torchvision.transforms.v2 as transforms
+import pyiqa
+import spaces
+from metrics.DeepDC import DeepDC
+from metrics.DeepWSD import DeepWSD
+from metrics.ADISTS import ADISTS
+
+@spaces.GPU(duration=5)
+class Evaluator:
+    def __init__(self, device):
+        self.device = device
+        self.metrics = {
+            "↓ MSE": torch.nn.functional.mse_loss,
+            "↓ L1": torch.nn.functional.l1_loss,
+            "↓ DISTS": DISTS().to(self.device),
+            "↓ LPIPS": pyiqa.create_metric("lpips", device=self.device),
+            "↑ PSNR": pyiqa.create_metric("psnr", device=self.device),
+            "↑ SSIM": pyiqa.create_metric("ssim", device=self.device),
+            "↑ MS-SSIM": pyiqa.create_metric("ms_ssim", device=self.device),
+            "↑ CW-SSIM": pyiqa.create_metric("cw_ssim", device=self.device),
+            "↑ FSIM": pyiqa.create_metric("fsim", device=self.device),
+            "↑ DeepDC": DeepDC().to(self.device),
+            "↑ DeepWSD": DeepWSD().to(self.device),
+            "↑ ADISTS": ADISTS().to(self.device),
+        }
+        self.transform = transforms.ToDtype(dtype=torch.float32, scale=True)
+
+    @torch.no_grad()
+    def evaluate(self, img_fname1, img_fname2):
+        img1 = self.transform(read_image(img_fname1)).unsqueeze(0).to(self.device)
+        img2 = self.transform(read_image(img_fname2)).unsqueeze(0).to(self.device)
+
+        # check images are the same size
+        if img1.shape != img2.shape:
+            raise gr.Error("Input images must have the same dimensions!")
+
+        return "\n".join(
+            f"{name:<10}: {float(metric(img1, img2).item()):3,.5f}"
+            for name, metric in self.metrics.items()
+        )
+
+
+@spaces.GPU(duration=1)
+def get_evaluator():
+    """Returns a singleton Evaluator instance per worker/session."""
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if not hasattr(get_evaluator, "evaluator"):
+        get_evaluator.evaluator = Evaluator(device)
+    return get_evaluator.evaluator
+
+
+@spaces.GPU(duration=5)
+def compute_similarity(img_fname1, img_fname2):
+    return get_evaluator().evaluate(img_fname1, img_fname2)
+
+
+if __name__ == "__main__":
+    examples = [
+    ["examples/01_1.jpg", "examples/01_1.jpg"],  # Add an extra example for identical images
+    *[[f"examples/{i:02d}_1.jpg", f"examples/{i:02d}_2.jpg"] for i in range(1, 10)],
+    ]
+
+    # Gradio UI
+    custom_css = """
+    .center-header {
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        margin: 0 0 10px 0;
+    }
+    .monospace-text {
+        font-family: 'Courier New', Courier, monospace;
+    }
+    """
+    with gr.Blocks(title="FR-IQA", css=custom_css) as demo:
+        gr.Markdown(f"""
+        <div class='center-header'><h1>FR-IQA</h1></div>
+        Upload two images to compute various Full-Reference IQA metrics for measuring similarity between them.<br>
+        <b>Note</b>: Images must be of the same size.
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                img_fname1 = gr.Image(type="filepath", label="Image#1", height=512, width=512)
+            with gr.Column(scale=2):
+                img_fname2 = gr.Image(type="filepath", label="Image#2", height=512, width=512)
+            with gr.Column(scale=1):
+                metrics_output = gr.Textbox(label="Metrics Output", lines=22, elem_classes="monospace-text", show_copy_button=True)
+
+        with gr.Row():
+            submit_btn = gr.Button("Compute!")
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                gr.Examples(
+                    examples=examples,
+                    inputs=[img_fname1, img_fname2],
+                    fn=compute_similarity,
+                    outputs=metrics_output,
+                    label="Example Image Pairs (all images are 1024×768)",
+                    cache_examples=False,
+                    examples_per_page=5
+                )
+            with gr.Column(scale=2):
+                gr.Markdown("""
+                <div class='center-header'><h3>Acknowledgements</h3></div>
+
+                - Example images are from the [TryOffDiff](https://rizavelioglu.github.io/tryoffdiff) paper, which are sampled from VITON-HD dataset.
+                - We use the [IQA-PyTorch](https://github.com/chaofengc/IQA-PyTorch) library for computing the metrics.
+
+
+                """)
+
+        submit_btn.click(
+            fn=compute_similarity,
+            inputs=[img_fname1, img_fname2],
+            outputs=[metrics_output]
+        )
+
+    demo.launch(share=False, ssr_mode=False)

metrics/ADISTS.py

@@ -0,0 +1,157 @@
+"""Taken from:
+https://github.com/dingkeyan93/A-DISTS/blob/3d20592648625df2e451c9aba25bbaf3c7952ac8/A-DISTS.py
+"""
+import torch
+import torch.nn as nn
+import numpy as np
+import torchvision
+from torchvision import models,transforms
+from torch.nn.functional import normalize
+import torch.nn.functional as F
+import math
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class Downsample(nn.Module):
+    def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
+        super(Downsample, self).__init__()
+        self.padding = (filter_size - 2 )//2
+        self.stride = stride
+        self.channels = channels
+        a = np.hanning(filter_size)[1:-1]
+        g = torch.Tensor(a[:,None]*a[None,:])
+        g = g/torch.sum(g)
+        self.register_buffer('filter', g[None,None,:,:].repeat((self.channels,1,1,1)))
+        # print (g)   
+
+    def forward(self, input):
+        input = input**2
+        out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1])
+        return (out+1e-12).sqrt()
+
+class ADISTS(torch.nn.Module):
+    def __init__(self, window_size=21):
+        super(ADISTS, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=True).features
+
+        self.stage1 = torch.nn.Sequential()
+        self.stage2 = torch.nn.Sequential()
+        self.stage3 = torch.nn.Sequential()
+        self.stage4 = torch.nn.Sequential()
+        self.stage5 = torch.nn.Sequential()
+
+        for x in range(0,4):
+            self.stage1.add_module(str(x), vgg_pretrained_features[x])
+        self.stage2.add_module(str(4), Downsample(channels=64))
+        for x in range(5, 9):
+            self.stage2.add_module(str(x), vgg_pretrained_features[x])
+        self.stage3.add_module(str(9), Downsample(channels=128))
+        for x in range(10, 16):
+            self.stage3.add_module(str(x), vgg_pretrained_features[x])
+        self.stage4.add_module(str(16), Downsample(channels=256))
+        for x in range(17, 23):
+            self.stage4.add_module(str(x), vgg_pretrained_features[x])
+        self.stage5.add_module(str(23), Downsample(channels=512))
+        for x in range(24, 30):
+            self.stage5.add_module(str(x), vgg_pretrained_features[x])
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1,-1,1,1))
+        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1,-1,1,1))
+
+        self.chns = [3,64,128,256,512,512]
+        self.windows=nn.ParameterList()
+        self.window_size = window_size
+        for k in range(len(self.chns)):
+            self.windows.append(self.create_window(self.window_size, self.window_size/3, self.chns[k]))
+
+
+    def compute_prob(self, x, k):
+
+        theta = [[0, 0], 
+                [1.0, 0.29],
+                [2.0, 0.52],
+                [2.95, 0.56],
+                [0.97, 0.25],
+                [0.21, 0.10]]
+
+        ps = 1/(1+torch.exp(-(x-theta[k][0])/theta[k][1]))
+        pt = 1 - ps
+        return ps, pt
+
+
+    def gaussian(self,window_size, sigma):
+        gauss = torch.Tensor([math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+        return gauss/gauss.sum()
+
+    def create_window(self,window_size, window_sigma, channel):
+        _1D_window = self.gaussian(window_size, window_sigma).unsqueeze(1)
+        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+        return nn.Parameter(window,requires_grad=False)
+
+    def forward_once(self, x):
+        h = (x-self.mean)/self.std
+        h = self.stage1(h)
+        h_relu1_2 = h
+        h = self.stage2(h)
+        h_relu2_2 = h
+        h = self.stage3(h)
+        h_relu3_3 = h
+        h = self.stage4(h)
+        h_relu4_3 = h
+        if len(self.chns)==6:
+            h = self.stage5(h)
+            h_relu5_3 = h
+            outs = [x,h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3]
+        else:
+            outs = [x,h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3]
+        return outs
+  
+    def forward(self, x, y, as_loss=False):
+        assert x.shape == y.shape
+        if as_loss:
+            feats_x = self.forward_once(x)
+            feats_y = self.forward_once(y)   
+        else:
+            with torch.no_grad():
+                feats_x = self.forward_once(x)
+                feats_y = self.forward_once(y) 
+    
+        D = 0 
+        c1 = 1e-6
+        c2 = 1e-6
+        pad = nn.ReflectionPad2d(0)
+        for k in range(len(self.chns)-1,-1,-1):
+            try:
+                x_mean = F.conv2d(pad(feats_x[k]), self.windows[k], stride =1, padding = 0, groups = self.chns[k])
+                y_mean = F.conv2d(pad(feats_y[k]), self.windows[k], stride =1, padding = 0, groups = self.chns[k])
+                x_var = F.conv2d(pad(feats_x[k]**2), self.windows[k], stride =1, padding = 0, groups = self.chns[k]) - x_mean**2
+                y_var = F.conv2d(pad(feats_y[k]**2), self.windows[k], stride =1, padding = 0, groups = self.chns[k]) - y_mean**2
+                xy_cov = F.conv2d(pad(feats_x[k]*feats_y[k]), self.windows[k], stride =1, padding = 0, groups = self.chns[k]) - x_mean*y_mean
+            except:
+                x_mean = feats_x[k].mean([2,3], keepdim=True)
+                y_mean = feats_y[k].mean([2,3], keepdim=True)
+                x_var = ((feats_x[k]-x_mean)**2).mean([2,3], keepdim=True)
+                y_var = ((feats_y[k]-y_mean)**2).mean([2,3], keepdim=True)
+                xy_cov = (feats_x[k]*feats_y[k]).mean([2,3],keepdim=True) - x_mean*y_mean
+
+            T = (2*x_mean*y_mean+c1)/(x_mean**2+y_mean**2+c1)
+            S = (2*xy_cov+c2)/(x_var+y_var+c2)
+
+            if k>0:
+                ratio = torch.mean(x_var/(x_mean+1e-12),dim=1,keepdim=True)
+                ps, pt = self.compute_prob(ratio,k)
+
+            D_map = pt*T+ps*S
+            # D = D + D_map.mean([2,3]).mean(1)/len(self.chns)
+            D = D + D_map.mean([2,3]).sum(1)/sum(self.chns)
+
+        if as_loss:
+            return 1-D.mean()
+        else:
+            return 1-D
+  

metrics/DeepDC.py

@@ -0,0 +1,134 @@
+"""Taken from: https://github.com/h4nwei/DeepDC
+"""
+import torch
+import torch.nn as nn
+from typing import Optional
+from collections import OrderedDict
+import numpy as np
+import torch.nn.functional as F
+import math
+import torchvision
+from torchvision import models, transforms
+
+
+names = {'vgg19': ['image', 'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',
+                   'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+                   'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2',
+                   'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',
+                   'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2',
+                   'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',
+                   'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2',
+                   'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5'],}
+
+class MultiVGGFeaturesExtractor(nn.Module):
+    def __init__(self, target_features=('conv1_2', 'conv2_2', 'conv3_4', 'conv4_4', 'conv5_4'), use_input_norm=False, requires_grad=False): # ALL FALSE is the best for COS_Similarity; Correlation: use_norm = True
+        super(MultiVGGFeaturesExtractor, self).__init__()
+        self.use_input_norm = use_input_norm
+        self.target_features = target_features
+
+
+        model = torchvision.models.vgg19(pretrained=True)
+        names_key = 'vgg19'
+
+        if self.use_input_norm:
+            mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+            std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+            self.register_buffer('mean', mean)
+            self.register_buffer('std', std)
+
+        self.target_indexes = [names[names_key].index(k) - 1 for k in self.target_features]
+        self.features = nn.Sequential(*list(model.features.children())[:(max(self.target_indexes) + 1)])
+
+        if not requires_grad:
+            for k, v in self.features.named_parameters():
+                v.requires_grad = False
+            self.features.eval()
+
+    def forward(self, x):
+        # assume input range is [0, 1]
+        if self.use_input_norm:
+            x = (x - self.mean) / self.std
+
+        y = OrderedDict()
+        if 'image' in self.target_features:
+            y.update({"image": x})
+        for key, layer in self.features._modules.items():
+            x = layer(x)
+            # x = self._normalize_tensor(x)
+            if int(key)  in self.target_indexes:
+                y.update({self.target_features[self.target_indexes.index(int(key))]: x})
+        return y
+
+    def _normalize_tensor(sefl, in_feat, eps=1e-10):
+        norm_factor = torch.sqrt(torch.sum(in_feat**2, dim=1, keepdim=True))
+        return in_feat / (norm_factor + eps)
+
+
+class DeepDC(nn.Module):
+    def __init__(self, features_to_compute=('conv1_2', 'conv2_2', 'conv3_4', 'conv4_4', 'conv5_4')):
+        super(DeepDC, self).__init__()
+        self.MSE = torch.nn.MSELoss()
+        self.features_extractor = MultiVGGFeaturesExtractor(target_features=features_to_compute).eval()
+
+    def forward(self, x, y):
+        r"""Compute IQA using DeepDC model.
+
+        Args:
+            - x: An input tensor with (N, C, H, W) shape. RGB channel order for colour images.
+            - y: An reference tensor with (N, C, H, W) shape. RGB channel order for colour images.
+
+        Returns:
+            Value of DeepDC model.
+
+        """
+        targets, inputs = x, y
+        inputs_fea = self.features_extractor(inputs)
+        
+        with torch.no_grad():
+            targets_fea =self.features_extractor(targets)
+
+        dc_scores = []
+
+        for _, key in enumerate(inputs_fea.keys()):
+            inputs_dcdm = self._DCDM(inputs_fea[key])
+            targets_dcdm = self._DCDM(targets_fea[key])
+            dc_scores.append(self.Distance_Correlation(inputs_dcdm, targets_dcdm))
+
+        dc_scores = torch.stack(dc_scores, dim=1)
+
+        score = 1 - dc_scores.mean(dim=1, keepdim=True)
+
+        return  score
+
+
+    # double-centered distance matrix (dcdm)
+    def _DCDM(self, x):
+        if len(x.shape)==4:
+            batchSize, dim, h, w = x.data.shape
+            M = h * w
+        elif len(x.shape)==3:
+            batchSize, M, dim = x.data.shape
+        x = x.reshape(batchSize, dim, M)
+        t = torch.log((1. / (torch.tensor(dim) * torch.tensor(dim))) )
+
+        I = torch.eye(dim, dim, device=x.device).view(1, dim, dim).repeat(batchSize, 1, 1).type(x.dtype)
+        I_M = torch.ones(batchSize, dim, dim, device=x.device).type(x.dtype)
+        x_pow2 = x.bmm(x.transpose(1, 2))
+        dcov = I_M.bmm(x_pow2 * I) + (x_pow2 * I).bmm(I_M) - 2 * x_pow2
+
+        dcov = torch.clamp(dcov, min=0.0)
+        dcov = torch.exp(t)* dcov
+        dcov = torch.sqrt(dcov + 1e-5)
+        dcdm = dcov - 1. / dim * dcov.bmm(I_M) - 1. / dim * I_M.bmm(dcov) + 1. / (dim * dim) * I_M.bmm(dcov).bmm(I_M)
+
+        return dcdm
+
+
+    def Distance_Correlation(self, matrix_A, matrix_B):
+
+        Gamma_XY = torch.sum(matrix_A * matrix_B, dim=[1,2])
+        Gamma_XX = torch.sum(matrix_A * matrix_A, dim=[1,2])
+        Gamma_YY = torch.sum(matrix_B * matrix_B, dim=[1,2])
+        c = 1e-6
+        correlation_r = (Gamma_XY + c) / (torch.sqrt(Gamma_XX * Gamma_YY) + c)
+        return correlation_r

metrics/DeepWSD.py

@@ -0,0 +1,169 @@
+"""Taken from: https://github.com/Buka-Xing/DeepWSD/blob/main/utils.py
+"""
+import numpy as np
+import os
+import sys
+import torch
+from torchvision import models,transforms
+import torch.nn as nn
+import torch.nn.functional as F
+import inspect
+from ot.lp import wasserstein_1d
+
+
+# Process input of VGG16 to make it close to 256
+def downsample(img1, img2, maxSize = 256):
+    _,channels,H,W = img1.shape
+    f = int(max(1,np.round(max(H,W)/maxSize)))
+
+    aveKernel = (torch.ones(channels,1,f,f)/f**2).to(img1.device)
+    img1 = F.conv2d(img1, aveKernel, stride=f, padding = 0, groups = channels)
+    img2 = F.conv2d(img2, aveKernel, stride=f, padding = 0, groups = channels)
+    # For an extremely Large image, the larger window will use to increase the receptive field.
+    if f >= 5:
+        win = 16
+    else:
+        win = 4
+    return img1, img2, win, f
+
+
+# Use L2pooling for VGG16 networks.
+# Original Maxpooling will generate distortions in color channels during optimization.
+class L2pooling(nn.Module):
+    def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
+        super(L2pooling, self).__init__()
+        self.padding = (filter_size - 2 )//2
+        self.stride = stride
+        self.channels = channels
+        a = np.hanning(filter_size)[1:-1]
+
+        g = torch.Tensor(a[:,None]*a[None,:])
+        g = g/torch.sum(g)
+        self.register_buffer('filter', g[None,None,:,:].repeat((self.channels,1,1,1)))
+
+    def forward(self, input):
+        input = input**2
+        out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1])
+        return (out+1e-12).sqrt()
+
+
+def ws_distance(X,Y,P=2,win=4):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    chn_num = X.shape[1]
+    X_sum = X.sum().sum()
+    Y_sum = Y.sum().sum()
+
+    X_patch   = torch.reshape(X,[win,win,chn_num,-1])
+    Y_patch   = torch.reshape(Y,[win,win,chn_num,-1])
+    patch_num = (X.shape[2]//win) * (X.shape[3]//win)
+
+    X_1D = torch.reshape(X_patch,[-1,chn_num*patch_num])
+    Y_1D = torch.reshape(Y_patch,[-1,chn_num*patch_num])
+
+    X_1D_pdf = X_1D / (X_sum + 1e-6)
+    Y_1D_pdf = Y_1D / (Y_sum + 1e-6)
+
+    interval = np.arange(0, X_1D.shape[0], 1)
+    all_samples = torch.from_numpy(interval).to(device).repeat([patch_num*chn_num,1]).t()
+
+    X_pdf = X_1D * X_1D_pdf
+    Y_pdf = Y_1D * Y_1D_pdf
+
+    wsd   = wasserstein_1d(all_samples, all_samples, X_pdf, Y_pdf, P)
+
+    L2 = ((X_1D - Y_1D) ** 2).sum(dim=0)
+    w  =  (1 / ( torch.sqrt(torch.exp( (- 1/(wsd+10) ))) * (wsd+10)**2))
+
+    final = wsd + L2 * w
+    # final = wsd
+
+    return final.sum()
+
+
+class DeepWSD(torch.nn.Module):
+
+    def __init__(self, channels=3, load_weights=True):
+        assert channels == 3
+        super(DeepWSD, self).__init__()
+        self.window = 4
+
+        vgg_pretrained_features = models.vgg16(pretrained=True).features
+        self.stage1 = torch.nn.Sequential()
+        self.stage2 = torch.nn.Sequential()
+        self.stage3 = torch.nn.Sequential()
+        self.stage4 = torch.nn.Sequential()
+        self.stage5 = torch.nn.Sequential()
+
+        # Rewrite the output layer of every block in the VGG network: maxpool->l2pool
+        for x in range(0, 4):
+            self.stage1.add_module(str(x), vgg_pretrained_features[x])
+        self.stage2.add_module(str(4), L2pooling(channels=64))
+        for x in range(5, 9):
+            self.stage2.add_module(str(x), vgg_pretrained_features[x])
+        self.stage3.add_module(str(9), L2pooling(channels=128))
+        for x in range(10, 16):
+            self.stage3.add_module(str(x), vgg_pretrained_features[x])
+        self.stage4.add_module(str(16), L2pooling(channels=256))
+        for x in range(17, 23):
+            self.stage4.add_module(str(x), vgg_pretrained_features[x])
+        self.stage5.add_module(str(23), L2pooling(channels=512))
+        for x in range(24, 30):
+            self.stage5.add_module(str(x), vgg_pretrained_features[x])
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.chns = [3, 64, 128, 256, 512, 512]
+
+    def forward_once(self, x):
+        h = x
+        h = self.stage1(h)
+        h_relu1_2 = h
+        h = self.stage2(h)
+        h_relu2_2 = h
+        h = self.stage3(h)
+        h_relu3_3 = h
+        h = self.stage4(h)
+        h_relu4_3 = h
+        h = self.stage5(h)
+        h_relu5_3 = h
+        return [x, h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3]
+
+    def forward(self, x, y, as_loss=False, resize=True):
+        assert x.shape == y.shape
+        if resize:
+            x, y, window, f = downsample(x, y)
+        if as_loss:
+            feats0 = self.forward_once(x)
+            feats1 = self.forward_once(y)
+        else:
+            with torch.no_grad():
+                feats0 = self.forward_once(x)
+                feats1 = self.forward_once(y)
+        score = 0
+        layer_score=[]
+        # To see score of each layer, use debugging mode of Pycharm.
+
+        for k in range(len(self.chns)):
+            row_padding = round(feats0[k].size(2) / window) * window - feats0[k].size(2)
+            column_padding = round(feats0[k].size(3) / window) * window - feats0[k].size(3)
+
+            pad = nn.ZeroPad2d((column_padding, 0, 0, row_padding))
+            feats0_k = pad(feats0[k])
+            feats1_k = pad(feats1[k])
+
+            tmp = ws_distance(feats0_k, feats1_k, win=window)
+            layer_score.append(torch.log(tmp + 1))
+            score = score + tmp
+        score = score / (k+1)
+
+        # For optimization, the logrithm will not use.
+        if as_loss:
+            return score
+        # We find use log**2 output will lead to higher PLCC results, thus we provide two output strategies
+        # They will only affect PLCC of Quality Assessment Results.
+        elif f==1:
+            return torch.log(score + 1)
+        else:
+            return torch.log(score + 1)**2

requirements.txt

@@ -0,0 +1,5 @@
+DISTS-pytorch
+torch
+torchvision
+pyiqa
+POT