dump shiet

.gitignore

@@ -0,0 +1 @@
+weights/best_weight.pth

Dockerfile

@@ -0,0 +1,21 @@
+FROM python:3.9
+
+
+
+USER root
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN apt-get update
+RUN apt-get install ffmpeg libsm6 libxext6  -y
+
+RUN pip install --upgrade pip
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"] 
+

MeasureV1.py

@@ -0,0 +1,123 @@
+import glob
+import os
+import time
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import cv2
+import argparse
+
+from natsort import natsort
+from skimage.metrics import structural_similarity as compare_ssim
+from skimage.metrics import peak_signal_noise_ratio as compare_psnr
+import lpips
+
+
+class Measure():
+    def __init__(self, net='alex', use_gpu=False):
+        self.device = 'cuda' if use_gpu else 'cpu'
+        self.model = lpips.LPIPS(net=net)
+        self.model.to(self.device)
+
+    def measure(self, imgA, imgB):
+        if not all([s1 == s2 for s1, s2 in zip(imgA.shape, imgB.shape)]):
+            raise RuntimeError("Image sizes not the same.")
+        return [float(f(imgA, imgB)) for f in [self.psnr, self.ssim, self.lpips]]
+
+    def lpips(self, imgA, imgB, model=None):
+        tA = t(imgA).to(self.device)
+        tB = t(imgB).to(self.device)
+        dist01 = self.model.forward(tA, tB).item()
+        return dist01
+
+    def ssim(self, imgA, imgB):
+        # multichannel: If True, treat the last dimension of the array as channels. Similarity calculations are done independently for each channel then averaged.
+        score, diff = compare_ssim(imgA, imgB, full=True, multichannel=True)
+        return score
+
+    def psnr(self, imgA, imgB):
+        psnr = compare_psnr(imgA, imgB)
+        return psnr
+
+
+def t(img):
+    def to_4d(img):
+        assert len(img.shape) == 3
+        assert img.dtype == np.uint8
+        img_new = np.expand_dims(img, axis=0)
+        assert len(img_new.shape) == 4
+        return img_new
+
+    def to_CHW(img):
+        return np.transpose(img, [2, 0, 1])
+
+    def to_tensor(img):
+        return torch.Tensor(img)
+
+    return to_tensor(to_4d(to_CHW(img))) / 127.5 - 1
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+def imread(path):
+    return cv2.imread(path)[:, :, [2, 1, 0]]
+
+
+def format_result(psnr, ssim, lpips):
+    return f'{psnr:0.2f}, {ssim:0.3f}, {lpips:0.3f}'
+
+def measure_dirs(dirA, dirB, use_gpu, verbose=False):
+    if verbose:
+        vprint = lambda x: print(x)
+    else:
+        vprint = lambda x: None
+
+
+    t_init = time.time()
+
+    paths_A = fiFindByWildcard(os.path.join(dirA, f'*.{type}'))
+    paths_B = fiFindByWildcard(os.path.join(dirB, f'*.{type}'))
+
+    vprint("Comparing: ")
+    vprint(dirA)
+    vprint(dirB)
+
+    measure = Measure(use_gpu=use_gpu)
+
+    results = []
+    for pathA, pathB in zip(paths_A, paths_B):
+        result = OrderedDict()
+
+        t = time.time()
+        result['psnr'], result['ssim'], result['lpips'] = measure.measure(imread(pathA), imread(pathB))
+        d = time.time() - t
+        vprint(f"{pathA.split('/')[-1]}, {pathB.split('/')[-1]}, {format_result(**result)}, {d:0.1f}")
+
+        results.append(result)
+
+    psnr = np.mean([result['psnr'] for result in results])
+    ssim = np.mean([result['ssim'] for result in results])
+    lpips = np.mean([result['lpips'] for result in results])
+
+    vprint(f"Final Result: {format_result(psnr, ssim, lpips)}, {time.time() - t_init:0.1f}s")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-dirA', default='', type=str)
+    parser.add_argument('-dirB', default='', type=str)
+    parser.add_argument('-type', default='png')
+    parser.add_argument('--use_gpu', action='store_true', default=False)
+    args = parser.parse_args()
+
+    dirA = args.dirA
+    dirB = args.dirB
+    type = args.type
+    use_gpu = args.use_gpu
+
+    if len(dirA) > 0 and len(dirB) > 0:
+        measure_dirs(dirA, dirB, use_gpu=use_gpu, verbose=True)
+

README.md

@@ -1,11 +1,9 @@
 ---
 title: USR DA
-emoji: 📊
+emoji: 😻
 colorFrom: gray
-colorTo: red
-sdk: gradio
-sdk_version: 3.14.0
-app_file: app.py
+colorTo: indigo
+sdk: docker
 pinned: false
 ---
 

app.py

