init

.gitignore

@@ -0,0 +1,12 @@
+datasets
+train_lines.txt
+__pycache__
+.vscode
+flagged
+gradio_queue.db
+runs
+model_data/ResNeSt.pt
+model_data/G_FFHQ.pth
+dataset/FFHQ
+dataset/MFFHQ
+dataset/MFFHQ.7z

app.py

@@ -0,0 +1,38 @@
+'''
+Author: Egrt
+Date: 2022-01-13 13:34:10
+LastEditors: [egrt]
+LastEditTime: 2022-05-04 12:26:53
+FilePath: \MaskGAN\app.py
+'''
+
+from PIL import Image
+from maskgan import MASKGAN
+import gradio as gr
+import os
+maskgan = MASKGAN()
+
+# --------模型推理---------- #
+def inference(img):
+    lr_shape = [112, 112]
+    img = img.resize(tuple(lr_shape), Image.BICUBIC)
+    r_image = maskgan.generate_1x1_image(img)
+    return r_image
+
+# --------网页信息---------- #  
+title = "MaskGAN:融合无监督的口罩遮挡人脸修复"
+description = "使用生成对抗网络对口罩遮挡人脸进行修复,能够有效的恢复被遮挡区域人脸。  @西南科技大学智能控制与图像处理研究室"
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.10257' target='_blank'>MaskGAN: Face Restoration Using Swin Transformer</a> | <a href='https://github.com/JingyunLiang/SwinIR' target='_blank'>Github Repo</a></p>"
+example_img_dir  = 'img'
+example_img_name = os.listdir(example_img_dir)
+examples=[[os.path.join(example_img_dir, image_path)] for image_path in example_img_name if image_path.endswith('.jpg')]
+gr.Interface(
+    inference, 
+    [gr.inputs.Image(type="pil", label="Input")], 
+    gr.outputs.Image(type="pil", label="Output"),
+    title=title,
+    description=description,
+    article=article,
+    enable_queue=True,
+    examples=examples
+    ).launch(debug=True)

img/README.md

@@ -0,0 +1,8 @@
+<!--
+ * @Author: Egrt
+ * @Date: 2022-04-09 11:02:07
+ * @LastEditors: Egrt
+ * @LastEditTime: 2022-04-09 11:02:07
+ * @FilePath: \MaskGAN\img\README.md
+-->
+存放示例照片

maskgan.py

@@ -0,0 +1,91 @@
+'''
+Author: Egrt
+Date: 2022-04-07 14:00:52
+LastEditors: [egrt]
+LastEditTime: 2022-05-04 11:47:21
+FilePath: \MaskGAN\maskgan.py
+'''
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from PIL import Image
+from models.SwinIR import Generator
+from utils.utils import cvtColor, preprocess_input
+
+
+class MASKGAN(object):
+    #-----------------------------------------#
+    #   注意修改model_path
+    #-----------------------------------------#
+    _defaults = {
+        #-----------------------------------------------#
+        #   model_path指向logs文件夹下的权值文件
+        #-----------------------------------------------#
+        "model_path"        : 'model_data/G_FFHQ.pth',
+        #-----------------------------------------------#
+        #   上采样的倍数，和训练时一样
+        #-----------------------------------------------#
+        "scale_factor"      : 1, 
+        #-----------------------------------------------#
+        #   hr_shape
+        #-----------------------------------------------#
+        "hr_shape"          : [112, 112], 
+        #-------------------------------#
+        #   是否使用Cuda
+        #   没有GPU可以设置成False
+        #-------------------------------#
+        "cuda"              : False,
+    }
+
+    #---------------------------------------------------#
+    #   初始化MASKGAN
+    #---------------------------------------------------#
+    def __init__(self, **kwargs):
+        self.__dict__.update(self._defaults)
+        for name, value in kwargs.items():
+            setattr(self, name, value)  
+        self.generate()
+
+    def generate(self):
+        self.net    = Generator(upscale=self.scale_factor, img_size=tuple(self.hr_shape),
+                   window_size=7, img_range=1., depths=[6, 6, 6, 6],
+                   embed_dim=96, num_heads=[6, 6, 6, 6], mlp_ratio=4, upsampler='pixelshuffledirect')
+
+        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.net    = torch.load(self.model_path, map_location=device)
+        self.net    = self.net.eval()
+        print('{} model, and classes loaded.'.format(self.model_path))
+
+        if self.cuda:
+            self.net = torch.nn.DataParallel(self.net)
+            cudnn.benchmark = True
+            self.net = self.net.cuda()
+        
+    def generate_1x1_image(self, image):
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度，并进行归一化
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]), 0)
+        
+        with torch.no_grad():
+            image_data = torch.from_numpy(image_data).type(torch.FloatTensor)
+            if self.cuda:
+                image_data = image_data.cuda()
+
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            hr_image = self.net(image_data)[0]
+            #---------------------------------------------------------#
+            #   将归一化的结果再转成rgb格式
+            #---------------------------------------------------------#
+            hr_image = (hr_image.cpu().data.numpy().transpose(1, 2, 0) * 0.5 + 0.5)
+            hr_image = np.clip(hr_image * 255, 0, 255)
+
+            hr_image = Image.fromarray(np.uint8(hr_image))
+            return hr_image

