Add files via upload

Author: Z7Gao

MSGlance/Global.py

@@ -0,0 +1,41 @@
+from .Local import LocalGlance
+import torch
+import torch.nn as nn
+from matplotlib import pyplot as plt
+
+class GlobalGlance(nn.Module):
+    def __init__(self, kernel_size=4, stride=4, repeat_time=10, patch_height=64, patch_width=64, drop=None, sigma=1.5):
+        super(GlobalGlance, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.repeat_time = repeat_time
+        self.patch_height = patch_height
+        self.patch_width = patch_width
+        self.sigma = sigma
+        self.ssim_loss = LocalGlance(window_size=self.kernel_size, stride=self.stride, drop=drop, sigma=self.sigma)
+    
+    def forward(self, src_vec, tar_vec):
+        batch_size, channels, height, width = src_vec.size()
+        loss = 0.0
+
+        for batch in range(batch_size):
+            index_list = []
+            for i in range(self.repeat_time):
+                if i == 0:
+                    tmp_index = torch.arange(height * width)
+                else:
+                    tmp_index = torch.randperm(height * width)
+                index_list.append(tmp_index)
+            
+            res_index = torch.cat(index_list)
+            rows = res_index // width
+            cols = res_index % width
+            tar_all = tar_vec[batch, :, rows, cols].view(channels, -1, self.patch_height, self.patch_width * self.repeat_time)
+            src_all = src_vec[batch, :, rows, cols].view(channels, -1, self.patch_height, self.patch_width * self.repeat_time)
+            tar_mag = torch.clip(tar_all, 0, 1)*255
+            src_mag = torch.clip(src_all, 0, 1)*255
+
+            loss += (1 - self.ssim_loss(src_mag, tar_mag))
+
+        loss /= batch_size
+        return loss

MSGlance/Local.py

@@ -0,0 +1,121 @@
+'''SSIM in PyTorch.
+
+The source code is adopted from:
+https://github.com/Po-Hsun-Su/pytorch-ssim
+
+
+Reference:
+[1] Wang Z, Bovik A C, Sheikh H R, et al.
+    Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing
+'''
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+import math
+
+
+# def gaussian(window_size, sigma):
+#     gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+#     return gauss/gauss.sum()
+
+def uniform(window_size,sigma):
+    uniform_tensor = torch.ones(window_size)
+    return uniform_tensor / uniform_tensor.sum()
+
+
+def create_window(window_size, channel, sigma=1.5):
+    _1D_window = uniform(window_size, sigma).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+def _ssim(img1, img2, window, window_size, channel, size_average = True, stride=None, drop=None):
+    mu1 = F.conv2d(img1, window, padding = (window_size-1)//2, groups = channel, stride=stride)
+    mu2 = F.conv2d(img2, window, padding = (window_size-1)//2, groups = channel, stride=stride)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1*mu2
+
+    sigma1_sq = F.conv2d(img1*img1, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_sq
+    sigma2_sq = F.conv2d(img2*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu2_sq
+    sigma12 = F.conv2d(img1*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_mu2
+
+    C1 = 0.01**2
+    C2 = 0.03**2
+    C3 = C2/2
+
+    L = (2*mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)
+    C = (2*torch.sqrt(sigma1_sq)*torch.sqrt(sigma2_sq) + C2) / (sigma1_sq + sigma2_sq + C2)
+    S = (sigma12 + C3) / (torch.sqrt(sigma1_sq)*torch.sqrt(sigma2_sq) + C3)
+
+    if drop == "L":
+        ssim_map = C*S
+    elif drop == "C":
+        ssim_map = L*S
+    elif drop == "S":
+        ssim_map = L*C
+    elif drop == "LC":
+        ssim_map = S
+    else:
+        ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+class LocalGlance(torch.nn.Module):
+    def __init__(self, window_size = 3, size_average = True, stride=3, drop=None, sigma=1.5, channel=1):
+        super(LocalGlance, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = channel
+        self.stride = stride
+        self.window = create_window(window_size, self.channel, sigma)
+        self.drop = drop
+        self.sigma = sigma
+
+
+    def forward(self, img1, img2):
+        """
+        img1, img2: torch.Tensor([b,2,h,w]) - 2表示复数的实部和虚部
+        """
+        # 计算幅值
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel, self.sigma)
+
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+
+            self.window = window
+            self.channel = channel
+
+        return _ssim(img1, img2, window, self.window_size, channel, self.size_average, stride=self.stride, drop=self.drop)       
+
+
+def ssim(img1, img2, window_size = 11, size_average = True, sigma=1.5):
+    (_, channel, _, _) = img1.size()
+    window = create_window(window_size, channel, sigma)
+
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+
+    return _ssim(img1, img2, window, window_size, channel, size_average)
+
+
+
+
+
+
+
+

README.md

@@ -1 +1,15 @@
-
+# [MS-Glance] WACV25 submission 1940
+
+This repo contains the code of MSGlance.
+
+## Content
+
+The source code is located in the MSGlance folder.
+
+We provide an easy-to-use Jupyter notebook, `explore_MSGlance.ipynb`, which demonstrates how to use MS-Glance for INR fitting with SIREN. Additionally, it includes the script used to generate Figure 1 in the Supplementary Materials.
+
+## Usage
+
+You can simply upload the files to your Google Drive and run them with Colab, which comes with all the necessary pre-installed packages to execute the `explore_MSGlance.ipynb` notebook.
+
+I’d be happy to assist you with this, but doing so would reveal my username.