@@ -0,0 +1,99 @@
+import glob
+import io
+import os
+
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn as nn
+import wget
+from torchvision.transforms import Compose, ToTensor
+
+from model import decoder, encoder
+
+WEIGHT_PATH = './weights/best_weight.pth'
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Model(object):
+    def __init__(self) -> None:
+        self.model_Enc = encoder.Encoder_RRDB(num_feat=64).to(device=DEVICE)
+        self.model_Dec_SR = decoder.Decoder_SR_RRDB(num_in_ch=64).to(device=DEVICE)
+        self.preprocess = Compose([ToTensor()])
+        self.load_model()
+
+    def load_model(self, weight_path=WEIGHT_PATH):
+        if not os.path.isfile("./weights/best_weight.pth"):
+            response = wget.download("https://raw.githubusercontent.com/hungnguyen2611/super-resolution/master/weights/best_weight.pth", "./weights/best_weight.pth")
+        weight = torch.load(weight_path)
+        print("[LOADING] Loading encoder...")
+        self.model_Enc.load_state_dict(weight['model_Enc'])
+        print("[LOADING] Loading decoder...")
+        self.model_Dec_SR.load_state_dict(weight['model_Dec_SR'])
+        print("[LOADING] Loading done!")
+        self.model_Enc.eval()
+        self.model_Dec_SR.eval()
+
+    def predict(self, img):
+        with torch.no_grad():
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = self.preprocess(img)
+            img = img.unsqueeze(0)
+            img = img.to(DEVICE)
+
+            feat = self.model_Enc(img)
+            out = self.model_Dec_SR(feat)
+            min_max = (0, 1)
+            out = out.detach()[0].float().cpu()
+
+            out = out.squeeze().float().cpu().clamp_(*min_max)
+            out = (out - min_max[0]) / (min_max[1] - min_max[0])
+            out = out.numpy()
+            out = np.transpose(out[[2, 1, 0], :, :], (1, 2, 0))
+
+            out = (out*255.0).round()
+            out = out.astype(np.uint8)
+            return out
+
+model = Model()
+
+def predict(img):
+    global model
+    img.save("test/1.png", "PNG")
+    image = cv2.imread("test/1.png", cv2.IMREAD_COLOR)
+    out = model.predict(img=image)
+
+    cv2.imwrite(f'images_uploaded/1.png', out)
+    return f"images_uploaded/1.png"
+
+
+
+
+if __name__ == '__main__':
+    title = "Super-Resolution Demo USR-DA Unofficial 🚀🚀🔥"
+    description = ''' 
+<br>
+**This Demo expects low-quality and low-resolution images**
+**We are looking for collaborators! Collaborator** 
+</br>
+'''
+    article = "<p style='text-align: center'><a href='https://openaccess.thecvf.com/content/ICCV2021/papers/Wang_Unsupervised_Real-World_Super-Resolution_A_Domain_Adaptation_Perspective_ICCV_2021_paper.pdf' target='_blank'>Unsupervised Real-World Super-Resolution: A Domain Adaptation Perspective</a> | <a href='https://github.com/hungnguyen2611/super-resolution.git' target='_blank'>Github Repo</a></p>"
+    examples= glob.glob("testsets/*.png")
+    gr.Interface(
+        predict, 
+        gr.inputs.Image(type="pil", label="Input").style(height=260),
+        gr.inputs.Image(type="pil", label="Ouput").style(height=240),
+        title=title,
+        description=description,
+        article=article,
+        examples=examples,
+        ).launch(enable_queue=True)
+
+
+
+
+    
+
+
+

compare.py

@@ -0,0 +1,47 @@
+import glob
+import os
+import time
+from collections import OrderedDict
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from natsort import natsort
+from tqdm import tqdm
+
+
+def fiFindByWildcard(wildcard):
+    return natsort.natsorted(glob.glob(wildcard, recursive=True))
+
+
+
+
+if __name__ == "__main__":
+    out_data_path = fiFindByWildcard("./results_crop (1)/out/*")
+    gt_data_path = fiFindByWildcard("./results_crop (1)/target/*")
+    source_data_path = fiFindByWildcard("./results_crop (1)/source/*")
+
+    for src_path, out_path, gt_path in tqdm(list(zip(source_data_path, out_data_path, gt_data_path))):
+        fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
+        ax1.set_title("Bicubic")
+        ax2.set_title("Baseline")
+        ax3.set_title("Ground truth")
+        
+        src = cv2.imread(src_path)[:, :, [2, 1, 0]]
+        out = cv2.imread(out_path)[:, :, [2, 1, 0]]
+        gt = cv2.imread(gt_path)[:, :, [2, 1, 0]]
+        src = cv2.resize(src, None, fx=4, fy=4, interpolation=cv2.INTER_CUBIC)
+        ax1.set_yticklabels([])
+        ax1.set_xticklabels([])
+        ax2.set_yticklabels([])
+        ax2.set_xticklabels([])
+        ax3.set_yticklabels([])
+        ax3.set_xticklabels([])
+        ax1.imshow(src)
+        ax2.imshow(out)
+        ax3.imshow(gt)
+        fig.savefig(f"./result_compare_crop_new/{os.path.basename(gt_path)}", bbox_inches='tight' , dpi=1200)
+        plt.close()
+
+
+
+

crop_test.py

@@ -0,0 +1,5 @@
+import cv2
+
+
+
+

flagged/log.csv

@@ -0,0 +1,2 @@
+Input,Ouput,flag,username,timestamp
+/home/ds/Documents/SR/USR_DA/USR-DA/USR-DA/flagged/Input/tmplp_isgr5.jpg,,,,2022-12-18 14:45:29.533016

inference.py

@@ -0,0 +1,97 @@
+import os
+import numpy as np
+import cv2
+import torch.nn as nn
+from tqdm import tqdm
+import torch
+import torchvision
+
+from model import encoder, decoder
+from opt.option import args
+
+
+# device setting
+if args.gpu_id is not None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+    print('using GPU 0')
+else:
+    print('use --gpu_id to specify GPU ID to use')
+    exit()
+
+
+# make directory for saving weights
+if not os.path.exists(args.results):
+    os.mkdir(args.results)
+
+
+# numpy array -> torch tensor
+class ToTensor(object):
+    def __call__(self, sample):
+        sample = np.transpose(sample, (2, 0, 1))
+        sample = torch.from_numpy(sample)
+        return sample
+
+
+# create model
+# model_Enc = encoder.Encoder().cuda()
+# model_Dec_SR = decoder.Decoder_SR().cuda()
+model_Enc = encoder.Encoder_RRDB(num_feat=args.n_hidden_feats).cuda()
+model_Dec_SR = decoder.Decoder_SR_RRDB(num_in_ch=args.n_hidden_feats).cuda()
+
+model_Enc = nn.DataParallel(model_Enc)
+#model_Dec_Id = nn.DataParallel(model_Dec_Id)
+model_Dec_SR = nn.DataParallel(model_Dec_SR)
+
+# load weights
+checkpoint = torch.load(args.weights)
+model_Enc.load_state_dict(checkpoint['model_Enc'])
+model_Dec_SR.load_state_dict(checkpoint['model_Dec_SR'])
+model_Enc.eval()
+model_Dec_SR.eval()
+
+# input transform
+transforms = torchvision.transforms.Compose([ToTensor()])
+
+
+filenames = os.listdir(args.dir_test)
+filenames.sort()
+with torch.no_grad():
+    for filename in tqdm(filenames):
+        img_name = os.path.join(args.dir_test, filename)
+        ext = os.path.splitext(img_name)[-1]
+        if ext in ['.png', '.jpg']:
+            img = cv2.imread(img_name)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            #img = cv2.resize(img, ((img.shape[1] // 4),(img.shape[0] // 4)))
+            img = np.array(img).astype('float32') / 255
+            # img = img[0:256, 0:256, :]
+            
+            img = transforms(img)
+            img = torch.tensor(img.cuda()).unsqueeze(0)
+
+            # inference output
+            feat = model_Enc(img)
+            out = model_Dec_SR(feat)
+
+            min_max = (0, 1)
+            out = out.detach()[0].float().cpu()
+
+            out = out.squeeze().float().cpu().clamp_(*min_max)
+            out = (out - min_max[0]) / (min_max[1] - min_max[0])
+            out = out.numpy()
+            out = np.transpose(out[[2, 1, 0], :, :], (1, 2, 0))
+
+            out = (out*255.0).round()
+            out = out.astype(np.uint8)
+
+            # result image save (b x c x h x w (torch tensor) -> h x w x c (numpy array))
+            # out = out.data.cpu().squeeze().numpy()
+            # out = np.clip(out, 0, 1)
+            # out = np.transpose(out, (1, 2, 0))
+            print(args.results, filename)
+            cv2.imwrite('%s_out.png' %(os.path.join(args.results, filename)[:-4]), out)
+
+
+
+
+