model_data/README.md

@@ -0,0 +1,8 @@
+<!--
+ * @Author: Egrt
+ * @Date: 2022-04-11 20:53:42
+ * @LastEditors: Egrt
+ * @LastEditTime: 2022-04-11 20:53:42
+ * @FilePath: \MaskGAN\model_data\README.md
+-->
+存放模型权重

utils/__init__.py

@@ -0,0 +1 @@
+from .utils import *

utils/dataloader.py

@@ -0,0 +1,101 @@
+from random import randint
+
+import cv2
+import numpy as np
+from PIL import Image
+from torch.utils.data.dataset import Dataset
+
+from .utils import cvtColor, preprocess_input
+
+def look_image(image_name, image):
+    image = np.array(image)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    cv2.imshow(image_name, image)
+    cv2.waitKey(0)
+
+
+def get_new_img_size(width, height, img_min_side=600):
+    if width <= height:
+        f = float(img_min_side) / width
+        resized_height = int(f * height)
+        resized_width = int(img_min_side)
+    else:
+        f = float(img_min_side) / height
+        resized_width = int(f * width)
+        resized_height = int(img_min_side)
+
+    return resized_width, resized_height
+
+class MASKGANDataset(Dataset):
+    def __init__(self, train_lines, lr_shape, hr_shape):
+        super(MASKGANDataset, self).__init__()
+
+        self.train_lines    = train_lines
+        self.train_batches  = len(train_lines)
+
+        self.lr_shape       = lr_shape
+        self.hr_shape       = hr_shape
+
+    def __len__(self):
+        return self.train_batches
+
+    def __getitem__(self, index):
+        index = index % self.train_batches
+        image_list = self.train_lines[index].split(' ')
+        image_origin = Image.open(image_list[0])
+        image_masked = Image.open(image_list[1].split()[0])
+
+        image_origin, image_masked = self.get_random_data(image_origin, image_masked, self.hr_shape)
+
+        image_origin = image_origin.resize((self.hr_shape[1], self.hr_shape[0]), Image.BICUBIC)
+        image_masked = image_masked.resize((self.lr_shape[1], self.lr_shape[0]), Image.BICUBIC)
+        # look_image('origin', image_origin)
+        # look_image('masked', image_masked)
+        image_origin = np.transpose(preprocess_input(np.array(image_origin, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
+        image_masked = np.transpose(preprocess_input(np.array(image_masked, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
+
+        return np.array(image_masked), np.array(image_origin)
+
+    def rand(self, a=0, b=1):
+        return np.random.rand()*(b-a) + a
+
+    def get_random_data(self, image_origin, image_masked, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
+        #------------------------------#
+        #   读取图像并转换成RGB图像
+        #------------------------------#
+        image_origin   = cvtColor(image_origin)
+        image_masked   = cvtColor(image_masked)
+
+        #------------------------------------------#
+        #   色域扭曲
+        #------------------------------------------#
+        hue = self.rand(-hue, hue)
+        sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat)
+        val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val)
+
+        x = cv2.cvtColor(np.array(image_origin,np.float32)/255, cv2.COLOR_RGB2HSV)
+        x[..., 1] *= sat
+        x[..., 2] *= val
+        x[x[:,:, 0]>360, 0] = 360
+        x[:, :, 1:][x[:, :, 1:]>1] = 1
+        x[x<0] = 0
+        image_data_origin = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
+
+        x = cv2.cvtColor(np.array(image_masked,np.float32)/255, cv2.COLOR_RGB2HSV)
+        x[..., 1] *= sat
+        x[..., 2] *= val
+        x[x[:,:, 0]>360, 0] = 360
+        x[:, :, 1:][x[:, :, 1:]>1] = 1
+        x[x<0] = 0
+        image_data_masked = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
+
+        return Image.fromarray(np.uint8(image_data_origin)), Image.fromarray(np.uint8(image_data_masked))
+
+        
+def MASKGAN_dataset_collate(batch):
+    images_l = []
+    images_h = []
+    for img_l, img_h in batch:
+        images_l.append(img_l)
+        images_h.append(img_h)
+    return np.array(images_l), np.array(images_h)

utils/utils.py

@@ -0,0 +1,164 @@
+import itertools
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from torch.nn import functional as F
+import cv2
+import distutils.util
+
+def show_result(num_epoch, G_net, imgs_lr, imgs_hr):
+    with torch.no_grad():
+        test_images = G_net(imgs_lr)
+
+        fig, ax = plt.subplots(1, 3)
+
+        for j in itertools.product(range(3)):
+            ax[j].get_xaxis().set_visible(False)
+            ax[j].get_yaxis().set_visible(False)
+        ax[0].cla()
+        ax[0].imshow(np.transpose(np.clip(imgs_lr.cpu().numpy()[0] * 0.5 + 0.5, 0, 1), [1,2,0]))
+
+        ax[1].cla()
+        ax[1].imshow(np.transpose(np.clip(test_images.cpu().numpy()[0] * 0.5 + 0.5, 0, 1), [1,2,0]))
+
+        ax[2].cla()
+        ax[2].imshow(np.transpose(np.clip(imgs_hr.cpu().numpy()[0] * 0.5 + 0.5, 0, 1), [1,2,0]))
+
+        label = 'Epoch {0}'.format(num_epoch)
+        fig.text(0.5, 0.04, label, ha='center')
+        plt.savefig("results/train_out/epoch_" + str(num_epoch) + "_results.png")
+        plt.close('all')  #避免内存泄漏
+
+#---------------------------------------------------------#
+#   将图像转换成RGB图像，防止灰度图在预测时报错。
+#   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+#---------------------------------------------------------#
+def cvtColor(image):
+    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
+        return image 
+    else:
+        image = image.convert('RGB')
+        return image 
+
+def preprocess_input(image, mean, std):
+    image = (image/255 - mean)/std
+    return image
+
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+def print_arguments(args):
+    print("-----------  Configuration Arguments -----------")
+    for arg, value in sorted(vars(args).items()):
+        print("%s: %s" % (arg, value))
+    print("------------------------------------------------")
+
+
+def add_arguments(argname, type, default, help, argparser, **kwargs):
+    type = distutils.util.strtobool if type == bool else type
+    argparser.add_argument("--" + argname,
+                           default=default,
+                           type=type,
+                           help=help + ' 默认: %(default)s.',
+                           **kwargs)
+
+def filter2D(img, kernel):
+    """PyTorch version of cv2.filter2D
+
+    Args:
+        img (Tensor): (b, c, h, w)
+        kernel (Tensor): (b, k, k)
+    """
+    k = kernel.size(-1)
+    b, c, h, w = img.size()
+    if k % 2 == 1:
+        img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
+    else:
+        raise ValueError('Wrong kernel size')
+
+    ph, pw = img.size()[-2:]
+
+    if kernel.size(0) == 1:
+        # apply the same kernel to all batch images
+        img = img.view(b * c, 1, ph, pw)
+        kernel = kernel.view(1, 1, k, k)
+        return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
+    else:
+        img = img.view(1, b * c, ph, pw)
+        kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
+        return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
+
+
+def usm_sharp(img, weight=0.5, radius=50, threshold=10):
+    """USM sharpening.
+
+    Input image: I; Blurry image: B.
+    1. sharp = I + weight * (I - B)
+    2. Mask = 1 if abs(I - B) > threshold, else: 0
+    3. Blur mask:
+    4. Out = Mask * sharp + (1 - Mask) * I
+
+
+    Args:
+        img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
+        weight (float): Sharp weight. Default: 1.
+        radius (float): Kernel size of Gaussian blur. Default: 50.
+        threshold (int):
+    """
+    if radius % 2 == 0:
+        radius += 1
+    blur = cv2.GaussianBlur(img, (radius, radius), 0)
+    residual = img - blur
+    mask = np.abs(residual) * 255 > threshold
+    mask = mask.astype('float32')
+    soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
+
+    sharp = img + weight * residual
+    sharp = np.clip(sharp, 0, 1)
+    return soft_mask * sharp + (1 - soft_mask) * img
+
+
+class USMSharp(torch.nn.Module):
+
+    def __init__(self, radius=50, sigma=0):
+        super(USMSharp, self).__init__()
+        if radius % 2 == 0:
+            radius += 1
+        self.radius = radius
+        kernel = cv2.getGaussianKernel(radius, sigma)
+        kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0)
+        self.register_buffer('kernel', kernel)
+
+    def forward(self, img, weight=0.5, threshold=10):
+        blur = filter2D(img, self.kernel)
+        residual = img - blur
+
+        mask = torch.abs(residual) * 255 > threshold
+        mask = mask.float()
+        soft_mask = filter2D(mask, self.kernel)
+        sharp = img + weight * residual
+        sharp = torch.clip(sharp, 0, 1)
+        return soft_mask * sharp + (1 - soft_mask) * img
+
+class USMSharp_npy():
+
+    def __init__(self, radius=50, sigma=0):
+        super(USMSharp_npy, self).__init__()
+        if radius % 2 == 0:
+            radius += 1
+        self.radius = radius
+        kernel = cv2.getGaussianKernel(radius, sigma)
+        self.kernel = np.dot(kernel, kernel.transpose()).astype(np.float32)
+
+    def filt(self, img, weight=0.5, threshold=10):
+        blur = cv2.filter2D(img, -1, self.kernel)
+        residual = img - blur
+
+        mask = np.abs(residual) * 255 > threshold
+        mask = mask.astype(np.float32)
+        soft_mask = cv2.filter2D(mask, -1, self.kernel)
+        sharp = img + weight * residual
+        sharp = np.clip(sharp, 0, 1)
+        return soft_mask * sharp + (1 - soft_mask) * img
+

utils/utils_fit.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn.functional as F
+from models.SwinIR import compute_gradient_penalty
+from tqdm import tqdm
+
+from .utils import get_lr, show_result
+from .utils_metrics import PSNR, SSIM
+
+
+
+def fit_one_epoch(writer, G_model_train, D_model_train, G_model, D_model, VGG_feature_model, ResNeSt_model, G_optimizer, D_optimizer, BCEWithLogits_loss, L1_loss, Face_loss, epoch, epoch_size, gen, Epoch, cuda, batch_size, save_interval):
+    G_total_loss = 0
+    D_total_loss = 0
+    G_total_PSNR = 0
+    G_total_SSIM = 0
+
+    with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3, ncols=150) as pbar:
+        for iteration, batch in enumerate(gen):
+            if iteration >= epoch_size:
+                break
+            
+            with torch.no_grad():
+                lr_images, hr_images    = batch
+                lr_images, hr_images    = torch.from_numpy(lr_images).type(torch.FloatTensor), torch.from_numpy(hr_images).type(torch.FloatTensor)
+                y_real, y_fake          = torch.ones(batch_size), torch.zeros(batch_size)
+                if cuda:
+                    lr_images, hr_images, y_real, y_fake  = lr_images.cuda(), hr_images.cuda(), y_real.cuda(), y_fake.cuda()
+            
+            #-------------------------------------------------#
+            #   训练判别器
+            #-------------------------------------------------#
+            D_optimizer.zero_grad()
+
+            D_result_r              = D_model_train(hr_images)
+
+            G_result                = G_model_train(lr_images)
+            D_result_f              = D_model_train(G_result).squeeze()
+            D_result_rf             = D_result_r - D_result_f.mean()
+            D_result_fr             = D_result_f - D_result_r.mean()
+            D_train_loss_rf         = BCEWithLogits_loss(D_result_rf, y_real)
+            D_train_loss_fr         = BCEWithLogits_loss(D_result_fr, y_fake)
+            gradient_penalty        = compute_gradient_penalty(D_model_train, hr_images, G_result)
+            D_train_loss            = 10 * gradient_penalty + (D_train_loss_rf + D_train_loss_fr) / 2
+            D_train_loss.backward()
+
+            D_optimizer.step()
+
+            #-------------------------------------------------#
+            #   训练生成器
+            #-------------------------------------------------#
+            G_optimizer.zero_grad()
+            
+            G_result                = G_model_train(lr_images)
+            image_loss              = L1_loss(G_result, hr_images)
+
+            D_result_r              = D_model_train(hr_images)
+            D_result_f              = D_model_train(G_result).squeeze()
+            D_result_rf             = D_result_r - D_result_f.mean()
+            D_result_fr             = D_result_f - D_result_r.mean()
+            D_train_loss_rf         = BCEWithLogits_loss(D_result_rf, y_fake)
+            D_train_loss_fr         = BCEWithLogits_loss(D_result_fr, y_real)
+            adversarial_loss        = (D_train_loss_rf + D_train_loss_fr) / 2
+
+            perception_loss         = L1_loss(VGG_feature_model(G_result), VGG_feature_model(hr_images))
+            # 进行下采样以适配人脸识别网络
+            G_result_face           = F.interpolate(G_result, size=(112, 112), mode='bicubic', align_corners=True)
+            hr_images_face          = F.interpolate(hr_images, size=(112, 112), mode='bicubic', align_corners=True)
+            face_loss               = torch.mean(1. - Face_loss(ResNeSt_model(G_result_face), ResNeSt_model(hr_images_face)))
+            G_train_loss            = 3.0 * image_loss + 1.0 * adversarial_loss + 0.9 * perception_loss + 2.5 * face_loss
+
+            G_train_loss.backward()
+            G_optimizer.step()
+            
+            G_total_loss            += G_train_loss.item()
+            D_total_loss            += D_train_loss.item()
+
+            with torch.no_grad():
+                G_total_PSNR        += PSNR(G_result, hr_images).item()
+                G_total_SSIM        += SSIM(G_result, hr_images).item()
+                
+            pbar.set_postfix(**{'G_loss'    : G_total_loss / (iteration + 1), 
+                                'D_loss'    : D_total_loss / (iteration + 1), 
+                                'G_PSNR'    : G_total_PSNR / (iteration + 1), 
+                                'G_SSIM'    : G_total_SSIM / (iteration + 1), 
+                                'lr'        : get_lr(G_optimizer)})
+            pbar.update(1)
+
+            if iteration % save_interval == 0:
+                show_result(epoch + 1, G_model_train, lr_images, hr_images)
+        writer.add_scalar('G_loss', G_total_loss / (iteration + 1), epoch + 1)
+        writer.add_scalar('D_loss', D_total_loss / (iteration + 1), epoch + 1)
+        writer.add_scalar('G_PSNR', G_total_PSNR / (iteration + 1), epoch + 1)
+        writer.add_scalar('G_SSIM', G_total_SSIM / (iteration + 1), epoch + 1)
+        writer.add_scalar('lr',     get_lr(G_optimizer),            epoch + 1)
+    print('Epoch:'+ str(epoch + 1) + '/' + str(Epoch))
+    print('G Loss: %.4f || D Loss: %.4f ' % (G_total_loss / epoch_size, D_total_loss / epoch_size))
+    print('Saving state, iter:', str(epoch+1))
+    # 保存模型权重
+    torch.save(G_model, 'logs/G_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size))
+    torch.save(D_model, 'logs/D_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size))

utils/utils_metrics.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn.functional as F
+from math import exp
+import numpy as np
+ 
+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 create_window(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).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 window
+
+def SSIM(img1, img2, window_size=11, window=None, size_average=True, full=False):
+    img1 = (img1 * 0.5 + 0.5) * 255
+    img2 = (img2 * 0.5 + 0.5) * 255
+    min_val = 0
+    max_val = 255
+    L = max_val - min_val
+    img2 = torch.clamp(img2, 0.0, 255.0)
+ 
+    padd = 0
+    (_, channel, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window(real_size, channel=channel).to(img1.device)
+ 
+    mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
+ 
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+ 
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2
+ 
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+ 
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+ 
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+ 
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+ 
+    if full:
+        return ret, cs
+    return ret
+ 
+def tf_log10(x):
+    numerator = torch.log(x)
+    denominator = torch.log(torch.tensor(10.0))
+    return numerator / denominator
+
+def PSNR(img1, img2):
+    img1 = (img1 * 0.5 + 0.5) * 255
+    img2 = (img2 * 0.5 + 0.5) * 255
+    max_pixel = 255.0
+    img2 = torch.clamp(img2, 0.0, 255.0)
+    return 10.0 * tf_log10((max_pixel ** 2) / (torch.mean(torch.pow(img2 - img1, 2))))