model/decoder.py

@@ -0,0 +1,240 @@
+import torch
+import torch.nn as nn
+from torch.nn import init as init
+import torch.nn.functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
+
+class Decoder_Identity(nn.Module):
+    def __init__(self):
+        super(Decoder_Identity, self).__init__()
+
+        self.conv_up_2 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU()
+        )
+
+        self.conv_up_1 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU()
+        )
+
+        self.conv_last = nn.Sequential(
+            nn.Conv2d(in_channels=16, out_channels=3, kernel_size=1, bias=True),
+            nn.ReLU()
+        )
+
+    def forward(self, feat):
+        featmap_2 = self.conv_up_2(feat)
+        featmap_1 = self.conv_up_1(featmap_2)
+        out = self.conv_last(featmap_1)
+
+        return out
+
+
+class Decoder_SR(nn.Module):
+    def __init__(self, scale=4):
+        super(Decoder_SR, self).__init__()
+
+        self.scale = scale
+
+        self.conv_up_2 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU()
+        )
+
+        self.conv_up_1 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU()
+        )
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # upsampling
+        self.upsample_1 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True)
+        self.upsample_2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True)
+
+        self.HR_conv = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True)
+        self.conv_last = nn.Conv2d(in_channels=16, out_channels=3, kernel_size=3, padding=1, bias=True)
+
+    def forward(self, feat):
+        featmap_2 = self.conv_up_2(feat)
+        featmap_1 = self.conv_up_1(featmap_2)
+
+        if self.scale == 4:
+            featmap = self.lrelu(self.upsample_1(F.interpolate(featmap_1, scale_factor=2, mode='nearest')))
+            featmap = self.lrelu(self.upsample_2(F.interpolate(featmap, scale_factor=2, mode='nearest')))
+        elif self.scale == 2:
+            featmap = self.lrelu(self.upsample_1(F.interpolate(featmap_1, scale_factor=2, mode='nearest')))
+
+
+        out = self.conv_last(self.lrelu(self.HR_conv(featmap)))
+
+        return out
+
+
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat=64, num_grow_ch=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat, num_grow_ch=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class Decoder_Id_RRDB(nn.Module):
+    def __init__(self, num_in_ch, num_out_ch=3, scale=4, num_feat=64, num_block=10, num_grow_ch=32):
+        super(Decoder_Id_RRDB, self).__init__()
+
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+    
+    def forward(self, x):
+
+        feat = self.conv_first(x)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out
+
+
+class Decoder_SR_RRDB(nn.Module):
+    def __init__(self, num_in_ch, num_out_ch=3, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
+        super(Decoder_SR_RRDB, self).__init__()
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+
+        feat = self.conv_first(x)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out

model/discriminator.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+
+class DiscriminatorVGG(nn.Module):
+    def __init__(self, in_ch=3, image_size=128, d=64):
+        super(DiscriminatorVGG, self).__init__()
+        self.feature_map_size = image_size // 32
+        self.d = d
+
+        self.features = nn.Sequential(
+            nn.Conv2d(in_ch, d, kernel_size=3, stride=1, padding=1),  # input is 3 x 128 x 128
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d, d, kernel_size=3, stride=2, padding=1, bias=False),  # state size. 64 x 64 x 64
+            nn.BatchNorm2d(d),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d, d*2, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(d*2),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*2, d*2, kernel_size=3, stride=2, padding=1, bias=False),  # state size. 128 x 32 x 32
+            nn.BatchNorm2d(d*2),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*2, d*4, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(d*4),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*4, d*4, kernel_size=3, stride=2, padding=1, bias=False),  # state size. 256 x 16 x 16
+            nn.BatchNorm2d(d*4),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*4, d*8, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(d*8),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*8, d*8, kernel_size=3, stride=2, padding=1, bias=False),  # state size. 512 x 8 x 8
+            nn.BatchNorm2d(d*8),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*8, d*8, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(d*8),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+
+            nn.Conv2d(d*8, d*8, kernel_size=3, stride=2, padding=1, bias=False),  # state size. 512 x 4 x 4
+            nn.BatchNorm2d(d*8),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Linear((self.d*8) * self.feature_map_size * self.feature_map_size, 100),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Linear(100, 1)
+        )
+    
+    def forward(self, x):
+        out = self.features(x)
+        out = torch.flatten(out, 1)
+        out = self.classifier(out)
+
+        return out
+
+
+class UNetDiscriminator(nn.Module):
+    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
+        super(UNetDiscriminator, self).__init__()
+        self.skip_connection = skip_connection
+        norm = spectral_norm
+        self.num_in_ch = num_in_ch
+        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
+
+        self.conv1 = norm(nn.Conv2d(num_feat, num_feat*2, kernel_size=4, stride=2, padding=1, bias=False))
+        self.conv2 = norm(nn.Conv2d(num_feat*2, num_feat*4, kernel_size=4, stride=2, padding=1, bias=False))
+        self.conv3 = norm(nn.Conv2d(num_feat*4, num_feat*8, kernel_size=4, stride=2, padding=1, bias=False))
+
+        # upsample
+        self.conv4 = norm(nn.Conv2d(num_feat*8, num_feat*4, kernel_size=3, stride=1, padding=1, bias=False))
+        self.conv5 = norm(nn.Conv2d(num_feat*4, num_feat*2, kernel_size=3, stride=1, padding=1, bias=False))
+        self.conv6 = norm(nn.Conv2d(num_feat*2, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
+
+        # extra
+        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
+        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1, bias=False))
+
+        self.conv9 = nn.Conv2d(num_feat, 1, kernel_size=3, stride=1, padding=1)
+
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+    
+    def forward(self, x):
+        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
+        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
+        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
+        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
+
+        # upsample
+        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear')
+        x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x4 = x4 + x2
+        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear')
+        x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x5 = x5 + x1
+        x5 = F.interpolate(x5, scale_factor=2, mode='bilinear')
+        x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
+
+        if self.skip_connection:
+            x6 = x6 + x0
+
+        # extra
+        out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
+        out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
+        out = self.conv9(out)
+        out = self.avg_pool(out)
+        out = torch.flatten(out, 1)
+        return out

model/encoder.py

@@ -0,0 +1,67 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()
+
+        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
+
+        self.conv_featmap_1 = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+        )
+
+        self.conv_featmap_2 = nn.Sequential(
+            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+        )
+
+        self.conv_featmap_3 = nn.Sequential(
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, bias=True),
+            nn.ReLU(),
+        )
+    
+    def forward(self, img):
+        featmap_1 = self.conv_featmap_1(img)
+        featmap_1_down = self.maxpool(featmap_1)
+
+        featmap_2 = self.conv_featmap_2(featmap_1_down)
+        featmap_2_down = self.maxpool(featmap_2)
+
+        featmap_3 = self.conv_featmap_3(featmap_2_down)
+        
+        return featmap_3
+
+
+class Encoder_RRDB(nn.Module):
+    def __init__(self, num_feat=16):
+        super(Encoder_RRDB, self).__init__()
+        self.conv_featmap = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=num_feat, kernel_size=3, padding=1, bias=True),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Conv2d(in_channels=num_feat, out_channels=num_feat, kernel_size=3, padding=1, bias=True),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Conv2d(in_channels=num_feat, out_channels=num_feat, kernel_size=3, padding=1, bias=True),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Conv2d(in_channels=num_feat, out_channels=num_feat, kernel_size=3, padding=1, bias=True),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Conv2d(in_channels=num_feat, out_channels=num_feat, kernel_size=3, padding=1, bias=True),
+        )
+    
+    def forward(self, img):
+        featmap = self.conv_featmap(img)
+        
+        return featmap

opt/option.py

@@ -0,0 +1,69 @@
+import argparse
+
+
+parser = argparse.ArgumentParser(description='BebyGAN')
+
+# Hardware specifications
+parser.add_argument('--gpu_id', type=str, default = "0", help='specify GPU ID to use')
+parser.add_argument('--num_workers', type=int, default=4)
+
+# Data specifications
+parser.add_argument('--dir_data', type=str, default='./dataset', help='dataset root directory')
+parser.add_argument('--scale', type=int, default=4, help='super resolution scale')
+parser.add_argument('--patch_size', type=int, default=64, help='LR patch size') # default = 128 (in the paper)
+
+# Train specifications
+parser.add_argument('--epochs', type=int, default=35000, help='total epochs')
+parser.add_argument('--batch_size', type=int, default=1, help='size of each batch') # default = 8 (in the paper)
+
+# Optimizer specificaions
+parser.add_argument('--lr_G', type=float, default=1e-4, help='initial learning rate of generator')
+parser.add_argument('--lr_D', type=float, default=1e-4, help='initial learning rate of discriminator')
+parser.add_argument('--beta1', type=float, default=0.9, help='ADAM beta1')
+parser.add_argument('--beta2', type=float, default=0.99, help='ADAM beta2')
+parser.add_argument('--weight_decay', type=float, default=0.0, help='weight decay')
+
+# Scheduler specifications
+parser.add_argument('--interval1', type=int, default=2.5e5, help='1st step size (iteration)')
+parser.add_argument('--interval2', type=int, default=3.5e5, help='2nd step size (iteration)')
+parser.add_argument('--interval3', type=int, default=4.5e5, help='3rd step size (iteration)')
+parser.add_argument('--interval4', type=int, default=5.5e5, help='4th step size (iteration)')
+parser.add_argument('--gamma_G', type=float, default=0.5, help='generator learning rate decay ratio')
+parser.add_argument('--gamma_D', type=float, default=0.5, help='discriminator learning rate decay ratio')
+
+# Train specificaions
+parser.add_argument('--snap_path', type=str, default='./weights', help='path to save model weights')
+parser.add_argument('--save_freq', type=str, default=5, help='save model frequency (epoch)')
+# Logger 
+parser.add_argument('--log_interval', type=int, default=20)
+# checkpoint
+parser.add_argument('--checkpoint', type=str, default=None, help='load checkpoint')
+# pretrained
+parser.add_argument('--pretrained', type=str, default=None)
+# Optimizer specifications 
+parser.add_argument('--lambda_align', type=float, default=0.01, help='L1 loss weight')
+parser.add_argument('--lambda_rec', type=float, default=1.0, help='back-projection loss weight')
+parser.add_argument('--lambda_res', type=float, default=1.0, help='perceptual loss weight')
+parser.add_argument('--lambda_sty', type=float, default=0.01, help='style loss weight')
+parser.add_argument('--lambda_idt', type=float, default=0.01, help='identity loss weight')
+parser.add_argument('--lambda_cyc', type=float, default=1, help='cycle loss weight')
+
+parser.add_argument('--lambda_percept', type=float, default=0.01, help='perceptual loss weight')
+parser.add_argument('--lambda_adv', type=float, default=0.01, help='adversarial loss weight')
+
+# generator & discriminator specifications
+parser.add_argument('--n_disc', type=int, default=1, help='number of iteration for discriminator update in one epoch')
+parser.add_argument('--n_gen', type=int, default=2, help='number of iteration for generator update in one epoch')
+
+# encoder & decoder specifications
+parser.add_argument('--n_hidden_feats', type=int, default=64, help='number of feature vectors in hidden layer')
+parser.add_argument('--n_sr_feats', type=int, default=64, help='number of feature vectors in RRDB layer')
+# eval spec
+parser.add_argument('--phase', type=str, default='train')
+
+# test specifications
+parser.add_argument('--weights', type=str, default = "/data4/anhdh4/SR2/USR_DA-main/weights/epoch_1660.pth",help='load weights for test')
+parser.add_argument('--dir_test', type=str, default = "/data4/anhdh4/SR2/NTIRE2020/valid_source_crop",help='directory of test images')
+parser.add_argument('--results', type=str, default='./results1660/', help='directory of test results')
+
+args = parser.parse_args()

requirements.txt

@@ -0,0 +1,10 @@
+fastapi
+opencv-python
+scipy
+wget
+scikit-image
+torch==1.13.0
+torchmetrics==0.11.0
+torchvision==0.14.0
+tqdm
+uvicorn

train.py

@@ -0,0 +1,474 @@
+import os
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+from torch.utils.data import DataLoader
+from torchvision import transforms
+from tqdm import tqdm
+
+from data.LQGT_dataset import LQGTDataset, LQGTValDataset
+from model import decoder, discriminator, encoder
+from opt.option import args
+from util.utils import (RandCrop, RandHorizontalFlip, RandRotate, ToTensor, RandCrop_pair,
+                        VGG19PerceptualLoss)
+
+from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+wandb.init(project='SR', config=args)
+
+
+
+# device setting
+if args.gpu_id is not None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+    print('using GPU 0')
+else:
+    print('use --gpu_id to specify GPU ID to use')
+    exit()
+
+device = torch.device('cuda')
+
+# make directory for saving weights
+if not os.path.exists(args.snap_path):
+    os.mkdir(args.snap_path)
+
+print("Loading dataset...")
+# load training dataset
+train_dataset = LQGTDataset(
+    db_path=args.dir_data,
+    transform=transforms.Compose([RandCrop(args.patch_size, args.scale), RandHorizontalFlip(), RandRotate(), ToTensor()])
+)
+
+val_dataset = LQGTValDataset(
+    db_path=args.dir_data,
+    transform=transforms.Compose([RandCrop_pair(args.patch_size, args.scale), ToTensor()])
+)
+
+train_loader = DataLoader(
+    train_dataset,
+    batch_size=args.batch_size,
+    num_workers=args.num_workers,
+    drop_last=True,
+    shuffle=True
+)
+
+val_loader = DataLoader(
+    val_dataset,
+    batch_size=args.batch_size,
+    num_workers=args.num_workers,
+    shuffle=False
+)
+
+
+print("Create model")
+model_Disc_feat = discriminator.DiscriminatorVGG(in_ch=args.n_hidden_feats, image_size=args.patch_size).to(device)
+model_Disc_img_LR = discriminator.DiscriminatorVGG(in_ch=3, image_size=args.patch_size).to(device)
+model_Disc_img_HR = discriminator.DiscriminatorVGG(in_ch=3, image_size=args.scale*args.patch_size).to(device)
+# define model (generator)
+model_Enc = encoder.Encoder_RRDB(num_feat=args.n_hidden_feats).to(device)
+model_Dec_Id = decoder.Decoder_Id_RRDB(num_in_ch=args.n_hidden_feats).to(device)
+model_Dec_SR = decoder.Decoder_SR_RRDB(num_in_ch=args.n_hidden_feats).to(device)
+
+# define model (discriminator)
+
+# model_Disc_feat = discriminator.UNetDiscriminator(num_in_ch=64).to(device)
+# model_Disc_img_LR = discriminator.UNetDiscriminator(num_in_ch=3).to(device)
+# model_Disc_img_HR = discriminator.UNetDiscriminator(num_in_ch=3).to(device)
+
+# wandb logging
+wandb.watch(model_Disc_feat)
+wandb.watch(model_Disc_img_LR)
+wandb.watch(model_Enc)
+wandb.watch(model_Dec_Id)
+wandb.watch(model_Dec_SR)
+
+
+print("Define Loss")
+# loss
+loss_L1 = nn.L1Loss().to(device)
+loss_MSE = nn.MSELoss().to(device)
+loss_adversarial = nn.BCEWithLogitsLoss().to(device)
+loss_percept = VGG19PerceptualLoss().to(device)
+
+
+print("Define Optimizer")
+# optimizer 
+params_G = list(model_Enc.parameters()) + list(model_Dec_Id.parameters()) + list(model_Dec_SR.parameters())
+optimizer_G = optim.Adam(
+    params_G,
+    lr=args.lr_G,
+    betas=(args.beta1, args.beta2),
+    weight_decay=args.weight_decay,
+    amsgrad=True
+)
+params_D = list(model_Disc_feat.parameters()) + list(model_Disc_img_LR.parameters()) + list(model_Disc_img_HR.parameters())
+optimizer_D = optim.Adam(
+    params_D,
+    lr=args.lr_D,
+    betas=(args.beta1, args.beta2),
+    weight_decay=args.weight_decay,
+    amsgrad=True
+)
+
+print("Define Scheduler")
+# Scheduler
+iter_indices = [args.interval1, args.interval2, args.interval3]
+scheduler_G = optim.lr_scheduler.MultiStepLR(
+    optimizer=optimizer_G,
+    milestones=iter_indices,
+    gamma=0.5
+)
+scheduler_D = optim.lr_scheduler.MultiStepLR(
+    optimizer=optimizer_D,
+    milestones=iter_indices,
+    gamma=0.5
+)
+
+# print("Data Parallel")
+model_Enc = nn.DataParallel(model_Enc)
+model_Dec_Id = nn.DataParallel(model_Dec_Id)
+model_Dec_SR = nn.DataParallel(model_Dec_SR)
+
+# define model (discriminator)
+#model_Disc_feat = nn.DataParallel(model_Disc_feat)
+#model_Disc_img_LR = nn.DataParallel(model_Disc_img_LR)
+#model_Disc_img_HR = nn.DataParallel(model_Disc_img_HR)
+
+print("Load model weight")
+# load model weights & optimzer % scheduler
+if args.checkpoint is not None:
+    checkpoint = torch.load(args.checkpoint)
+
+    model_Enc.load_state_dict(checkpoint['model_Enc'])
+    model_Dec_Id.load_state_dict(checkpoint['model_Dec_Id'])
+    model_Dec_SR.load_state_dict(checkpoint['model_Dec_SR'])
+    model_Disc_feat.load_state_dict(checkpoint['model_Disc_feat'])
+    model_Disc_img_LR.load_state_dict(checkpoint['model_Disc_img_LR'])
+    model_Disc_img_HR.load_state_dict(checkpoint['model_Disc_img_HR'])
+
+    optimizer_D.load_state_dict(checkpoint['optimizer_D'])
+    optimizer_G.load_state_dict(checkpoint['optimizer_G'])
+
+    scheduler_D.load_state_dict(checkpoint['scheduler_D'])
+    scheduler_G.load_state_dict(checkpoint['scheduler_G'])
+
+    start_epoch = checkpoint['epoch']
+else:
+    start_epoch = 0
+
+
+if args.pretrained is not None:
+    ckpt = torch.load(args.pretrained)
+    ckpt["params"]["conv_first.weight"] = ckpt["params"]["conv_first.weight"][:,0,:,:].expand(64,64,3,3)
+    model_Dec_SR.load_state_dict(ckpt["params"])
+    
+
+
+
+
+
+
+# model_Enc = model_Enc.to(device)
+# model_Dec_Id = model_Dec_Id.to(device)
+# model_Dec_SR = model_Dec_SR.to(device)
+
+# # define model (discriminator)
+# model_Disc_feat = model_Disc_feat.to(device)
+# model_Disc_img_LR = model_Disc_img_LR.to(device)
+# model_Disc_img_HR =model_Disc_img_HR.to(device)
+# training
+
+PSNR = PeakSignalNoiseRatio().to(device)
+SSIM = StructuralSimilarityIndexMeasure().to(device)
+LPIPS = LearnedPerceptualImagePatchSimilarity().to(device)
+
+if args.phase == "train":
+  for epoch in range(start_epoch, args.epochs):
+      # generator
+      model_Enc.train()
+      model_Dec_Id.train()
+      model_Dec_SR.train()
+
+      # discriminator
+      model_Disc_feat.train()
+      model_Disc_img_LR.train()
+      model_Disc_img_HR.train()
+      
+      running_loss_D_total = 0.0
+      running_loss_G_total = 0.0
+
+      running_loss_align = 0.0
+      running_loss_rec = 0.0
+      running_loss_res = 0.0
+      running_loss_sty = 0.0
+      running_loss_idt = 0.0
+      running_loss_cyc = 0.0
+
+      iter = 0    
+
+      for data in tqdm(train_loader):
+          iter += 1
+
+          ########################
+          #       data load      #
+          ########################
+          X_t, Y_s = data['img_LQ'], data['img_GT']
+
+          ds4 = nn.Upsample(scale_factor=1/args.scale, mode='bicubic')
+          X_s = ds4(Y_s)
+
+          X_t = X_t.cuda(non_blocking=True)
+          X_s = X_s.cuda(non_blocking=True)
+          Y_s = Y_s.cuda(non_blocking=True)
+
+          # real label and fake label
+          batch_size = X_t.size(0)
+          real_label = torch.full((batch_size, 1), 1, dtype=X_t.dtype).cuda(non_blocking=True)
+          fake_label = torch.full((batch_size, 1), 0, dtype=X_t.dtype).cuda(non_blocking=True)
+
+
+          ########################
+          # (1) Update D network #
+          ########################
+          model_Disc_feat.zero_grad()
+          model_Disc_img_LR.zero_grad()
+          model_Disc_img_HR.zero_grad()
+
+          for i in range(args.n_disc):
+              # generator output (feature domain)
+              F_t = model_Enc(X_t)
+              F_s = model_Enc(X_s)
+
+              # 1. feature aligment loss (discriminator)
+              # output of discriminator (feature domain) (b x c(=1) x h x w)
+              output_Disc_F_t = model_Disc_feat(F_t.detach())
+              output_Disc_F_s = model_Disc_feat(F_s.detach())
+              # discriminator loss (feature domain)
+              loss_Disc_F_t = loss_MSE(output_Disc_F_t, fake_label)
+              loss_Disc_F_s = loss_MSE(output_Disc_F_s, real_label)
+              loss_Disc_feat_align = (loss_Disc_F_t + loss_Disc_F_s) / 2
+
+              # 2. SR reconstruction loss (discriminator)
+              # generator output (image domain)
+              Y_s_s = model_Dec_SR(F_s)
+              # output of discriminator (image domain)
+              output_Disc_Y_s_s = model_Disc_img_HR(Y_s_s.detach())
+              output_Disc_Y_s = model_Disc_img_HR(Y_s)
+              # discriminator loss (image domain)
+              loss_Disc_Y_s_s = loss_MSE(output_Disc_Y_s_s, fake_label)
+              loss_Disc_Y_s = loss_MSE(output_Disc_Y_s, real_label)
+              loss_Disc_img_rec = (loss_Disc_Y_s_s + loss_Disc_Y_s) / 2
+
+              # 4. Target degradation style loss
+              # generator output (image domain)
+              X_s_t = model_Dec_Id(F_s)
+              # output of discriminator (image domain)
+              output_Disc_X_s_t = model_Disc_img_LR(X_s_t.detach())
+              output_Disc_X_t = model_Disc_img_LR(X_t)
+              # discriminator loss (image domain)
+              loss_Disc_X_s_t = loss_MSE(output_Disc_X_s_t, fake_label)
+              loss_Disc_X_t = loss_MSE(output_Disc_X_t, real_label)
+              loss_Disc_img_sty = (loss_Disc_X_s_t + loss_Disc_X_t) / 2
+
+              # 6. Cycle loss
+              # generator output (image domain)
+              Y_s_t_s = model_Dec_SR(model_Enc(model_Dec_Id(F_s)))
+              # output of discriminator (image domain)
+              output_Disc_Y_s_t_s = model_Disc_img_HR(Y_s_t_s.detach())
+              output_Disc_Y_s = model_Disc_img_HR(Y_s)
+              # discriminator loss (image domain)
+              loss_Disc_Y_s_t_s = loss_MSE(output_Disc_Y_s_t_s, fake_label)
+              loss_Disc_Y_s = loss_MSE(output_Disc_Y_s, real_label)
+              loss_Disc_img_cyc = (loss_Disc_Y_s_t_s + loss_Disc_Y_s) / 2
+
+              # discriminator weight update
+              loss_D_total = loss_Disc_feat_align + loss_Disc_img_rec + loss_Disc_img_sty + loss_Disc_img_cyc
+              loss_D_total.backward()
+              optimizer_D.step()
+
+
+
+          scheduler_D.step()
+
+
+          ########################
+          # (2) Update G network #
+          ########################
+          model_Enc.zero_grad()
+          model_Dec_Id.zero_grad()
+          model_Dec_SR.zero_grad()
+
+          for i in range(args.n_gen):
+              # generator output (feature domain)
+              F_t = model_Enc(X_t)
+              F_s = model_Enc(X_s)
+
+              # 1. feature alignment loss (generator)
+              # output of discriminator (feature domain)
+              output_Disc_F_t = model_Disc_feat(F_t)
+              output_Disc_F_s = model_Disc_feat(F_s)
+              # generator loss (feature domain)
+              loss_G_F_t = loss_MSE(output_Disc_F_t, (real_label + fake_label)/2)
+              loss_G_F_s = loss_MSE(output_Disc_F_s, (real_label + fake_label)/2)
+              L_align_E = loss_G_F_t + loss_G_F_s
+
+              # 2. SR reconstruction loss
+              # generator output (image domain)
+              Y_s_s = model_Dec_SR(F_s)
+              # output of discriminator (image domain)
+              output_Disc_Y_s_s = model_Disc_img_HR(Y_s_s)
+              # L1 loss
+              loss_L1_rec = loss_L1(Y_s.detach(), Y_s_s)
+              # perceptual loss
+              loss_percept_rec = loss_percept(Y_s.detach(), Y_s_s)
+              # adversatial loss
+              loss_G_Y_s_s = loss_MSE(output_Disc_Y_s_s, real_label)
+              L_rec_G_SR = loss_L1_rec + args.lambda_percept*loss_percept_rec + args.lambda_adv*loss_G_Y_s_s
+
+              # 3. Target LR restoration loss
+              X_t_t = model_Dec_Id(F_t)
+              L_res_G_t = loss_L1(X_t, X_t_t)
+
+              # 4. Target degredation style loss
+              # generator output (image domain)
+              X_s_t = model_Dec_Id(F_s)
+              # output of discriminator (img domain)
+              output_Disc_X_s_t = model_Disc_img_LR(X_s_t)
+              # generator loss (feature domain)
+              loss_G_X_s_t = loss_MSE(output_Disc_X_s_t, real_label)
+              L_sty_G_t = loss_G_X_s_t
+
+              # 5. Feature identity loss
+              F_s_tilda = model_Enc(model_Dec_Id(F_s))
+              L_idt_G_t = loss_L1(F_s, F_s_tilda)
+
+              # 6. Cycle loss
+              # generator output (image domain)
+              Y_s_t_s = model_Dec_SR(model_Enc(model_Dec_Id(F_s)))
+              # output of discriminator (image domain)
+              output_Disc_Y_s_t_s = model_Disc_img_HR(Y_s_t_s)
+              # L1 loss
+              loss_L1_cyc = loss_L1(Y_s.detach(), Y_s_t_s)
+              # perceptual loss
+              loss_percept_cyc = loss_percept(Y_s.detach(), Y_s_t_s)
+              # adversarial loss 
+              loss_Y_s_t_s = loss_MSE(output_Disc_Y_s_t_s, real_label)
+              L_cyc_G_t_G_SR = loss_L1_cyc + args.lambda_percept*loss_percept_cyc + args.lambda_adv*loss_Y_s_t_s
+
+              # generator weight update
+              loss_G_total = args.lambda_align*L_align_E + args.lambda_rec*L_rec_G_SR + args.lambda_res*L_res_G_t + args.lambda_sty*L_sty_G_t + args.lambda_idt*L_idt_G_t + args.lambda_cyc*L_cyc_G_t_G_SR
+              loss_G_total.backward()
+              optimizer_G.step()
+          scheduler_G.step()
+
+
+          ########################
+          #     compute loss     #
+          ########################
+          running_loss_D_total += loss_D_total.item()
+          running_loss_G_total += loss_G_total.item()
+
+          running_loss_align += L_align_E.item()
+          running_loss_rec += L_rec_G_SR.item()
+          running_loss_res += L_res_G_t.item()
+          running_loss_sty += L_sty_G_t.item()
+          running_loss_idt += L_idt_G_t.item()
+          running_loss_cyc += L_cyc_G_t_G_SR.item()
+          if iter % args.log_interval == 0:
+              wandb.log(
+                  {
+                      "loss_D_total_step": running_loss_D_total/iter,
+                      "loss_G_total_step": running_loss_G_total/iter,
+                      "loss_align_step": running_loss_align/iter,
+                      "loss_rec_step": running_loss_rec/iter,
+                      "loss_res_step": running_loss_res/iter,
+                      "loss_sty_step": running_loss_sty/iter,
+                      "loss_idt_step": running_loss_idt/iter,
+                      "loss_cyc_step": running_loss_cyc/iter,
+                  }
+              )
+      ### EVALUATE ###
+      total_PSNR = 0
+      total_SSIM = 0
+      total_LPIPS = 0
+      val_iter = 0
+      with torch.no_grad():
+          model_Enc.eval()
+          model_Dec_SR.eval()
+          for batch_idx, batch in enumerate(val_loader):
+              val_iter += 1
+              source = batch["img_LQ"].to(device)
+              target = batch["img_GT"].to(device)
+
+              feat = model_Enc(source)
+              out = model_Dec_SR(feat)
+
+              total_PSNR += PSNR(out, target)
+              total_SSIM += SSIM(out, target)
+              total_LPIPS += LPIPS(out, target)
+      
+      wandb.log(
+          {
+              "epoch": epoch,
+              "lr": optimizer_G.param_groups[0]['lr'],
+              "loss_D_total_epoch": running_loss_D_total/iter,
+              "loss_G_total_epoch": running_loss_G_total/iter,
+              "loss_align_epoch": running_loss_align/iter,
+              "loss_rec_epoch": running_loss_rec/iter,
+              "loss_res_epoch": running_loss_res/iter,
+              "loss_sty_epoch": running_loss_sty/iter,
+              "loss_idt_epoch": running_loss_idt/iter,
+              "loss_cyc_epoch": running_loss_cyc/iter,
+              "PSNR_val": total_PSNR/val_iter,
+              "SSIM_val": total_SSIM/val_iter,
+              "LPIPS_val": total_LPIPS/val_iter
+          }
+      )
+
+
+      if (epoch+1) % args.save_freq == 0:
+          weights_file_name = 'epoch_%d.pth' % (epoch+1)
+          weights_file = os.path.join(args.snap_path, weights_file_name)
+          torch.save({
+              'epoch': epoch,
+
+              'model_Enc': model_Enc.state_dict(),
+              'model_Dec_Id': model_Dec_Id.state_dict(),
+              'model_Dec_SR': model_Dec_SR.state_dict(),
+              'model_Disc_feat': model_Disc_feat.state_dict(),
+              'model_Disc_img_LR': model_Disc_img_LR.state_dict(),
+              'model_Disc_img_HR': model_Disc_img_HR.state_dict(),
+
+              'optimizer_D': optimizer_D.state_dict(),
+              'optimizer_G': optimizer_G.state_dict(),
+
+              'scheduler_D': scheduler_D.state_dict(),
+              'scheduler_G': scheduler_G.state_dict(),
+          }, weights_file)
+          print('save weights of epoch %d' % (epoch+1))
+else:
+  ### EVALUATE ###
+  total_PSNR = 0
+  total_SSIM = 0
+  total_LPIPS = 0
+  val_iter = 0
+  with torch.no_grad():
+      model_Enc.eval()
+      model_Dec_SR.eval()
+      for batch_idx, batch in enumerate(val_loader):
+          val_iter += 1
+          source = batch["img_LQ"].to(device)
+          target = batch["img_GT"].to(device)
+
+          feat = model_Enc(source)
+          out = model_Dec_SR(feat)
+
+          total_PSNR += PSNR(out, target)
+          total_SSIM += SSIM(out, target)
+          total_LPIPS += LPIPS(out, target)
+      print("PSNR_val: ", total_PSNR/val_iter)
+      print("SSIM_val: ", total_SSIM/val_iter)
+      print("LPIPS_val: ", total_LPIPS/val_iter)

util/utils.py

@@ -0,0 +1,133 @@
+import torch
+import torchvision
+
+import math
+import cv2
+import numpy as np
+from scipy.ndimage import rotate
+
+
+class RandCrop(object):
+    def __init__(self, crop_size, scale):
+        # if output size is tuple -> (height, width)
+        assert isinstance(crop_size, (int, tuple))
+        if isinstance(crop_size, int):
+            self.crop_size = (crop_size, crop_size)
+        else:
+            assert len(crop_size) == 2
+            self.crop_size = crop_size
+        
+        self.scale = scale
+
+    def __call__(self, sample):
+        # img_LQ: H x W x C (numpy array)
+        img_LQ, img_GT = sample['img_LQ'], sample['img_GT']
+
+        h, w, c = img_LQ.shape
+        new_h, new_w = self.crop_size
+        top = np.random.randint(0, h - new_h)
+        left = np.random.randint(0, w - new_w)
+        img_LQ_crop = img_LQ[top: top+new_h, left: left+new_w, :]
+
+        h, w, c = img_GT.shape
+        top = np.random.randint(0, h - self.scale*new_h)
+        left = np.random.randint(0, w - self.scale*new_w)
+        img_GT_crop = img_GT[top: top + self.scale*new_h, left: left + self.scale*new_w, :]
+
+        sample = {'img_LQ': img_LQ_crop, 'img_GT': img_GT_crop}
+        return sample
+
+
+class RandRotate(object):
+    def __call__(self, sample):
+        # img_LQ: H x W x C (numpy array)
+        img_LQ, img_GT = sample['img_LQ'], sample['img_GT']
+
+        prob_rotate = np.random.random()
+        if prob_rotate < 0.25:
+            img_LQ = rotate(img_LQ, 90).copy()
+            img_GT = rotate(img_GT, 90).copy()
+        elif prob_rotate < 0.5:
+            img_LQ = rotate(img_LQ, 90).copy()
+            img_GT = rotate(img_GT, 90).copy()
+        elif prob_rotate < 0.75:
+            img_LQ = rotate(img_LQ, 90).copy()
+            img_GT = rotate(img_GT, 90).copy()
+        
+        sample = {'img_LQ': img_LQ, 'img_GT': img_GT}
+        return sample
+
+
+class RandHorizontalFlip(object):
+    def __call__(self, sample):
+        # img_LQ: H x W x C (numpy array)
+        img_LQ, img_GT = sample['img_LQ'], sample['img_GT']
+
+        prob_lr = np.random.random()
+        if prob_lr < 0.5:
+            img_LQ = np.fliplr(img_LQ).copy()
+            img_GT = np.fliplr(img_GT).copy()
+        
+        sample = {'img_LQ': img_LQ, 'img_GT': img_GT}
+        return sample
+
+
+class ToTensor(object):
+    def __call__(self, sample):
+        # img_LQ : H x W x C (numpy array) -> C x H x W (torch tensor)
+        img_LQ, img_GT = sample['img_LQ'], sample['img_GT']
+
+        img_LQ = img_LQ.transpose((2, 0, 1))
+        img_GT = img_GT.transpose((2, 0, 1))
+
+        img_LQ = torch.from_numpy(img_LQ)
+        img_GT = torch.from_numpy(img_GT)
+
+        sample = {'img_LQ': img_LQ, 'img_GT': img_GT}
+        return sample
+
+
+class VGG19PerceptualLoss(torch.nn.Module):
+    def __init__(self, feature_layer=35):
+        super(VGG19PerceptualLoss, self).__init__()
+        model = torchvision.models.vgg19(weights=torchvision.models.VGG19_Weights.DEFAULT)
+        self.features = torch.nn.Sequential(*list(model.features.children())[:feature_layer]).eval()
+        # Freeze parameters
+        for name, param in self.features.named_parameters():
+            param.requires_grad = False
+    
+    def forward(self, source, target):
+        vgg_loss = torch.nn.functional.l1_loss(self.features(source), self.features(target))
+
+        return vgg_loss
+        
+
+class RandCrop_pair(object):
+    def __init__(self, crop_size, scale):
+        # if output size is tuple -> (height, width)
+        assert isinstance(crop_size, (int, tuple))
+        if isinstance(crop_size, int):
+            self.crop_size = (crop_size, crop_size)
+        else:
+            assert len(crop_size) == 2
+            self.crop_size = crop_size
+        
+        self.scale = scale
+
+    def __call__(self, sample):
+        # img_LQ: H x W x C (numpy array)
+        img_LQ, img_GT = sample['img_LQ'], sample['img_GT']
+
+        h, w, c = img_LQ.shape
+        new_h, new_w = self.crop_size
+        top = np.random.randint(0, h - new_h)
+        left = np.random.randint(0, w - new_w)
+        img_LQ_crop = img_LQ[top: top+new_h, left: left+new_w, :]
+
+        h, w, c = img_GT.shape
+        top = self.scale*top
+        left = self.scale*left
+        img_GT_crop = img_GT[top: top + self.scale*new_h, left: left + self.scale*new_w, :]
+
+        sample = {'img_LQ': img_LQ_crop, 'img_GT': img_GT_crop}
+        return sample