+
+ >
+
+>
+For more results, please refer to the [paper](https://arxiv.org/abs/2108.05301) and [supp](https://github.com/JingyunLiang/HCFlow/releases/tag/v0.0) for details.
+
+## Citation
+ @inproceedings{liang21hcflow,
+ title={Hierarchical Conditional Flow: A Unified Framework for Image Super-Resolution and Image Rescaling},
+ author={Liang, Jingyun and Lugmayr, Andreas and Zhang, Kai and Danelljan, Martin and Van Gool, Luc and Timofte, Radu},
+ booktitle={IEEE Conference on International Conference on Computer Vision},
+ year={2021}
+ }
+
+
+## License & Acknowledgement
+
+This project is released under the Apache 2.0 license. The codes are based on [MMSR](https://github.com/open-mmlab/mmediting), [SRFlow](https://github.com/andreas128/SRFlow), [IRN](https://github.com/pkuxmq/Invertible-Image-Rescaling) and [Glow-pytorch](https://github.com/chaiyujin/glow-pytorch). Please also follow their licenses. Thanks for their great works.
+
diff --git a/codes/data/GTLQ_dataset.py b/codes/data/GTLQ_dataset.py
new file mode 100644
index 0000000..ffea133
--- /dev/null
+++ b/codes/data/GTLQ_dataset.py
@@ -0,0 +1,129 @@
+import random
+import numpy as np
+import cv2
+import lmdb
+import torch
+import torch.utils.data as data
+import data.util as util
+import sys
+import os
+
+try:
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ from data.util import imresize_np
+ from utils import util as utils
+except ImportError:
+ pass
+
+
+class GTLQDataset(data.Dataset):
+ '''
+ Load HR-LR image pairs.
+ '''
+
+ def __init__(self, opt):
+ super(GTLQDataset, self).__init__()
+ self.opt = opt
+ self.LR_paths, self.GT_paths = None, None
+ self.LR_env, self.GT_env = None, None # environment for lmdb
+ self.scale = opt['scale']
+ if self.opt['phase'] == 'train':
+ self.GT_size = opt['GT_size']
+ self.LR_size = self.GT_size // self.scale
+
+ # read image list from lmdb or image files
+ if opt['data_type'] == 'lmdb':
+ self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
+ self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
+ elif opt['data_type'] == 'img':
+ self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
+ self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
+ else:
+ print('Error: data_type is not matched in Dataset')
+ assert self.GT_paths, 'Error: GT paths are empty.'
+ if self.LR_paths and self.GT_paths:
+ assert len(self.LR_paths) == len(
+ self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(
+ len(self.LR_paths), len(self.GT_paths))
+
+ def _init_lmdb(self):
+ # https://github.com/chainer/chainermn/issues/129
+ self.GT_env = lmdb.open(self.opt['dataroot_GT'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ if self.opt['dataroot_LQ'] is not None:
+ self.LR_env = lmdb.open(self.opt['dataroot_LQ'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ else:
+ self.LR_env = 'No lmdb input for LR'
+
+ def __getitem__(self, index):
+ if self.opt['data_type'] == 'lmdb':
+ if (self.GT_env is None) or (self.LR_env is None):
+ self._init_lmdb()
+
+ if self.opt['data_type'] == 'lmdb':
+ resolution = [int(s) for s in self.GT_sizes[index].split('_')]
+ else:
+ resolution = None
+
+
+ # loading code from srflow test
+ # img_GT = cv2.imread(GT_path)[:, :, [2, 1, 0]]
+ # img_GT = torch.Tensor(img_GT.transpose([2, 0, 1]).astype(np.float32)) / 255
+ # img_LR = cv2.imread(LR_path)[:, :, [2, 1, 0]]
+ # pad_factor = 2
+ # h, w, c = img_LR.shape
+ # img_LR = impad(img_LR, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
+ # right=int(np.ceil(w / pad_factor) * pad_factor - w))
+ # img_LR = torch.Tensor(img_LR.transpose([2, 0, 1]).astype(np.float32)) / 255
+
+
+ # get GT and LR image
+ GT_path = self.GT_paths[index]
+ LR_path = self.LR_paths[index]
+ # LR_path = GT_path.replace('HR', 'LR_bicubic/X4').replace('.png','x{}.png'.format(self.scale))
+ img_GT = util.read_img(self.GT_env, GT_path, resolution) # return: Numpy float32, HWC, BGR, [0,1]
+ img_LR = util.read_img(self.LR_env, LR_path, resolution)
+
+
+ if self.opt['phase'] == 'train':
+ # crop
+ H, W, C = img_LR.shape
+ rnd_top_LR = random.randint(0, max(0, H - self.LR_size))
+ rnd_left_LR = random.randint(0, max(0, W - self.LR_size))
+ rnd_top_GT = rnd_top_LR * self.scale
+ rnd_left_GT = rnd_left_LR * self.scale
+
+ img_GT = img_GT[rnd_top_GT:rnd_top_GT + self.GT_size, rnd_left_GT:rnd_left_GT + self.GT_size, :]
+ img_LR = img_LR[rnd_top_LR:rnd_top_LR + self.LR_size, rnd_left_LR:rnd_left_LR + self.LR_size, :]
+
+ # augmentation - flip, rotate
+ img_GT, img_LR = util.augment([img_GT, img_LR], self.opt['use_flip'],
+ self.opt['use_rot'], self.opt['mode'])
+
+ # change color space if necessary, deal with gray image
+ if self.opt['color']:
+ img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
+ img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0]
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ if img_GT.shape[2] == 3:
+ img_GT = img_GT[:, :, [2, 1, 0]]
+ if img_LR.shape[2] == 3:
+ img_LR = img_LR[:, :, [2, 1, 0]]
+ img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
+ img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+
+
+ # modcrop
+ _, H, W = img_LR.size()
+ img_GT = img_GT[:, :H*self.scale, :W*self.scale]
+
+ return {'LQ': img_LR, 'GT': img_GT, 'LQ_path': LR_path, 'GT_path': GT_path}
+
+ def __len__(self):
+ return len(self.GT_paths)
+
+
+def impad(img, top=0, bottom=0, left=0, right=0, color=255):
+ return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect')
diff --git a/codes/data/GTLQnpy_dataset.py b/codes/data/GTLQnpy_dataset.py
new file mode 100644
index 0000000..89eae1d
--- /dev/null
+++ b/codes/data/GTLQnpy_dataset.py
@@ -0,0 +1,82 @@
+import random
+import numpy as np
+import cv2
+import lmdb
+import torch
+import torch.utils.data as data
+import data.util as util
+import sys
+import os
+
+try:
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ from data.util import imresize_np
+ from utils import util as utils
+except ImportError:
+ pass
+
+
+class GTLQnpyDataset(data.Dataset):
+ '''
+ Load HR-LR image npy pairs. Make sure HR-LR images are in the same order.
+ '''
+
+ def __init__(self, opt):
+ super(GTLQnpyDataset, self).__init__()
+ self.opt = opt
+ self.LR_paths, self.GT_paths = None, None
+ self.scale = opt['scale']
+ if self.opt['phase'] == 'train':
+ self.GT_size = opt['GT_size']
+ self.LR_size = self.GT_size // self.scale
+
+ self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
+ self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
+
+ assert self.GT_paths, 'Error: GT paths are empty.'
+ if self.LR_paths and self.GT_paths:
+ assert len(self.LR_paths) == len(
+ self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(
+ len(self.LR_paths), len(self.GT_paths))
+
+ def __getitem__(self, index):
+ # get GT and LR image
+ GT_path = self.GT_paths[index]
+ # LR_path = self.LR_paths[index]
+ LR_path = GT_path.replace('DIV2K+Flickr2K_HR', 'DIV2K+Flickr2K_LR_bicubic/X4').replace('.npy','x{}.npy'.format(self.scale))
+ img_GT = util.read_img_fromnpy(np.load(GT_path))
+ img_LR = util.read_img_fromnpy(np.load(LR_path)) # return: Numpy float32, HWC, BGR, [0,1]
+
+ if self.opt['phase'] == 'train':
+ # crop
+ H, W, C = img_LR.shape
+ rnd_top_LR = random.randint(0, max(0, H - self.LR_size))
+ rnd_left_LR = random.randint(0, max(0, W - self.LR_size))
+ rnd_top_GT = rnd_top_LR * self.scale
+ rnd_left_GT = rnd_left_LR * self.scale
+
+ img_GT = img_GT[rnd_top_GT:rnd_top_GT + self.GT_size, rnd_left_GT:rnd_left_GT + self.GT_size, :]
+ img_LR = img_LR[rnd_top_LR:rnd_top_LR + self.LR_size, rnd_left_LR:rnd_left_LR + self.LR_size, :]
+
+ # augmentation - flip, rotate
+ img_GT, img_LR = util.augment([img_GT, img_LR], self.opt['use_flip'],
+ self.opt['use_rot'], self.opt['mode'])
+
+ # change color space if necessary, deal with gray image
+ if self.opt['color']:
+ img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
+ img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0]
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ if img_GT.shape[2] == 3:
+ img_GT = img_GT[:, :, [2, 1, 0]]
+ if img_LR.shape[2] == 3:
+ img_LR = img_LR[:, :, [2, 1, 0]]
+ img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
+ img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+
+ return {'LQ': img_LR, 'GT': img_GT, 'LQ_path': LR_path, 'GT_path': GT_path}
+
+ def __len__(self):
+ return len(self.GT_paths)
+
diff --git a/codes/data/GTLQx_dataset.py b/codes/data/GTLQx_dataset.py
new file mode 100644
index 0000000..1bc5c15
--- /dev/null
+++ b/codes/data/GTLQx_dataset.py
@@ -0,0 +1,129 @@
+import random
+import numpy as np
+import cv2
+import lmdb
+import torch
+import torch.utils.data as data
+import data.util as util
+import sys
+import os
+
+try:
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ from data.util import imresize_np
+ from utils import util as utils
+except ImportError:
+ pass
+
+
+class GTLQxDataset(data.Dataset):
+ '''
+ Load HR-LR image pairs.
+ '''
+
+ def __init__(self, opt):
+ super(GTLQxDataset, self).__init__()
+ self.opt = opt
+ self.LR_paths, self.GT_paths = None, None
+ self.LR_env, self.GT_env = None, None # environment for lmdb
+ self.scale = opt['scale']
+ if self.opt['phase'] == 'train':
+ self.GT_size = opt['GT_size']
+ self.LR_size = self.GT_size // self.scale
+
+ # read image list from lmdb or image files
+ if opt['data_type'] == 'lmdb':
+ self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
+ self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
+ elif opt['data_type'] == 'img':
+ self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
+ self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
+ else:
+ print('Error: data_type is not matched in Dataset')
+ assert self.GT_paths, 'Error: GT paths are empty.'
+ if self.LR_paths and self.GT_paths:
+ assert len(self.LR_paths) == len(
+ self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(
+ len(self.LR_paths), len(self.GT_paths))
+
+ def _init_lmdb(self):
+ # https://github.com/chainer/chainermn/issues/129
+ self.GT_env = lmdb.open(self.opt['dataroot_GT'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ if self.opt['dataroot_LQ'] is not None:
+ self.LR_env = lmdb.open(self.opt['dataroot_LQ'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ else:
+ self.LR_env = 'No lmdb input for LR'
+
+ def __getitem__(self, index):
+ if self.opt['data_type'] == 'lmdb':
+ if (self.GT_env is None) or (self.LR_env is None):
+ self._init_lmdb()
+
+ if self.opt['data_type'] == 'lmdb':
+ resolution = [int(s) for s in self.GT_sizes[index].split('_')]
+ else:
+ resolution = None
+
+
+ # loading code from srflow test
+ # img_GT = cv2.imread(GT_path)[:, :, [2, 1, 0]]
+ # img_GT = torch.Tensor(img_GT.transpose([2, 0, 1]).astype(np.float32)) / 255
+ # img_LR = cv2.imread(LR_path)[:, :, [2, 1, 0]]
+ # pad_factor = 2
+ # h, w, c = img_LR.shape
+ # img_LR = impad(img_LR, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
+ # right=int(np.ceil(w / pad_factor) * pad_factor - w))
+ # img_LR = torch.Tensor(img_LR.transpose([2, 0, 1]).astype(np.float32)) / 255
+
+
+ # get GT and LR image
+ GT_path = self.GT_paths[index]
+ # LR_path = self.LR_paths[index]
+ LR_path = GT_path.replace('HR', 'LR_bicubic/X4').replace('.png','x{}.png'.format(self.scale))
+ img_GT = util.read_img(self.GT_env, GT_path, resolution) # return: Numpy float32, HWC, BGR, [0,1]
+ img_LR = util.read_img(self.LR_env, LR_path, resolution)
+
+
+ if self.opt['phase'] == 'train':
+ # crop
+ H, W, C = img_LR.shape
+ rnd_top_LR = random.randint(0, max(0, H - self.LR_size))
+ rnd_left_LR = random.randint(0, max(0, W - self.LR_size))
+ rnd_top_GT = rnd_top_LR * self.scale
+ rnd_left_GT = rnd_left_LR * self.scale
+
+ img_GT = img_GT[rnd_top_GT:rnd_top_GT + self.GT_size, rnd_left_GT:rnd_left_GT + self.GT_size, :]
+ img_LR = img_LR[rnd_top_LR:rnd_top_LR + self.LR_size, rnd_left_LR:rnd_left_LR + self.LR_size, :]
+
+ # augmentation - flip, rotate
+ img_GT, img_LR = util.augment([img_GT, img_LR], self.opt['use_flip'],
+ self.opt['use_rot'], self.opt['mode'])
+
+ # change color space if necessary, deal with gray image
+ if self.opt['color']:
+ img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
+ img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0]
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ if img_GT.shape[2] == 3:
+ img_GT = img_GT[:, :, [2, 1, 0]]
+ if img_LR.shape[2] == 3:
+ img_LR = img_LR[:, :, [2, 1, 0]]
+ img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
+ img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+
+
+ # modcrop
+ _, H, W = img_LR.size()
+ img_GT = img_GT[:, :H*self.scale, :W*self.scale]
+
+ return {'LQ': img_LR, 'GT': img_GT, 'LQ_path': LR_path, 'GT_path': GT_path}
+
+ def __len__(self):
+ return len(self.GT_paths)
+
+
+def impad(img, top=0, bottom=0, left=0, right=0, color=255):
+ return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect')
diff --git a/codes/data/GT_dataset.py b/codes/data/GT_dataset.py
new file mode 100644
index 0000000..7bd5294
--- /dev/null
+++ b/codes/data/GT_dataset.py
@@ -0,0 +1,131 @@
+import random
+import numpy as np
+import cv2
+import lmdb
+import torch
+import torch.utils.data as data
+import data.util as util
+import sys
+import os
+
+try:
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ from data.util import imresize_np
+ from utils import util as utils
+except ImportError:
+ pass
+
+
+class GTDataset(data.Dataset):
+ '''
+ Load GT images only. 30s faster than LQGTKer (90s for 200iter).
+ '''
+
+ def __init__(self, opt):
+ super(GTDataset, self).__init__()
+ self.opt = opt
+ self.LR_paths, self.GT_paths = None, None
+ self.LR_env, self.GT_env = None, None # environment for lmdb
+ self.LR_size, self.GT_size = opt['LR_size'], opt['GT_size']
+
+ # read image list from lmdb or image files
+ if opt['data_type'] == 'lmdb':
+ self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
+ self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
+ elif opt['data_type'] == 'img':
+ self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
+ self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
+ else:
+ print('Error: data_type is not matched in Dataset')
+ assert self.GT_paths, 'Error: GT paths are empty.'
+ if self.LR_paths and self.GT_paths:
+ assert len(self.LR_paths) == len(
+ self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(
+ len(self.LR_paths), len(self.GT_paths))
+ self.random_scale_list = [1]
+
+ def _init_lmdb(self):
+ # https://github.com/chainer/chainermn/issues/129
+ self.GT_env = lmdb.open(self.opt['dataroot_GT'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ if self.opt['dataroot_LQ'] is not None:
+ self.LR_env = lmdb.open(self.opt['dataroot_LQ'], readonly=True, lock=False, readahead=False,
+ meminit=False)
+ else:
+ self.LR_env = 'No lmdb input for LR'
+
+ def __getitem__(self, index):
+ if self.opt['data_type'] == 'lmdb':
+ if (self.GT_env is None) or (self.LR_env is None):
+ self._init_lmdb()
+
+ GT_path, LR_path = None, None
+ scale = self.opt['scale']
+ GT_size = self.opt['GT_size']
+
+ # get GT image
+ GT_path = self.GT_paths[index]
+ if self.opt['data_type'] == 'lmdb':
+ resolution = [int(s) for s in self.GT_sizes[index].split('_')]
+ else:
+ resolution = None
+ img_GT = util.read_img(self.GT_env, GT_path, resolution) # return: Numpy float32, HWC, BGR, [0,1]
+
+ # modcrop in the validation / test phase
+ img_GT = util.modcrop(img_GT, scale)
+
+ # get LR image
+ if self.LR_paths: # LR exist
+ raise ValueError('GTker_dataset.py doesn Not allow LR input.')
+
+ else: # down-sampling on-the-fly
+ # randomly scale during training
+ if self.opt['phase'] == 'train':
+ random_scale = random.choice(self.random_scale_list)
+ if random_scale != 1:
+ H_s, W_s, _ = img_GT.shape
+ H_s = _mod(H_s, random_scale, scale, GT_size)
+ W_s = _mod(W_s, random_scale, scale, GT_size)
+ img_GT = cv2.resize(np.copy(img_GT), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
+
+ # force to 3 channels
+ if img_GT.ndim == 2:
+ img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
+
+ if self.opt['phase'] == 'train':
+ H, W, C = img_GT.shape
+
+ # randomly crop on HR, more positions than first crop on LR and HR simultaneously
+ rnd_h_GT = random.randint(0, max(0, H - GT_size))
+ rnd_w_GT = random.randint(0, max(0, W - GT_size))
+ img_GT = img_GT[rnd_h_GT:rnd_h_GT + GT_size, rnd_w_GT:rnd_w_GT + GT_size, :]
+
+ # augmentation - flip, rotate
+ img_GT = util.augment(img_GT, self.opt['use_flip'],
+ self.opt['use_rot'], self.opt['mode'])
+
+ # change color space if necessary, deal with gray image
+ if self.opt['color']:
+ img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ if img_GT.shape[2] == 3:
+ img_GT = img_GT[:, :, [2, 1, 0]]
+ img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
+
+ if LR_path is None:
+ LR_path = GT_path
+
+ # don't need LR because it's generated from HR batches.
+ img_LR = torch.ones(1, 1, 1)
+
+ return {'LQ': img_LR, 'GT': img_GT, 'LQ_path': LR_path, 'GT_path': GT_path}
+
+ def __len__(self):
+ return len(self.GT_paths)
+
+
+def _mod(n, random_scale, scale, thres):
+ rlt = int(n * random_scale)
+ rlt = (rlt // scale) * scale
+ return thres if rlt < thres else rlt
diff --git a/codes/data/LQ_dataset.py b/codes/data/LQ_dataset.py
new file mode 100644
index 0000000..d11ae9c
--- /dev/null
+++ b/codes/data/LQ_dataset.py
@@ -0,0 +1,106 @@
+import random
+import numpy as np
+import cv2
+import lmdb
+import torch
+import torch.nn.functional as F
+import torch.utils.data as data
+import data.util as util
+import sys
+import os
+
+try:
+ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ from data.util import imresize_np
+ from utils import util as utils
+except ImportError:
+ pass
+
+
+class LQDataset(data.Dataset):
+ '''
+ Load LR images only, e.g. real-world images
+ '''
+
+ def __init__(self, opt):
+ super(LQDataset, self).__init__()
+ self.opt = opt
+ self.LR_paths, self.GT_paths = None, None
+ self.LR_env, self.GT_env = None, None # environment for lmdb
+ self.LR_size, self.GT_size = opt['LR_size'], opt['GT_size']
+
+ # read image list from lmdb or image files
+ if opt['data_type'] == 'lmdb':
+ self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
+ self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
+ elif opt['data_type'] == 'img':
+ self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
+ self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
+ else:
+ print('Error: data_type is not matched in Dataset')
+ assert self.LR_paths, 'Error: LQ paths are empty.'
+ if self.LR_paths and self.GT_paths:
+ assert len(self.LR_paths) == len(
+ self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(
+ len(self.LR_paths), len(self.GT_paths))
+ self.random_scale_list = [1]
+
+ def __getitem__(self, index):
+
+ GT_path, LQ_path = None, None
+
+ # get GT image
+ LQ_path = self.LR_paths[index]
+ img_LQ = util.read_img(None, LQ_path, None) # return: Numpy float32, HWC, BGR, [0,1]
+
+ if self.GT_paths: # LR exist
+ raise ValueError('LQ_dataset.py doesn Not allow HR input.')
+
+ else:
+ # force to 3 channels
+ if img_LQ.ndim == 2:
+ img_LQ = cv2.cvtColor(img_LQ, cv2.COLOR_GRAY2BGR)
+
+ # change color space if necessary, deal with gray image
+ if self.opt['color']:
+ img_LQ = util.channel_convert(img_LQ.shape[2], self.opt['color'], [img_LQ])[0]
+
+ # BGR to RGB, HWC to CHW, numpy to tensor
+ if img_LQ.shape[2] == 3:
+ img_LQ = img_LQ[:, :, [2, 1, 0]]
+ img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
+
+ if GT_path is None:
+ GT_path = LQ_path
+
+ # don't need LR because it's generated from HR batches.
+ img_GT = torch.ones(1, 1, 1)
+
+ # deal with the image margins for real-world images
+ img_LQ = img_LQ.unsqueeze(0)
+ x_gt = F.interpolate(img_LQ, scale_factor=self.opt['scale'], mode='nearest')
+ if self.opt['scale'] == 4:
+ real_crop = 3
+ elif self.opt['scale'] == 2:
+ real_crop = 6
+ elif self.opt['scale'] == 1:
+ real_crop = 11
+ assert real_crop * self.opt['scale'] * 2 > self.opt['kernel_size']
+ x_gt = F.pad(x_gt, pad=(
+ real_crop * self.opt['scale'], real_crop * self.opt['scale'], real_crop * self.opt['scale'],
+ real_crop * self.opt['scale']), mode='replicate') # 'constant', 'reflect', 'replicate' or 'circular
+
+ kernel_gt, sigma_gt = utils.stable_batch_kernel(1, l=self.opt['kernel_size'], sig=10, sig1=0, sig2=0,
+ theta=0, rate_iso=1, scale=self.opt['scale'],
+ tensor=True) # generate kernel [BHW], y [BCHW]
+
+ blur_layer = utils.BatchBlur(l=self.opt['kernel_size'], padmode='zero')
+ sample_layer = utils.BatchSubsample(scale=self.opt['scale'])
+ y_blurred = sample_layer(blur_layer(x_gt, kernel_gt))
+ y_blurred[:, :, real_crop:-real_crop, real_crop:-real_crop] = img_LQ
+ img_LQ = y_blurred.squeeze(0)
+
+ return {'LQ': img_LQ, 'GT': img_GT, 'LQ_path': LQ_path, 'GT_path': GT_path}
+
+ def __len__(self):
+ return len(self.LR_paths)
diff --git a/codes/data/LRHR_PKL_dataset.py b/codes/data/LRHR_PKL_dataset.py
new file mode 100644
index 0000000..5a25e8a
--- /dev/null
+++ b/codes/data/LRHR_PKL_dataset.py
@@ -0,0 +1,193 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file contains content licensed by https://github.com/xinntao/BasicSR/blob/master/LICENSE/LICENSE
+
+import os
+# import subprocess
+import torch.utils.data as data
+import numpy as np
+import time
+import torch
+
+import pickle
+
+
+class LRHR_PKLDataset(data.Dataset):
+ def __init__(self, opt):
+ super(LRHR_PKLDataset, self).__init__()
+ self.opt = opt
+ self.crop_size = opt.get("GT_size", None)
+ self.scale = None
+ self.random_scale_list = [1]
+
+ hr_file_path = opt["dataroot_GT"]
+ lr_file_path = opt["dataroot_LQ"]
+ y_labels_file_path = opt['dataroot_y_labels']
+
+ gpu = True
+ augment = True
+
+ self.use_flip = opt["use_flip"] if "use_flip" in opt.keys() else False
+ self.use_rot = opt["use_rot"] if "use_rot" in opt.keys() else False
+ self.use_crop = opt["use_crop"] if "use_crop" in opt.keys() else False
+ self.center_crop_hr_size = opt.get("center_crop_hr_size", None)
+
+ n_max = opt["n_max"] if "n_max" in opt.keys() else int(1e8)
+
+ t = time.time()
+ self.lr_images = self.load_pkls(lr_file_path, n_max)
+ self.hr_images = self.load_pkls(hr_file_path, n_max)
+
+ min_val_hr = np.min([i.min() for i in self.hr_images[:20]])
+ max_val_hr = np.max([i.max() for i in self.hr_images[:20]])
+
+ min_val_lr = np.min([i.min() for i in self.lr_images[:20]])
+ max_val_lr = np.max([i.max() for i in self.lr_images[:20]])
+
+ t = time.time() - t
+ print("Loaded {} HR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+ format(len(self.hr_images), min_val_hr, max_val_hr, t, hr_file_path))
+ print("Loaded {} LR images with [{:.2f}, {:.2f}] in {:.2f}s from {}".
+ format(len(self.lr_images), min_val_lr, max_val_lr, t, lr_file_path))
+
+ self.gpu = gpu
+ self.augment = augment
+
+ self.measures = None
+
+ # todo: this is very slow (~15min) when using nn.DistributedDataParallel(), and we have to set n_worker=0
+
+ # # save as png
+ # import cv2
+ # for i in range(400):
+ # img = self.hr_images[i]
+ # img = np.transpose(img, [1,2,0])
+ # cv2.imwrite("/cluster/work/cvl/jinliang/log/srflow_experiments/ICCV21/baseline_results/CelebA_HR_8X/{}.png".format(i), img[:,:,[2, 1, 0]])
+ # img = self.lr_images[i]
+ # img = np.transpose(img, [1,2,0])
+ # cv2.imwrite("/cluster/work/cvl/jinliang/log/srflow_experiments/ICCV21/baseline_results/CelebA_LR_8X/{}.png".format(i), img[:,:,[2, 1, 0]])
+ # raise NotImplementedError
+
+ def load_pkls(self, path, n_max):
+ assert os.path.isfile(path), path
+ images = []
+ with open(path, "rb") as f:
+ images += pickle.load(f)
+ assert len(images) > 0, path
+ images = images[:n_max]
+ images = [np.transpose(image, [2, 0, 1]) for image in images]
+ return images
+
+ def __len__(self):
+ return len(self.hr_images)
+
+ def __getitem__(self, item):
+
+ hr = self.hr_images[item]
+ lr = self.lr_images[item]
+
+ if self.scale == None:
+ self.scale = hr.shape[1] // lr.shape[1]
+ assert hr.shape[1] == self.scale * lr.shape[1], ('non-fractional ratio', lr.shape, hr.shape)
+
+ if self.use_crop:
+ hr, lr = random_crop(hr, lr, self.crop_size, self.scale, self.use_crop)
+
+ if self.center_crop_hr_size:
+ hr, lr = center_crop(hr, self.center_crop_hr_size), center_crop(lr, self.center_crop_hr_size // self.scale)
+
+ if self.use_flip:
+ hr, lr = random_flip(hr, lr)
+
+ if self.use_rot:
+ hr, lr = random_rotation(hr, lr)
+
+ hr = hr / 255.0
+ lr = lr / 255.0
+
+ if self.measures is None or np.random.random() < 0.05:
+ if self.measures is None:
+ self.measures = {}
+ self.measures['hr_means'] = np.mean(hr)
+ self.measures['hr_stds'] = np.std(hr)
+ self.measures['lr_means'] = np.mean(lr)
+ self.measures['lr_stds'] = np.std(lr)
+
+ hr = torch.Tensor(hr)
+ lr = torch.Tensor(lr)
+
+ # if self.gpu:
+ # hr = hr.cuda()
+ # lr = lr.cuda()
+
+ return {'LQ': lr, 'GT': hr, 'LQ_path': str(item), 'GT_path': str(item)}
+
+ def print_and_reset(self, tag):
+ m = self.measures
+ kvs = []
+ for k in sorted(m.keys()):
+ kvs.append("{}={:.2f}".format(k, m[k]))
+ print("[KPI] " + tag + ": " + ", ".join(kvs))
+ self.measures = None
+
+
+def random_flip(img, seg):
+ random_choice = np.random.choice([True, False])
+ img = img if random_choice else np.flip(img, 2).copy()
+ seg = seg if random_choice else np.flip(seg, 2).copy()
+ return img, seg
+
+
+def random_rotation(img, seg):
+ random_choice = np.random.choice([0, 1, 3])
+ img = np.rot90(img, random_choice, axes=(1, 2)).copy()
+ seg = np.rot90(seg, random_choice, axes=(1, 2)).copy()
+ return img, seg
+
+
+def random_crop(hr, lr, size_hr, scale, random):
+ size_lr = size_hr // scale
+
+ size_lr_x = lr.shape[1]
+ size_lr_y = lr.shape[2]
+
+ start_x_lr = np.random.randint(low=0, high=(size_lr_x - size_lr) + 1) if size_lr_x > size_lr else 0
+ start_y_lr = np.random.randint(low=0, high=(size_lr_y - size_lr) + 1) if size_lr_y > size_lr else 0
+
+ # LR Patch
+ lr_patch = lr[:, start_x_lr:start_x_lr + size_lr, start_y_lr:start_y_lr + size_lr]
+
+ # HR Patch
+ start_x_hr = start_x_lr * scale
+ start_y_hr = start_y_lr * scale
+ hr_patch = hr[:, start_x_hr:start_x_hr + size_hr, start_y_hr:start_y_hr + size_hr]
+
+ return hr_patch, lr_patch
+
+
+def center_crop(img, size):
+ assert img.shape[1] == img.shape[2], img.shape
+ border_double = img.shape[1] - size
+ assert border_double % 2 == 0, (img.shape, size)
+ border = border_double // 2
+ return img[:, border:-border, border:-border]
+
+
+def center_crop_tensor(img, size):
+ assert img.shape[2] == img.shape[3], img.shape
+ border_double = img.shape[2] - size
+ assert border_double % 2 == 0, (img.shape, size)
+ border = border_double // 2
+ return img[:, :, border:-border, border:-border]
diff --git a/codes/data/__init__.py b/codes/data/__init__.py
new file mode 100644
index 0000000..fd4ebb7
--- /dev/null
+++ b/codes/data/__init__.py
@@ -0,0 +1,54 @@
+'''create dataset and dataloader'''
+import logging
+import torch
+import torch.utils.data
+
+
+def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
+ phase = dataset_opt['phase']
+ if phase == 'train':
+ if opt['dist']:
+ world_size = torch.distributed.get_world_size()
+ num_workers = dataset_opt['n_workers']
+ assert dataset_opt['batch_size'] % world_size == 0
+ batch_size = dataset_opt['batch_size'] // world_size
+ shuffle = False
+ else:
+ num_workers = dataset_opt['n_workers'] * len(opt['gpu_ids'])
+ batch_size = dataset_opt['batch_size']
+ shuffle = True
+ return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
+ num_workers=num_workers, sampler=sampler, drop_last=True,
+ pin_memory=True)
+ else:
+ return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0,
+ pin_memory=True)
+
+
+def create_dataset(dataset_opt):
+ mode = dataset_opt['mode']
+ if mode == 'GT': # load HR image and generate LR on-the-fly
+ from data.GT_dataset import GTDataset as D
+ dataset = D(dataset_opt)
+ elif mode == 'GTLQ': # load generated HR-LR image pairs
+ from data.GTLQ_dataset import GTLQDataset as D
+ dataset = D(dataset_opt)
+ elif mode == 'GTLQx': # load generated HR-LR image pairs, and replace with x4
+ from data.GTLQx_dataset import GTLQxDataset as D
+ dataset = D(dataset_opt)
+ elif mode == 'LQ': # load LR image for testing
+ from data.LQ_dataset import LQDataset as D
+ dataset = D(dataset_opt)
+ elif mode == 'LRHR_PKL':
+ from data.LRHR_PKL_dataset import LRHR_PKLDataset as D
+ dataset = D(dataset_opt)
+ elif mode == 'GTLQnpy': # load generated HR-LR image pairs
+ from data.GTLQnpy_dataset import GTLQnpyDataset as D
+ dataset = D(dataset_opt)
+ else:
+ raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
+
+ logger = logging.getLogger('base')
+ logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
+ dataset_opt['name']))
+ return dataset
diff --git a/codes/data/data_sampler.py b/codes/data/data_sampler.py
new file mode 100644
index 0000000..ba9a72b
--- /dev/null
+++ b/codes/data/data_sampler.py
@@ -0,0 +1,109 @@
+"""
+Modified from torch.utils.data.distributed.DistributedSampler
+Support enlarging the dataset for *iter-oriented* training, for saving time when restart the
+dataloader after each epoch
+"""
+import math
+import torch
+from torch.utils.data.sampler import Sampler
+import torch.distributed as dist
+
+
+class DistIterSampler(Sampler):
+ """Sampler that restricts data loading to a subset of the dataset.
+
+ It is especially useful in conjunction with
+ :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+ process can pass a DistributedSampler instance as a DataLoader sampler,
+ and load a subset of the original dataset that is exclusive to it.
+
+ .. note::
+ Dataset is assumed to be of constant size.
+
+ Arguments:
+ dataset: Dataset used for sampling.
+ num_replicas (optional): Number of processes participating in
+ distributed training.
+ rank (optional): Rank of the current process within num_replicas.
+ """
+
+ def __init__(self, dataset, num_replicas=None, rank=None, ratio=100):
+ if num_replicas is None:
+ if not dist.is_available():
+ raise RuntimeError("Requires distributed package to be available")
+ num_replicas = dist.get_world_size()
+ if rank is None:
+ if not dist.is_available():
+ raise RuntimeError("Requires distributed package to be available")
+ rank = dist.get_rank()
+ self.dataset = dataset
+ self.num_replicas = num_replicas
+ self.rank = rank
+ self.epoch = 0
+ self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas))
+ self.total_size = self.num_samples * self.num_replicas
+
+ def __iter__(self):
+ # deterministically shuffle based on epoch
+ g = torch.Generator()
+ g.manual_seed(self.epoch)
+ indices = torch.randperm(self.total_size, generator=g).tolist() #Returns a random permutation of integers from 0 to n - 1
+
+ dsize = len(self.dataset)
+ indices = [v % dsize for v in indices]
+
+ # subsample
+ indices = indices[self.rank:self.total_size:self.num_replicas]
+ assert len(indices) == self.num_samples
+
+ return iter(indices)
+
+ def __len__(self):
+ return self.num_samples
+
+ def set_epoch(self, epoch):
+ self.epoch = epoch
+
+
+class EnlargedSampler(Sampler):
+ """Sampler that restricts data loading to a subset of the dataset.
+ Modified from torch.utils.data.distributed.DistributedSampler
+ Support enlarging the dataset for iteration-based training, for saving
+ time when restart the dataloader after each epoch
+ Args:
+ dataset (torch.utils.data.Dataset): Dataset used for sampling.
+ num_replicas (int | None): Number of processes participating in
+ the training. It is usually the world_size.
+ rank (int | None): Rank of the current process within num_replicas.
+ ratio (int): Enlarging ratio. Default: 1.
+ """
+
+ def __init__(self, dataset, num_replicas, rank, ratio=1):
+ self.dataset = dataset
+ self.num_replicas = num_replicas
+ self.rank = rank
+ self.epoch = 0
+ self.num_samples = math.ceil(
+ len(self.dataset) * ratio / self.num_replicas)
+ self.total_size = self.num_samples * self.num_replicas
+
+ def __iter__(self):
+ # deterministically shuffle based on epoch
+ g = torch.Generator()
+ g.manual_seed(self.epoch)
+ indices = torch.randperm(self.total_size, generator=g).tolist()
+
+ dataset_size = len(self.dataset)
+ indices = [v % dataset_size for v in indices]
+
+ # subsample
+ indices = indices[self.rank:self.total_size:self.num_replicas]
+ assert len(indices) == self.num_samples
+
+ return iter(indices)
+
+ def __len__(self):
+ return self.num_samples
+
+ def set_epoch(self, epoch):
+ self.epoch = epoch
\ No newline at end of file
diff --git a/codes/data/util.py b/codes/data/util.py
new file mode 100644
index 0000000..84666f0
--- /dev/null
+++ b/codes/data/util.py
@@ -0,0 +1,536 @@
+import os
+import math
+import pickle
+import random
+import numpy as np
+import torch
+import cv2
+
+####################
+# Files & IO
+####################
+
+###################### get image path list ######################
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.npy']
+
+
+def is_image_file(filename):
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def _get_paths_from_images(path):
+ '''get image path list from image folder'''
+ assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+ images = []
+ for dirpath, _, fnames in sorted(os.walk(path)):
+ for fname in sorted(fnames):
+ if is_image_file(fname):
+ img_path = os.path.join(dirpath, fname)
+ images.append(img_path)
+ assert images, '{:s} has no valid image file'.format(path)
+ return images
+
+
+def _get_paths_from_lmdb(dataroot):
+ '''get image path list from lmdb meta info'''
+ meta_info = pickle.load(open(os.path.join(dataroot, 'meta_info.pkl'), 'rb'))
+ paths = meta_info['keys']
+ sizes = meta_info['resolution']
+ if len(sizes) == 1:
+ sizes = sizes * len(paths)
+ return paths, sizes
+
+
+def get_image_paths(data_type, dataroot):
+ '''get image path list
+ support lmdb or image files'''
+ paths, sizes = None, None
+ if data_type == 'lmdb':
+ if dataroot is not None:
+ paths, sizes = _get_paths_from_lmdb(dataroot)
+ return paths, sizes
+ elif data_type == 'img':
+ if dataroot is not None:
+ paths = sorted(_get_paths_from_images(dataroot))
+ return paths
+ else:
+ raise NotImplementedError('data_type [{:s}] is not recognized.'.format(data_type))
+
+
+###################### read images ######################
+def _read_img_lmdb(env, key, size):
+ '''read image from lmdb with key (w/ and w/o fixed size)
+ size: (C, H, W) tuple'''
+ with env.begin(write=False) as txn:
+ buf = txn.get(key.encode('ascii'))
+ img_flat = np.frombuffer(buf, dtype=np.uint8)
+ C, H, W = size
+ img = img_flat.reshape(H, W, C)
+ return img
+
+
+def read_img(env, path, size=None):
+ '''read image by cv2 or from lmdb
+ return: Numpy float32, HWC, BGR, [0,1]'''
+ if env is None: # img
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+ else:
+ img = _read_img_lmdb(env, path, size)
+ img = img.astype(np.float32) / 255.
+ if img.ndim == 2:
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # added by jinliang
+ # img = np.expand_dims(img, axis=2)
+ # some images have 4 channels
+ if img.shape[2] > 3:
+ img = img[:, :, :3]
+ return img
+
+def read_img_fromcv2(img):
+ '''transform image in cv2 to the same output as read_img
+ return: Numpy float32, HWC, BGR, [0,1]'''
+ img = img.astype(np.float32) / 255.
+ if img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ # some images have 4 channels
+ if img.shape[2] > 3:
+ img = img[:, :, :3]
+ return img
+
+def read_img_fromnpy(img):
+ '''read image in sio to the same output as read_img
+ return: Numpy float32, HWC, BGR, [0,1]'''
+ img = img.astype(np.float32) / 255.
+ if img.ndim == 2:
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # added by jinliang
+ # img = np.expand_dims(img, axis=2)
+ # some images have 4 channels
+ if img.shape[2] > 3:
+ img = img[:, :, :3]
+ return img
+####################
+# image processing
+# process on numpy image
+####################
+
+
+def augment(img, hflip=True, rot=True, mode=None):
+ # horizontal flip OR rotate
+ hflip = hflip and random.random() < 0.5
+ vflip = rot and random.random() < 0.5
+ rot90 = rot and random.random() < 0.5
+
+ def _augment(img):
+ if hflip:
+ img = img[:, ::-1, :]
+ if vflip:
+ img = img[::-1, :, :]
+ if rot90:
+ img = img.transpose(1, 0, 2)
+ return img
+ if mode == 'LQ' or mode == 'GT' or mode == 'SRker' or mode == 'GTker' or mode == 'GTker_memory' or mode == 'GTLQnpy_noisy':
+ return _augment(img)
+ elif mode == 'GTLQ' or mode == 'GTLQnpy':
+ return [_augment(I) for I in img]
+ else:
+ raise NotImplementedError('{} dataloader has not been implemented yet.'.format(mode))
+
+
+def augment_flow(img_list, flow_list, hflip=True, rot=True):
+ # horizontal flip OR rotate
+ hflip = hflip and random.random() < 0.5
+ vflip = rot and random.random() < 0.5
+ rot90 = rot and random.random() < 0.5
+
+ def _augment(img):
+ if hflip:
+ img = img[:, ::-1, :]
+ if vflip:
+ img = img[::-1, :, :]
+ if rot90:
+ img = img.transpose(1, 0, 2)
+ return img
+
+ def _augment_flow(flow):
+ if hflip:
+ flow = flow[:, ::-1, :]
+ flow[:, :, 0] *= -1
+ if vflip:
+ flow = flow[::-1, :, :]
+ flow[:, :, 1] *= -1
+ if rot90:
+ flow = flow.transpose(1, 0, 2)
+ flow = flow[:, :, [1, 0]]
+ return flow
+
+ rlt_img_list = [_augment(img) for img in img_list]
+ rlt_flow_list = [_augment_flow(flow) for flow in flow_list]
+
+ return rlt_img_list, rlt_flow_list
+
+
+def channel_convert(in_c, tar_type, img_list):
+ # conversion among BGR, gray and y
+ if in_c == 3 and tar_type == 'gray': # BGR to gray
+ gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
+ return [np.expand_dims(img, axis=2) for img in gray_list]
+ elif in_c == 3 and tar_type == 'y': # BGR to y
+ y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
+ return [np.expand_dims(img, axis=2) for img in y_list]
+ elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR
+ return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
+ else:
+ return img_list
+
+
+def rgb2ycbcr(img, only_y=True):
+ '''same as matlab rgb2ycbcr
+ only_y: only return Y channel
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ if only_y:
+ rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
+ else:
+ rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+ [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+
+def bgr2ycbcr(img, only_y=True):
+ '''bgr version of rgb2ycbcr
+ only_y: only return Y channel
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ if only_y:
+ rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
+ else:
+ rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+ [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+
+def ycbcr2rgb(img):
+ '''same as matlab ycbcr2rgb
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+ [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+def modcrop(img_in, scale):
+ # img_in: Numpy, HWC or HW
+ img = np.copy(img_in)
+ if img.ndim == 2:
+ H, W = img.shape
+ H_r, W_r = H % scale, W % scale
+ img = img[:H - H_r, :W - W_r]
+ elif img.ndim == 3:
+ H, W, C = img.shape
+ H_r, W_r = H % scale, W % scale
+ img = img[:H - H_r, :W - W_r, :]
+ else:
+ raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
+ return img
+
+
+####################
+# Functions
+####################
+
+
+# matlab 'imresize' function, now only support 'bicubic'
+def cubic(x):
+ absx = torch.abs(x)
+ absx2 = absx**2
+ absx3 = absx**3
+ return (1.5 * absx3 - 2.5 * absx2 + 1) * (
+ (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * ((
+ (absx > 1) * (absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+ if (scale < 1) and (antialiasing):
+ # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
+ kernel_width = kernel_width / scale
+
+ # Output-space coordinates
+ x = torch.linspace(1, out_length, out_length)
+
+ # Input-space coordinates. Calculate the inverse mapping such that 0.5
+ # in output space maps to 0.5 in input space, and 0.5+scale in output
+ # space maps to 1.5 in input space.
+ u = x / scale + 0.5 * (1 - 1 / scale)
+
+ # What is the left-most pixel that can be involved in the computation?
+ left = torch.floor(u - kernel_width / 2)
+
+ # What is the maximum number of pixels that can be involved in the
+ # computation? Note: it's OK to use an extra pixel here; if the
+ # corresponding weights are all zero, it will be eliminated at the end
+ # of this function.
+ P = math.ceil(kernel_width) + 2
+
+ # The indices of the input pixels involved in computing the k-th output
+ # pixel are in row k of the indices matrix.
+ indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
+ 1, P).expand(out_length, P)
+
+ # The weights used to compute the k-th output pixel are in row k of the
+ # weights matrix.
+ distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
+ # apply cubic kernel
+ if (scale < 1) and (antialiasing):
+ weights = scale * cubic(distance_to_center * scale)
+ else:
+ weights = cubic(distance_to_center)
+ # Normalize the weights matrix so that each row sums to 1.
+ weights_sum = torch.sum(weights, 1).view(out_length, 1)
+ weights = weights / weights_sum.expand(out_length, P)
+
+ # If a column in weights is all zero, get rid of it. only consider the first and last column.
+ weights_zero_tmp = torch.sum((weights == 0), 0)
+ if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+ indices = indices.narrow(1, 1, P - 2)
+ weights = weights.narrow(1, 1, P - 2)
+ if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+ indices = indices.narrow(1, 0, P - 2)
+ weights = weights.narrow(1, 0, P - 2)
+ weights = weights.contiguous()
+ indices = indices.contiguous()
+ sym_len_s = -indices.min() + 1
+ sym_len_e = indices.max() - in_length
+ indices = indices + sym_len_s - 1
+ return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+def imresize(img, scale, antialiasing=True):
+ # Now the scale should be the same for H and W
+ # input: img: CHW RGB [0,1]
+ # output: CHW RGB [0,1] w/o round
+
+ in_C, in_H, in_W = img.size()
+ _, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+ kernel_width = 4
+ kernel = 'cubic'
+
+ # Return the desired dimension order for performing the resize. The
+ # strategy is to perform the resize first along the dimension with the
+ # smallest scale factor.
+ # Now we do not support this.
+
+ # get weights and indices
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
+ # process H dimension
+ # symmetric copying
+ img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
+ img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
+
+ sym_patch = img[:, :sym_len_Hs, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+ sym_patch = img[:, -sym_len_He:, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+ out_1 = torch.FloatTensor(in_C, out_H, in_W)
+ kernel_width = weights_H.size(1)
+ for i in range(out_H):
+ idx = int(indices_H[i][0])
+ out_1[0, i, :] = img_aug[0, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+ out_1[1, i, :] = img_aug[1, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+ out_1[2, i, :] = img_aug[2, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+
+ # process W dimension
+ # symmetric copying
+ out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
+ out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
+
+ sym_patch = out_1[:, :, :sym_len_Ws]
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
+ out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+ sym_patch = out_1[:, :, -sym_len_We:]
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
+ out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+ out_2 = torch.FloatTensor(in_C, out_H, out_W)
+ kernel_width = weights_W.size(1)
+ for i in range(out_W):
+ idx = int(indices_W[i][0])
+ out_2[0, :, i] = out_1_aug[0, :, idx:idx + kernel_width].mv(weights_W[i])
+ out_2[1, :, i] = out_1_aug[1, :, idx:idx + kernel_width].mv(weights_W[i])
+ out_2[2, :, i] = out_1_aug[2, :, idx:idx + kernel_width].mv(weights_W[i])
+
+ return out_2
+
+
+def imresize_np(img, scale, antialiasing=True):
+ # Now the scale should be the same for H and W
+ # input: img: Numpy, HWC BGR [0,1]
+ # output: HWC BGR [0,1] w/o round
+ img = torch.from_numpy(img)
+
+ in_H, in_W, in_C = img.size()
+ _, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+ kernel_width = 4
+ kernel = 'cubic'
+
+ # Return the desired dimension order for performing the resize. The
+ # strategy is to perform the resize first along the dimension with the
+ # smallest scale factor.
+ # Now we do not support this.
+
+ # get weights and indices
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
+ # process H dimension
+ # symmetric copying
+ img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
+ img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
+
+ sym_patch = img[:sym_len_Hs, :, :]
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
+ img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+ sym_patch = img[-sym_len_He:, :, :]
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
+ img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+ out_1 = torch.FloatTensor(out_H, in_W, in_C)
+ kernel_width = weights_H.size(1)
+ for i in range(out_H):
+ idx = int(indices_H[i][0])
+ out_1[i, :, 0] = img_aug[idx:idx + kernel_width, :, 0].transpose(0, 1).mv(weights_H[i])
+ out_1[i, :, 1] = img_aug[idx:idx + kernel_width, :, 1].transpose(0, 1).mv(weights_H[i])
+ out_1[i, :, 2] = img_aug[idx:idx + kernel_width, :, 2].transpose(0, 1).mv(weights_H[i])
+
+ # process W dimension
+ # symmetric copying
+ out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
+ out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
+
+ sym_patch = out_1[:, :sym_len_Ws, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+ sym_patch = out_1[:, -sym_len_We:, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+ out_2 = torch.FloatTensor(out_H, out_W, in_C)
+ kernel_width = weights_W.size(1)
+ for i in range(out_W):
+ idx = int(indices_W[i][0])
+ out_2[:, i, 0] = out_1_aug[:, idx:idx + kernel_width, 0].mv(weights_W[i])
+ out_2[:, i, 1] = out_1_aug[:, idx:idx + kernel_width, 1].mv(weights_W[i])
+ out_2[:, i, 2] = out_1_aug[:, idx:idx + kernel_width, 2].mv(weights_W[i])
+
+ return out_2.numpy()
+
+
+ ### Load data kernel map ###
+def load_ker_map_list(path):
+ real_ker_map_list = []
+ batch_kermap = torch.load(path)
+ size_kermap = batch_kermap.size()
+ m = size_kermap[0]
+ for i in range(m):
+ real_ker_map_list.append(batch_kermap[i])
+
+ return real_ker_map_list
+
+
+def test_patchwise(model, L, patchsize=48, overlapsize=16, sf=1):
+ '''
+ Args:
+ model: deep model
+ L: input low-quality image, bxcxhxw
+ patchsize:
+ overlapsize:
+ sf: scale factor, 1 for denoising
+ Return:
+ E: output estimated image, bxcx(h*sf)x(w*sf)
+
+ see another version for segmentation here: https://github.com/fudan-zvg/SETR/blob/main/mmseg/models/segmentors/encoder_decoder.py#L169
+ '''
+ b, c, h, w = L.size()
+ stride = patchsize-overlapsize
+ h_idx_list = list(range(0, h-patchsize, stride)) + [h-patchsize]
+ w_idx_list = list(range(0, w-patchsize, stride)) + [w-patchsize]
+ E = torch.zeros(b, c, h*sf, w*sf).type_as(L)
+ W = torch.zeros_like(E)
+ for h_idx in h_idx_list:
+ for w_idx in w_idx_list:
+ in_patch = L[..., h_idx:h_idx+patchsize, w_idx:w_idx+patchsize]
+ out_patch = model(in_patch)
+ E[..., h_idx*sf:(h_idx+patchsize)*sf, w_idx*sf:(w_idx+patchsize)*sf].add_(out_patch)
+ W[..., h_idx*sf:(h_idx+patchsize)*sf, w_idx*sf:(w_idx+patchsize)*sf].add_(1)
+ return E.div_(W)
+
+
+if __name__ == '__main__':
+ # test imresize function
+ # read images
+ img = cv2.imread('test.png')
+ img = img * 1.0 / 255
+ img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+ # imresize
+ scale = 1 / 4
+ import time
+ total_time = 0
+ for i in range(10):
+ start_time = time.time()
+ rlt = imresize(img, scale, antialiasing=True)
+ use_time = time.time() - start_time
+ total_time += use_time
+ print('average time: {}'.format(total_time / 10))
+
+ import torchvision.utils
+ torchvision.utils.save_image((rlt * 255).round() / 255, 'rlt.png', nrow=1, padding=0,
+ normalize=False)
diff --git a/codes/models/HCFlow_Rescaling_model.py b/codes/models/HCFlow_Rescaling_model.py
new file mode 100644
index 0000000..20d3c04
--- /dev/null
+++ b/codes/models/HCFlow_Rescaling_model.py
@@ -0,0 +1,521 @@
+# base model for SRFlow, a flow-based model
+import logging
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from models.modules.loss import GANLoss
+from .base_model import BaseModel
+import utils.util as util
+from models.modules.loss import ReconstructionLoss
+from models.modules.Basic import Quantization
+# from torch.cuda.amp import autocast as autocast
+
+logger = logging.getLogger('base')
+
+
+class HCFLowRescalingModel(BaseModel):
+ def __init__(self, opt, step):
+ super(HCFLowRescalingModel, self).__init__(opt)
+ self.opt = opt
+
+ self.hr_size = util.opt_get(opt, ['datasets', 'train', 'GT_size'], 160)
+ self.lr_size = self.hr_size // opt['scale']
+
+ if opt['dist']:
+ self.rank = torch.distributed.get_rank()
+ else:
+ self.rank = -1 # non dist training
+
+ # define network and load pretrained models
+ self.netG = networks.define_G(opt, step).to(self.device)
+ if opt['dist']:
+ self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netG = DataParallel(self.netG)
+
+ self.Quantization = Quantization()
+
+ if self.is_train:
+ train_opt = opt['train']
+ self.netG.train()
+
+ # NLL weight
+ # self.cri_pix_hr = ReconstructionLoss(losstype=train_opt['pixel_criterion_hr'])
+ # self.cri_pix_lr = ReconstructionLoss(losstype=train_opt['pixel_criterion_lr'])
+ self.eps_std_reverse = train_opt['eps_std_reverse']
+
+ if train_opt['pixel_weight_hr'] > 0:
+ loss_type = train_opt['pixel_criterion_hr']
+ if loss_type == 'l1':
+ self.cri_pix_hr = nn.L1Loss().to(self.device)
+ elif loss_type == 'l2':
+ self.cri_pix_hr = nn.MSELoss().to(self.device)
+ else:
+ raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+ self.l_pix_w_hr = train_opt['pixel_weight_hr']
+ else:
+ logger.info('Remove HR pixel loss.')
+ self.cri_pix_hr = None
+
+ if train_opt['pixel_weight_lr'] > 0:
+ loss_type = train_opt['pixel_criterion_lr']
+ if loss_type == 'l1':
+ self.cri_pix_lr = nn.L1Loss().to(self.device)
+ elif loss_type == 'l2':
+ self.cri_pix_lr = nn.MSELoss().to(self.device)
+ else:
+ raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+ self.l_pix_w_lr = train_opt['pixel_weight_lr']
+ else:
+ logger.info('Remove LR pixel loss.')
+ self.cri_pix_lr = None
+
+ self.l_pix_w_hr = train_opt['pixel_weight_hr']
+ self.l_pix_w_lr = train_opt['pixel_weight_lr']
+ self.l_w_z = train_opt['weight_z']
+
+ # HR feature loss
+ if train_opt['feature_weight'] > 0:
+ l_fea_type = train_opt['feature_criterion']
+ if l_fea_type == 'l1':
+ self.cri_fea = nn.L1Loss().to(self.device)
+ elif l_fea_type == 'l2':
+ self.cri_fea = nn.MSELoss().to(self.device)
+ else:
+ raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
+ self.l_fea_w = train_opt['feature_weight']
+
+ # load VGG perceptual loss
+ self.netF = networks.define_F(opt, use_bn=False).to(self.device)
+ if opt['dist']:
+ self.netF = DistributedDataParallel(self.netF, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netF = DataParallel(self.netF)
+ else:
+ logger.info('Remove feature loss.')
+ self.cri_fea = None
+
+ # HR GAN loss
+ # put here to be compatible with PSNR version
+ self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1
+ self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0
+ if train_opt['gan_weight'] > 0:
+ self.cri_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
+ self.l_gan_w = train_opt['gan_weight']
+
+ # define GAN Discriminator
+ self.netD = networks.define_D(opt).to(self.device)
+ if opt['dist']:
+ self.netD = DistributedDataParallel(self.netD, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netD = DataParallel(self.netD)
+ self.netD.train()
+ else:
+ logger.info('Remove GAN loss.')
+ self.cri_gan = None
+
+ # gradient clip & norm
+ self.max_grad_clip = util.opt_get(train_opt, ['max_grad_clip'])
+ self.max_grad_norm = util.opt_get(train_opt, ['max_grad_norm'])
+
+ # optimizers
+ # G
+ wd_G = util.opt_get(train_opt, ['weight_decay_G'], 0)
+ optim_params = []
+ for k, v in self.netG.named_parameters(): # can optimize for a part of the model
+ if v.requires_grad:
+ # if v.requires_grad and ('additional_flow_steps' in k or 'additional_feature_steps' in k): # fixmainflow
+ # if v.requires_grad and ('additional_flow_steps' in k): # fix mainflowRRDB
+ optim_params.append(v)
+ else:
+ v.requires_grad = False
+ if self.rank <= 0:
+ logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+ self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
+ weight_decay=wd_G,
+ betas=(train_opt['beta1'], train_opt['beta2']))
+ self.optimizers.append(self.optimizer_G)
+
+ # D
+ if self.cri_gan:
+ wd_D = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
+ self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=train_opt['lr_D'],
+ weight_decay=wd_D,
+ betas=(train_opt['beta1_D'], train_opt['beta2_D']))
+ self.optimizers.append(self.optimizer_D)
+
+ # schedulers
+ if train_opt['lr_scheme'] == 'MultiStepLR':
+ for optimizer in self.optimizers:
+ self.schedulers.append(
+ lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+ restarts=train_opt['restarts'],
+ weights=train_opt['restart_weights'],
+ gamma=train_opt['lr_gamma'],
+ clear_state=train_opt['clear_state']))
+ elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+ for optimizer in self.optimizers:
+ self.schedulers.append(
+ lr_scheduler.CosineAnnealingLR_Restart(
+ optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+ restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+ else:
+ print('MultiStepLR learning rate scheme is enough.')
+
+ self.log_dict = OrderedDict()
+
+ # val
+ if 'val' in opt:
+ self.heats = opt['val']['heats']
+ self.n_sample = opt['val']['n_sample']
+ self.sr_mode = opt['val']['sr_mode']
+
+ self.print_network() # print network
+ self.load() # load G and D if needed
+
+ def init_model(self, scale=0.1):
+ # Common practise for initialization.
+ for layer in self.netG.modules():
+ if isinstance(layer, nn.Conv2d):
+ init.kaiming_normal_(layer.weight, a=0, mode='fan_in')
+ layer.weight.data *= scale # for residual block
+ if layer.bias is not None:
+ layer.bias.data.zero_()
+ elif isinstance(layer, nn.Linear):
+ init.kaiming_normal_(layer.weight, a=0, mode='fan_in')
+ layer.weight.data *= scale
+ if layer.bias is not None:
+ layer.bias.data.zero_()
+ elif isinstance(layer, nn.BatchNorm2d):
+ init.constant_(layer.weight, 1)
+ init.constant_(layer.bias.data, 0.0)
+
+ def feed_data(self, data, need_GT=True):
+ self.var_L = data['LQ'].to(self.device) # LQ
+ if need_GT:
+ self.real_H = data['GT'].to(self.device) # GT
+
+ def optimize_parameters(self, step):
+ # special initialization for actnorm; don't initialize when fine-tuned with gan
+ if step < self.opt['network_G']['act_norm_start_step'] and not (self.cri_pix_hr or self.cri_gan):
+ self.set_actnorm_init(inited=False)
+
+ # (1) G
+ fake_H = None
+ if (step % self.D_update_ratio == 0 and step > self.D_init_iters) or (not self.cri_gan):
+ self.optimizer_G.zero_grad()
+
+ fake_LR, fake_z1, fake_z2 = self.netG(hr=self.real_H, lr=self.var_L, u=None, reverse=False)
+ l_g_lr = self.l_pix_w_lr * self.cri_pix_lr(fake_LR, self.var_L)
+ l_g_z = self.l_w_z * (torch.cat([fake_z1.flatten(), fake_z2.flatten()],0)**2).mean()
+
+ fake_LR = self.Quantization(fake_LR)
+ fake_H = self.netG(lr=fake_LR, z=None, u=None, eps_std=self.eps_std_reverse, reverse=True)
+ l_g_hr = self.l_pix_w_hr * self.cri_pix_hr(fake_H, self.real_H)
+
+ l_g_total = 0
+ if torch.isfinite(l_g_lr):
+ l_g_total += l_g_lr
+ if torch.isfinite(l_g_z):
+ l_g_total += l_g_z
+ if torch.isfinite(l_g_hr):
+ l_g_total += l_g_hr
+
+ self.log_dict['l_g_lr'] = l_g_lr.item()
+ self.log_dict['l_g_z'] = l_g_z.item()
+ self.log_dict['l_g_hr'] = l_g_hr.item()
+
+
+ ########################
+ # feature loss
+ if self.cri_fea:
+ real_fea = self.netF(self.real_H).detach()
+ fake_fea = self.netF(fake_H)
+ l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
+ if torch.isfinite(l_g_fea):
+ l_g_total += l_g_fea
+ self.log_dict['l_g_fea'] = l_g_fea.item()
+
+ # gan loss
+ if self.cri_gan:
+ for p in self.netD.parameters():
+ p.requires_grad = False
+
+ pred_g_fake = self.netD(fake_H)
+ if self.opt['train']['gan_type'] in ['gan', 'lsgan', 'wgangp']:
+ l_g_gan = self.l_gan_w * self.cri_gan(pred_g_fake, True)
+ elif self.opt['train']['gan_type'] == 'ragan':
+ pred_d_real = self.netD(self.real_H).detach()
+ l_g_gan = self.l_gan_w * ( self.cri_gan(pred_d_real - torch.mean(pred_g_fake), False) +
+ self.cri_gan(pred_g_fake - torch.mean(pred_d_real), True)) / 2
+ if torch.isfinite(l_g_gan):
+ l_g_total += l_g_gan
+ self.log_dict['l_g_gan'] = l_g_gan.item()
+
+ if l_g_total != 0:
+ l_g_total.backward()
+ self.gradient_clip()
+ self.optimizer_G.step()
+
+ # (2) D
+ if self.cri_gan:
+ self.optimizer_G.zero_grad() # can help save memory
+
+ for p in self.netD.parameters():
+ p.requires_grad = True
+
+ # initialize D
+ if fake_H is None:
+ with torch.no_grad():
+ fake_H = self.netG(lr=self.var_L, z=None, u=None, eps_std=self.eps_std_reverse, reverse=True)
+
+ self.optimizer_D.zero_grad()
+ pred_d_real = self.netD(self.real_H)
+ pred_d_fake = self.netD(fake_H.detach()) # detach to avoid BP to G
+ if self.opt['train']['gan_type'] in ['gan', 'lsgan', 'wgangp']:
+ l_d_real = self.cri_gan(pred_d_real, True)
+ l_d_fake = self.cri_gan(pred_d_fake, False)
+ l_d_total = l_d_real + l_d_fake
+ elif self.opt['train']['gan_type'] == 'ragan':
+ l_d_real = self.cri_gan(pred_d_real - torch.mean(pred_d_fake), True)
+ l_d_fake = self.cri_gan(pred_d_fake - torch.mean(pred_d_real), False)
+ l_d_total = (l_d_real + l_d_fake) / 2
+
+ self.log_dict['l_d_real'] = l_d_real.item()
+ self.log_dict['l_d_fake'] = l_d_fake.item()
+ self.log_dict['D_real'] = torch.mean(pred_d_real.detach())
+ self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+
+ if torch.isfinite(l_d_total):
+ l_d_total.backward()
+ self.optimizer_D.step()
+
+ def gradient_clip(self):
+ # gradient clip & norm, is not used in SRFlow
+ if self.max_grad_clip is not None:
+ torch.nn.utils.clip_grad_value_(self.netG.parameters(), self.max_grad_clip)
+ if self.max_grad_norm is not None:
+ torch.nn.utils.clip_grad_norm_(self.netG.parameters(), self.max_grad_norm)
+
+ def test(self):
+ self.netG.eval()
+ self.fake_H = {}
+
+ with torch.no_grad():
+ # hr->lr+z, calculate nll
+ self.fake_L_from_H, fake_z1, fake_z2 = self.netG(hr=self.real_H, lr=self.var_L, u=None, reverse=False, training=False)
+ self.fake_L_from_H = self.Quantization(self.fake_L_from_H)
+
+ # lr+z->hr
+ for heat in self.heats:
+ for sample in range(self.n_sample):
+ # z = self.get_z(heat, seed=1, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+ self.fake_H[(heat, sample)] = self.netG(lr=self.fake_L_from_H,
+ z=None, u=None, eps_std=heat, reverse=True, training=False)
+
+ self.netG.train()
+
+ return fake_z1.mean().item()
+
+ def get_encode_nll(self, lq, hr, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ _, nll, _ = self.netG(hr=hr, lr=lq, reverse=False, y_label=y_label)
+ self.netG.train()
+ return nll.mean().item()
+
+ def get_sr(self, lq, heat=None, seed=None, z=None, epses=None, y_label=None):
+ return self.get_sr_with_z(lq, heat, seed, z, epses, y_label=y_label)[0]
+
+ def get_encode_z(self, lq, hr, epses=None, add_gt_noise=True, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ z, _, _ = self.netG(hr=hr, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise, y_label=y_label)
+ self.netG.train()
+ return z
+
+ def get_encode_z_and_nll(self, lq, hr, epses=None, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ z, nll, _ = self.netG(hr=hr, lr=lq, reverse=False, epses=epses, y_label=y_label)
+ self.netG.train()
+ return z, nll
+
+ def get_sr_with_z(self, lq, heat=None, seed=None, z=None, epses=None, y_label=None):
+ self.netG.eval()
+
+ z = self.get_z(heat, seed, batch_size=lq.shape[0], lr_shape=lq.shape,
+ y_label=None) if z is None and epses is None else z
+
+ with torch.no_grad():
+ sr, logdet = self.netG(lr=lq, z=z, eps_std=heat, reverse=True, epses=epses, y_label=None)
+ self.netG.train()
+ return sr, z
+
+ def get_z(self, heat, seed=None, batch_size=1, lr_shape=None, y_label=None):
+ if y_label is None:
+ pass
+ if seed: torch.manual_seed(seed)
+ if util.opt_get(self.opt, ['network_G', 'flowDownsampler', 'splitOff', 'enable']):
+ C, H, W = lr_shape[1], lr_shape[2], lr_shape[3]
+
+ size = (batch_size, C, H, W)
+ if heat == 0:
+ z = torch.zeros(size)
+ else:
+ z = torch.normal(mean=0, std=heat, size=size)
+ else:
+ L = util.opt_get(self.opt, ['network_G', 'flow', 'L']) or 3
+ fac = 2 ** (L - 3)
+ z_size = int(self.lr_size // (2 ** (L - 3)))
+ z = torch.normal(mean=0, std=heat, size=(batch_size, 3 * 8 * 8 * fac * fac, z_size, z_size))
+ return z.to(self.device)
+
+ def get_current_log(self):
+ return self.log_dict
+
+ def get_current_visuals(self, need_GT=True):
+ out_dict = OrderedDict()
+ out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+ for heat in self.heats:
+ for i in range(self.n_sample):
+ out_dict[('SR', heat, i)] = self.fake_H[(heat, i)].detach()[0].float().cpu()
+
+ if need_GT:
+ out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+ out_dict['LQ_fromH'] = self.fake_L_from_H.detach()[0].float().cpu()
+
+ return out_dict
+
+ def print_network(self):
+ s, n = self.get_network_description(self.netG)
+ if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+ self.netG.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netG.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+ logger.info(s)
+
+ if self.is_train:
+ # Discriminator
+ if self.cri_gan:
+ s, n = self.get_network_description(self.netD)
+ if isinstance(self.netD, nn.DataParallel) or isinstance(self.netD,
+ DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netD.__class__.__name__,
+ self.netD.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netD.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network D structure: {}, with parameters: {:,d}'.format(
+ net_struc_str, n))
+ logger.info(s)
+
+ # F, Perceptual Network
+ if self.cri_fea:
+ s, n = self.get_network_description(self.netF)
+ if isinstance(self.netF, nn.DataParallel) or isinstance(
+ self.netF, DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netF.__class__.__name__,
+ self.netF.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netF.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network F structure: {}, with parameters: {:,d}'.format(
+ net_struc_str, n))
+ logger.info(s)
+
+ def load(self):
+ # resume training automatically if resume_state=='auto'
+ _, get_resume_model_path = util.get_resume_paths(self.opt)
+ if get_resume_model_path is not None:
+ logger.info('Automatically loading model for G [{:s}] ...'.format(get_resume_model_path))
+ self.load_network(get_resume_model_path, self.netG, strict=True, submodule=None)
+ self.set_actnorm_init(inited=True)
+
+ if self.is_train and self.cri_gan:
+ get_resume_model_path = get_resume_model_path.replace('_G.pth', '_D.pth')
+ logger.info('Automatically loading model for D [{:s}] ...'.format(get_resume_model_path))
+ self.load_network(get_resume_model_path, self.netD, strict=True, submodule=None)
+ return
+
+ # resume training according to given paths (pretrain path has been overrided by resume path)
+ if self.opt.get('path') is not None:
+ load_path_G = self.opt['path']['pretrain_model_G']
+ if load_path_G is not None:
+ logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+ self.load_network(load_path_G, self.netG, self.opt['path'].get('strict_load', True))
+ self.set_actnorm_init(inited=True)
+
+ if self.is_train and self.cri_gan:
+ load_path_D = self.opt['path']['pretrain_model_D']
+ if load_path_D is not None:
+ logger.info('Loading model for D [{:s}] ...'.format(load_path_D))
+ self.load_network(load_path_D, self.netD, self.opt['path'].get('strict_load', True))
+
+ def save(self, iter_label):
+ self.save_network(self.netG, 'G', iter_label)
+ if self.cri_gan:
+ self.save_network(self.netD, 'D', iter_label)
+
+ def set_actnorm_init(self, inited=True):
+ for name, m in self.netG.named_modules():
+ if (m.__class__.__name__.find("ActNorm") >= 0):
+ m.inited = inited
+
+ # from glow
+ def generate_z(self, img):
+ self.eval()
+ B = self.hparams.Train.batch_size
+ x = img.unsqueeze(0).repeat(B, 1, 1, 1).cuda()
+ z,_, _ = self(x)
+ self.train()
+ return z[0].detach().cpu().numpy()
+
+ def generate_attr_deltaz(self, dataset):
+ assert "y_onehot" in dataset[0]
+ self.eval()
+ with torch.no_grad():
+ B = self.hparams.Train.batch_size
+ N = len(dataset)
+ attrs_pos_z = [[0, 0] for _ in range(self.y_classes)]
+ attrs_neg_z = [[0, 0] for _ in range(self.y_classes)]
+ for i in tqdm(range(0, N, B)):
+ j = min([i + B, N])
+ # generate z for data from i to j
+ xs = [dataset[k]["x"] for k in range(i, j)]
+ while len(xs) < B:
+ xs.append(dataset[0]["x"])
+ xs = torch.stack(xs).cuda()
+ zs, _, _ = self(xs)
+ for k in range(i, j):
+ z = zs[k - i].detach().cpu().numpy()
+ # append to different attrs
+ y = dataset[k]["y_onehot"]
+ for ai in range(self.y_classes):
+ if y[ai] > 0:
+ attrs_pos_z[ai][0] += z
+ attrs_pos_z[ai][1] += 1
+ else:
+ attrs_neg_z[ai][0] += z
+ attrs_neg_z[ai][1] += 1
+ # break
+ deltaz = []
+ for ai in range(self.y_classes):
+ if attrs_pos_z[ai][1] == 0:
+ attrs_pos_z[ai][1] = 1
+ if attrs_neg_z[ai][1] == 0:
+ attrs_neg_z[ai][1] = 1
+ z_pos = attrs_pos_z[ai][0] / float(attrs_pos_z[ai][1])
+ z_neg = attrs_neg_z[ai][0] / float(attrs_neg_z[ai][1])
+ deltaz.append(z_pos - z_neg)
+ self.train()
+ return deltaz
\ No newline at end of file
diff --git a/codes/models/HCFlow_SR_model.py b/codes/models/HCFlow_SR_model.py
new file mode 100644
index 0000000..7ed42e4
--- /dev/null
+++ b/codes/models/HCFlow_SR_model.py
@@ -0,0 +1,510 @@
+# base model for HCFlow
+import logging
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from models.modules.loss import GANLoss
+from .base_model import BaseModel
+import utils.util as util
+
+logger = logging.getLogger('base')
+
+
+class HCFlowSRModel(BaseModel):
+ def __init__(self, opt, step):
+ super(HCFlowSRModel, self).__init__(opt)
+ self.opt = opt
+
+ self.hr_size = util.opt_get(opt, ['datasets', 'train', 'GT_size'], 160)
+ self.lr_size = self.hr_size // opt['scale']
+
+ if opt['dist']:
+ self.rank = torch.distributed.get_rank()
+ else:
+ self.rank = -1 # non dist training
+
+ # define network and load pretrained models
+ self.netG = networks.define_G(opt, step).to(self.device)
+ if opt['dist']:
+ self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netG = DataParallel(self.netG)
+
+ if self.is_train:
+ train_opt = opt['train']
+ self.netG.train()
+
+ # NLL weight
+ self.l_nll_w = train_opt['nll_weight']
+ self.eps_std_reverse = train_opt['eps_std_reverse']
+
+ # HR pixel loss
+ if train_opt['pixel_weight_hr'] > 0:
+ loss_type = train_opt['pixel_criterion_hr']
+ if loss_type == 'l1':
+ self.cri_pix_hr = nn.L1Loss().to(self.device)
+ elif loss_type == 'l2':
+ self.cri_pix_hr = nn.MSELoss().to(self.device)
+ else:
+ raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+ self.l_pix_w_hr = train_opt['pixel_weight_hr']
+ else:
+ logger.info('Remove HR pixel loss.')
+ self.cri_pix_hr = None
+
+ # HR feature loss
+ if train_opt['feature_weight'] > 0:
+ l_fea_type = train_opt['feature_criterion']
+ if l_fea_type == 'l1':
+ self.cri_fea = nn.L1Loss().to(self.device)
+ elif l_fea_type == 'l2':
+ self.cri_fea = nn.MSELoss().to(self.device)
+ else:
+ raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
+ self.l_fea_w = train_opt['feature_weight']
+
+ # load VGG perceptual loss
+ self.netF = networks.define_F(opt, use_bn=False).to(self.device)
+ if opt['dist']:
+ self.netF = DistributedDataParallel(self.netF, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netF = DataParallel(self.netF)
+ else:
+ logger.info('Remove feature loss.')
+ self.cri_fea = None
+
+ # HR GAN loss
+ # put here to be compatible with PSNR version
+ self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1
+ self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0
+ if train_opt['gan_weight'] > 0:
+ self.cri_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
+ self.l_gan_w = train_opt['gan_weight']
+
+ # define GAN Discriminator
+ self.netD = networks.define_D(opt).to(self.device)
+ if opt['dist']:
+ self.netD = DistributedDataParallel(self.netD, device_ids=[torch.cuda.current_device()])
+ else:
+ self.netD = DataParallel(self.netD)
+ self.netD.train()
+ else:
+ logger.info('Remove GAN loss.')
+ self.cri_gan = None
+
+ # gradient clip & norm
+ self.max_grad_clip = util.opt_get(train_opt, ['max_grad_clip'])
+ self.max_grad_norm = util.opt_get(train_opt, ['max_grad_norm'])
+
+ # optimizers
+ # G
+ wd_G = util.opt_get(train_opt, ['weight_decay_G'], 0)
+ optim_params = []
+ for k, v in self.netG.named_parameters(): # can optimize for a part of the model
+ if v.requires_grad:
+ # if v.requires_grad and ('additional_flow_steps' in k or 'additional_feature_steps' in k): # fixmainflow
+ # if v.requires_grad and ('additional_flow_steps' in k): # fix mainflowRRDB
+ optim_params.append(v)
+ else:
+ v.requires_grad = False
+ if self.rank <= 0:
+ logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+ self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
+ weight_decay=wd_G,
+ betas=(train_opt['beta1'], train_opt['beta2']))
+ self.optimizers.append(self.optimizer_G)
+
+ # D
+ if self.cri_gan:
+ wd_D = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
+ self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=train_opt['lr_D'],
+ weight_decay=wd_D,
+ betas=(train_opt['beta1_D'], train_opt['beta2_D']))
+ self.optimizers.append(self.optimizer_D)
+
+ # schedulers
+ if train_opt['lr_scheme'] == 'MultiStepLR':
+ for optimizer in self.optimizers:
+ self.schedulers.append(
+ lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+ restarts=train_opt['restarts'],
+ weights=train_opt['restart_weights'],
+ gamma=train_opt['lr_gamma'],
+ clear_state=train_opt['clear_state']))
+ elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+ for optimizer in self.optimizers:
+ self.schedulers.append(
+ lr_scheduler.CosineAnnealingLR_Restart(
+ optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+ restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+ else:
+ print('MultiStepLR learning rate scheme is enough.')
+
+ self.log_dict = OrderedDict()
+
+ # val
+ if 'val' in opt:
+ self.heats = opt['val']['heats']
+ self.n_sample = opt['val']['n_sample']
+ self.sr_mode = opt['val']['sr_mode']
+
+ self.print_network() # print network
+ self.load() # load G and D if needed
+
+ def init_model(self, scale=0.1):
+ # Common practise for initialization.
+ for layer in self.netG.modules():
+ if isinstance(layer, nn.Conv2d):
+ init.kaiming_normal_(layer.weight, a=0, mode='fan_in')
+ layer.weight.data *= scale # for residual block
+ if layer.bias is not None:
+ layer.bias.data.zero_()
+ elif isinstance(layer, nn.Linear):
+ init.kaiming_normal_(layer.weight, a=0, mode='fan_in')
+ layer.weight.data *= scale
+ if layer.bias is not None:
+ layer.bias.data.zero_()
+ elif isinstance(layer, nn.BatchNorm2d):
+ init.constant_(layer.weight, 1)
+ init.constant_(layer.bias.data, 0.0)
+
+ def feed_data(self, data, need_GT=True):
+ self.var_L = data['LQ'].to(self.device) # LQ
+ if need_GT:
+ self.real_H = data['GT'].to(self.device) # GT
+
+ def optimize_parameters(self, step):
+ # special initialization for actnorm; don't initialize when fine-tuning
+ if step < self.opt['network_G']['act_norm_start_step'] and not (self.cri_pix_hr or self.cri_gan):
+ self.set_actnorm_init(inited=False)
+
+ # (1) G
+ fake_H = None
+ if (step % self.D_update_ratio == 0 and step > self.D_init_iters) or (not self.cri_gan):
+ # normal flow
+ l_g_total = 0
+
+ _, nll = self.netG(hr=self.real_H, lr=self.var_L, u=None, reverse=False)
+ nll = self.l_nll_w * nll.sum()
+ self.log_dict['nll'] = nll.item()
+ if not torch.isnan(nll).any():
+ l_g_total += nll
+ if l_g_total != 0:
+ l_g_total.backward()
+ self.gradient_clip()
+ self.optimizer_G.step()
+
+ # reverse flow (optimize NLL loss and HR loss seperately (takes less memory and more time, slightly better results))
+ self.optimizer_G.zero_grad()
+ if self.cri_pix_hr:
+ fake_H = self.netG(lr=self.var_L, z=None, u=None, eps_std=0.0, reverse=True)
+ l_g_total = 0
+ if not torch.isnan(fake_H).any():
+ # pixel loss
+ l_g_pix_hr = self.l_pix_w_hr * self.cri_pix_hr(fake_H, self.real_H)
+ l_g_total += l_g_pix_hr
+ self.log_dict['l_g_pix_hr'] = l_g_pix_hr.item()
+ if l_g_total != 0:
+ l_g_total.backward()
+ self.gradient_clip()
+ self.optimizer_G.step()
+
+
+ ########################
+
+ if self.cri_gan or self.cri_fea:
+ self.optimizer_G.zero_grad()
+ fake_H = self.netG(lr=self.var_L, z=None, u=None, eps_std=self.eps_std_reverse, reverse=True)
+ l_g_fea_gan = 0
+
+ # feature loss
+ if self.cri_fea:
+ real_fea = self.netF(self.real_H).detach()
+ fake_fea = self.netF(fake_H)
+ l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
+ l_g_fea_gan += l_g_fea
+ self.log_dict['l_g_fea'] = l_g_fea.item()
+
+ # gan loss
+ if self.cri_gan:
+ for p in self.netD.parameters():
+ p.requires_grad = False
+
+ pred_g_fake = self.netD(fake_H)
+ if self.opt['train']['gan_type'] in ['gan', 'lsgan', 'wgangp']:
+ l_g_gan = self.l_gan_w * self.cri_gan(pred_g_fake, True)
+ elif self.opt['train']['gan_type'] == 'ragan':
+ pred_d_real = self.netD(self.real_H).detach()
+ l_g_gan = self.l_gan_w * ( self.cri_gan(pred_d_real - torch.mean(pred_g_fake), False) +
+ self.cri_gan(pred_g_fake - torch.mean(pred_d_real), True)) / 2
+ l_g_fea_gan += l_g_gan
+ self.log_dict['l_g_gan'] = l_g_gan.item()
+
+ if not torch.isnan(l_g_fea_gan):
+ l_g_fea_gan.backward()
+ self.gradient_clip()
+ self.optimizer_G.step()
+
+ # (2) D
+ if self.cri_gan:
+ self.optimizer_G.zero_grad() # can help save memory
+
+ for p in self.netD.parameters():
+ p.requires_grad = True
+
+ # initialize D
+ if fake_H is None:
+ with torch.no_grad():
+ fake_H = self.netG(lr=self.var_L, z=None, u=None, eps_std=self.eps_std_reverse, reverse=True)
+
+ self.optimizer_D.zero_grad()
+ pred_d_real = self.netD(self.real_H)
+ pred_d_fake = self.netD(fake_H.detach()) # detach to avoid BP to G
+ if self.opt['train']['gan_type'] in ['gan', 'lsgan', 'wgangp']:
+ l_d_real = self.cri_gan(pred_d_real, True)
+ l_d_fake = self.cri_gan(pred_d_fake, False)
+ l_d_total = l_d_real + l_d_fake
+ elif self.opt['train']['gan_type'] == 'ragan':
+ l_d_real = self.cri_gan(pred_d_real - torch.mean(pred_d_fake), True)
+ l_d_fake = self.cri_gan(pred_d_fake - torch.mean(pred_d_real), False)
+ l_d_total = (l_d_real + l_d_fake) / 2
+
+ self.log_dict['l_d_real'] = l_d_real.item()
+ self.log_dict['l_d_fake'] = l_d_fake.item()
+ self.log_dict['D_real'] = torch.mean(pred_d_real.detach())
+ self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+
+ if not torch.isnan(l_d_total):
+ l_d_total.backward()
+ self.optimizer_D.step()
+
+ def gradient_clip(self):
+ # gradient clip & norm, is not used in SRFlow
+ if self.max_grad_clip is not None:
+ torch.nn.utils.clip_grad_value_(self.netG.parameters(), self.max_grad_clip)
+ if self.max_grad_norm is not None:
+ torch.nn.utils.clip_grad_norm_(self.netG.parameters(), self.max_grad_norm)
+
+ def test(self):
+ self.netG.eval()
+ self.fake_H = {}
+
+ with torch.no_grad():
+ # hr->lr+z, calculate nll
+ self.fake_L_from_H, nll = self.netG(hr=self.real_H, lr=self.var_L, u=None, reverse=False, training=False)
+ # nll = torch.zeros(1)
+
+ # lr+z->hr
+ for heat in self.heats:
+ for sample in range(self.n_sample):
+ # z = self.get_z(heat, seed=1, batch_size=self.var_L.shape[0], lr_shape=self.var_L.shape)
+ self.fake_H[(heat, sample)] = self.netG(lr=self.var_L,
+ z=None, u=None, eps_std=heat, reverse=True, training=False)
+
+ self.netG.train()
+
+ return nll.mean().item()
+
+ def get_encode_nll(self, lq, hr, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ _, nll, _ = self.netG(hr=hr, lr=lq, reverse=False, y_label=y_label)
+ self.netG.train()
+ return nll.mean().item()
+
+ def get_sr(self, lq, heat=None, seed=None, z=None, epses=None, y_label=None):
+ return self.get_sr_with_z(lq, heat, seed, z, epses, y_label=y_label)[0]
+
+ def get_encode_z(self, lq, hr, epses=None, add_gt_noise=True, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ z, _, _ = self.netG(hr=hr, lr=lq, reverse=False, epses=epses, add_gt_noise=add_gt_noise, y_label=y_label)
+ self.netG.train()
+ return z
+
+ def get_encode_z_and_nll(self, lq, hr, epses=None, y_label=None):
+ self.netG.eval()
+ with torch.no_grad():
+ z, nll, _ = self.netG(hr=hr, lr=lq, reverse=False, epses=epses, y_label=y_label)
+ self.netG.train()
+ return z, nll
+
+ def get_sr_with_z(self, lq, heat=None, seed=None, z=None, epses=None, y_label=None):
+ self.netG.eval()
+
+ z = self.get_z(heat, seed, batch_size=lq.shape[0], lr_shape=lq.shape,
+ y_label=None) if z is None and epses is None else z
+
+ with torch.no_grad():
+ sr, logdet = self.netG(lr=lq, z=z, eps_std=heat, reverse=True, epses=epses, y_label=None)
+ self.netG.train()
+ return sr, z
+
+ def get_z(self, heat, seed=None, batch_size=1, lr_shape=None, y_label=None):
+ if y_label is None:
+ pass
+ if seed: torch.manual_seed(seed)
+ if util.opt_get(self.opt, ['network_G', 'flowLR', 'splitOff', 'enable']):
+ C, H, W = lr_shape[1], lr_shape[2], lr_shape[3]
+
+ size = (batch_size, C, H, W)
+ if heat == 0:
+ z = torch.zeros(size)
+ else:
+ z = torch.normal(mean=0, std=heat, size=size)
+ else:
+ L = util.opt_get(self.opt, ['network_G', 'flow', 'L']) or 3
+ fac = 2 ** (L - 3)
+ z_size = int(self.lr_size // (2 ** (L - 3)))
+ z = torch.normal(mean=0, std=heat, size=(batch_size, 3 * 8 * 8 * fac * fac, z_size, z_size))
+ return z.to(self.device)
+
+ def get_current_log(self):
+ return self.log_dict
+
+ def get_current_visuals(self, need_GT=True):
+ out_dict = OrderedDict()
+ out_dict['LQ'] = self.var_L.detach()[0].float().cpu()
+ for heat in self.heats:
+ for i in range(self.n_sample):
+ out_dict[('SR', heat, i)] = self.fake_H[(heat, i)].detach()[0].float().cpu()
+
+ if need_GT:
+ out_dict['GT'] = self.real_H.detach()[0].float().cpu()
+ out_dict['LQ_fromH'] = self.fake_L_from_H.detach()[0].float().cpu()
+
+ return out_dict
+
+ def print_network(self):
+ s, n = self.get_network_description(self.netG)
+ if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+ self.netG.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netG.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+ logger.info(s)
+
+ if self.is_train:
+ # Discriminator
+ if self.cri_gan:
+ s, n = self.get_network_description(self.netD)
+ if isinstance(self.netD, nn.DataParallel) or isinstance(self.netD,
+ DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netD.__class__.__name__,
+ self.netD.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netD.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network D structure: {}, with parameters: {:,d}'.format(
+ net_struc_str, n))
+ logger.info(s)
+
+ # F, Perceptual Network
+ if self.cri_fea:
+ s, n = self.get_network_description(self.netF)
+ if isinstance(self.netF, nn.DataParallel) or isinstance(
+ self.netF, DistributedDataParallel):
+ net_struc_str = '{} - {}'.format(self.netF.__class__.__name__,
+ self.netF.module.__class__.__name__)
+ else:
+ net_struc_str = '{}'.format(self.netF.__class__.__name__)
+ if self.rank <= 0:
+ logger.info('Network F structure: {}, with parameters: {:,d}'.format(
+ net_struc_str, n))
+ logger.info(s)
+
+ def load(self):
+ # resume training automatically if resume_state=='auto'
+ _, get_resume_model_path = util.get_resume_paths(self.opt)
+ if get_resume_model_path is not None:
+ logger.info('Automatically loading model for G [{:s}] ...'.format(get_resume_model_path))
+ self.load_network(get_resume_model_path, self.netG, strict=True, submodule=None)
+ self.set_actnorm_init(inited=True)
+
+ if self.is_train and self.cri_gan:
+ get_resume_model_path = get_resume_model_path.replace('_G.pth', '_D.pth')
+ logger.info('Automatically loading model for D [{:s}] ...'.format(get_resume_model_path))
+ self.load_network(get_resume_model_path, self.netD, strict=True, submodule=None)
+ return
+
+ # resume training according to given paths (pretrain path has been overrided by resume path)
+ if self.opt.get('path') is not None:
+ load_path_G = self.opt['path']['pretrain_model_G']
+ if load_path_G is not None:
+ logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+ self.load_network(load_path_G, self.netG, self.opt['path'].get('strict_load', True))
+
+ self.set_actnorm_init(inited=True)
+
+ if self.is_train and self.cri_gan:
+ load_path_D = self.opt['path']['pretrain_model_D']
+ if load_path_D is not None:
+ logger.info('Loading model for D [{:s}] ...'.format(load_path_D))
+ self.load_network(load_path_D, self.netD, self.opt['path'].get('strict_load', True))
+
+ def save(self, iter_label):
+ self.save_network(self.netG, 'G', iter_label)
+ if self.cri_gan:
+ self.save_network(self.netD, 'D', iter_label)
+
+ def set_actnorm_init(self, inited=True):
+ for name, m in self.netG.named_modules():
+ if (m.__class__.__name__.find("ActNorm") >= 0):
+ m.inited = inited
+
+ # from glow
+ def generate_z(self, img):
+ self.eval()
+ B = self.hparams.Train.batch_size
+ x = img.unsqueeze(0).repeat(B, 1, 1, 1).cuda()
+ z,_, _ = self(x)
+ self.train()
+ return z[0].detach().cpu().numpy()
+
+ def generate_attr_deltaz(self, dataset):
+ assert "y_onehot" in dataset[0]
+ self.eval()
+ with torch.no_grad():
+ B = self.hparams.Train.batch_size
+ N = len(dataset)
+ attrs_pos_z = [[0, 0] for _ in range(self.y_classes)]
+ attrs_neg_z = [[0, 0] for _ in range(self.y_classes)]
+ for i in tqdm(range(0, N, B)):
+ j = min([i + B, N])
+ # generate z for data from i to j
+ xs = [dataset[k]["x"] for k in range(i, j)]
+ while len(xs) < B:
+ xs.append(dataset[0]["x"])
+ xs = torch.stack(xs).cuda()
+ zs, _, _ = self(xs)
+ for k in range(i, j):
+ z = zs[k - i].detach().cpu().numpy()
+ # append to different attrs
+ y = dataset[k]["y_onehot"]
+ for ai in range(self.y_classes):
+ if y[ai] > 0:
+ attrs_pos_z[ai][0] += z
+ attrs_pos_z[ai][1] += 1
+ else:
+ attrs_neg_z[ai][0] += z
+ attrs_neg_z[ai][1] += 1
+ # break
+ deltaz = []
+ for ai in range(self.y_classes):
+ if attrs_pos_z[ai][1] == 0:
+ attrs_pos_z[ai][1] = 1
+ if attrs_neg_z[ai][1] == 0:
+ attrs_neg_z[ai][1] = 1
+ z_pos = attrs_pos_z[ai][0] / float(attrs_pos_z[ai][1])
+ z_neg = attrs_neg_z[ai][0] / float(attrs_neg_z[ai][1])
+ deltaz.append(z_pos - z_neg)
+ self.train()
+ return deltaz
diff --git a/codes/models/__init__.py b/codes/models/__init__.py
new file mode 100644
index 0000000..9803258
--- /dev/null
+++ b/codes/models/__init__.py
@@ -0,0 +1,52 @@
+import importlib
+import logging
+import os
+
+try:
+ import local_config
+except:
+ local_config = None
+
+
+logger = logging.getLogger('base')
+
+
+def find_model_using_name(model_name):
+ # Given the option --model [modelname],
+ # the file "models/modelname_model.py"
+ # will be imported.
+ model_filename = "models." + model_name + "_model"
+ modellib = importlib.import_module(model_filename)
+
+ # In the file, the class called ModelNameModel() will
+ # be instantiated. It has to be a subclass of torch.nn.Module,
+ # and it is case-insensitive.
+ model = None
+ target_model_name = model_name.replace('_', '') + 'Model'
+ for name, cls in modellib.__dict__.items():
+ if name.lower() == target_model_name.lower():
+ model = cls
+
+ if model is None:
+ print(
+ "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s." % (
+ model_filename, target_model_name))
+ exit(0)
+
+ return model
+
+
+def create_model(opt, step=0, **opt_kwargs):
+ if local_config is not None:
+ opt['path']['pretrain_model_G'] = os.path.join(local_config.checkpoint_path, os.path.basename(opt['path']['results_root'] + '.pth'))
+
+ for k, v in opt_kwargs.items():
+ opt[k] = v
+
+ model = opt['model']
+
+ M = find_model_using_name(model)
+
+ m = M(opt, step)
+ logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
+ return m
diff --git a/codes/models/base_model.py b/codes/models/base_model.py
new file mode 100644
index 0000000..e5af191
--- /dev/null
+++ b/codes/models/base_model.py
@@ -0,0 +1,166 @@
+import os
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+import natsort
+import glob
+
+
+class BaseModel():
+ def __init__(self, opt):
+ self.opt = opt
+ self.device = torch.device('cuda' if opt.get('gpu_ids', None) is not None else 'cpu')
+ self.is_train = opt['is_train']
+ self.schedulers = []
+ self.optimizers = []
+
+ def feed_data(self, data):
+ pass
+
+ def optimize_parameters(self):
+ pass
+
+ def get_current_visuals(self):
+ pass
+
+ def get_current_losses(self):
+ pass
+
+ def print_network(self):
+ pass
+
+ def save(self, label):
+ pass
+
+ def load(self):
+ pass
+
+ def _set_lr(self, lr_groups_l):
+ ''' set learning rate for warmup,
+ lr_groups_l: list for lr_groups. each for a optimizer'''
+ for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
+ for param_group, lr in zip(optimizer.param_groups, lr_groups):
+ param_group['lr'] = lr
+
+ def _get_init_lr(self):
+ # get the initial lr, which is set by the scheduler
+ init_lr_groups_l = []
+ for optimizer in self.optimizers:
+ init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
+ return init_lr_groups_l
+
+ def update_learning_rate(self, cur_iter, warmup_iter=-1):
+ for scheduler in self.schedulers:
+ scheduler.step()
+ #### set up warm up learning rate
+ if cur_iter < warmup_iter:
+ # get initial lr for each group
+ init_lr_g_l = self._get_init_lr()
+ # modify warming-up learning rates
+ warm_up_lr_l = []
+ for init_lr_g in init_lr_g_l:
+ warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g])
+ # set learning rate
+ self._set_lr(warm_up_lr_l)
+
+ def get_current_learning_rate(self):
+ # return self.schedulers[0].get_lr()[0]
+ return self.optimizers[0].param_groups[0]['lr']
+
+ def get_network_description(self, network):
+ '''Get the string and total parameters of the network'''
+ if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+ network = network.module
+ s = str(network)
+ n = sum(map(lambda x: x.numel(), network.parameters()))
+ return s, n
+
+ def save_network(self, network, network_label, iter_label):
+ paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['models'], "*_{}.pth".format(network_label))),
+ reverse=True)
+ paths = [p for p in paths if
+ "latest_" not in p and not any([str(i * 5000) in p.split("/")[-1].split("_") for i in range(101)])]
+ if len(paths) > 2:
+ for path in paths[2:]:
+ os.remove(path)
+ save_filename = '{}_{}.pth'.format(iter_label, network_label)
+ save_path = os.path.join(self.opt['path']['models'], save_filename)
+ if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+ network = network.module
+ state_dict = network.state_dict()
+ for key, param in state_dict.items():
+ state_dict[key] = param.cpu()
+ torch.save(state_dict, save_path)
+
+ def load_network(self, load_path, network, strict=True, submodule=None):
+ if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+ network = network.module
+ if not (submodule is None or submodule.lower() == 'none'.lower()):
+ network = network.__getattr__(submodule)
+ load_net = torch.load(load_path)
+ load_net_clean = OrderedDict() # remove unnecessary 'module.'
+ for k, v in load_net.items():
+
+ # if 'flowLRNet' in k or 'quantization' in k:
+ # continue
+ # k = k.replace('flowDownpsamplerNet', 'flow')
+ #
+ # # step2
+ # k = k.replace('Split2d712', 'level0_condFlow')
+ # k = k.replace('layers.29', 'level1_condFlow')
+ # # k = k.replace('Split2d932', 'level2_condFlow')
+
+
+ if k.startswith('module.'):
+ load_net_clean[k[7:]] = v
+ else:
+ load_net_clean[k] = v
+
+ network.load_state_dict(load_net_clean, strict=strict)
+
+ # if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+ # network = network.module
+ # state_dict = network.state_dict()
+ # for key, param in state_dict.items():
+ # state_dict[key] = param.cpu()
+ # torch.save(state_dict, '../experiments/SR_DF2K_X4_HCFlow32.pth')
+
+ def save_training_state(self, epoch, iter_step):
+ '''Saves training state during training, which will be used for resuming'''
+ state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []}
+ for s in self.schedulers:
+ state['schedulers'].append(s.state_dict())
+ for o in self.optimizers:
+ state['optimizers'].append(o.state_dict())
+ save_filename = '{}.state'.format(iter_step)
+ save_path = os.path.join(self.opt['path']['training_state'], save_filename)
+
+ paths = natsort.natsorted(glob.glob(os.path.join(self.opt['path']['training_state'], "*.state")),
+ reverse=True)
+ paths = [p for p in paths if "latest_" not in p]
+ if len(paths) > 2:
+ for path in paths[2:]:
+ os.remove(path)
+
+ torch.save(state, save_path)
+
+ def resume_training(self, resume_state):
+ '''Resume the optimizers and schedulers for training'''
+ resume_optimizers = resume_state['optimizers']
+ resume_schedulers = resume_state['schedulers']
+ assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+ assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+ for i, o in enumerate(resume_optimizers):
+ self.optimizers[i].load_state_dict(o)
+ for i, s in enumerate(resume_schedulers):
+ # manually change lr milestones
+ # from collections import Counter
+ # s['milestones'] = Counter([100000, 150000, 200000, 250000, 300000, 350000, 400000]) # for multistage lr
+ # s['restarts'] = [120001, 240001, 360001] # for cosine_restart_lr
+
+ self.schedulers[i].load_state_dict(s)
+
+
+
+
diff --git a/codes/models/lr_scheduler.py b/codes/models/lr_scheduler.py
new file mode 100644
index 0000000..2304f0c
--- /dev/null
+++ b/codes/models/lr_scheduler.py
@@ -0,0 +1,144 @@
+import math
+from collections import Counter
+from collections import defaultdict
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class MultiStepLR_Restart(_LRScheduler):
+ def __init__(self, optimizer, milestones, restarts=None, weights=None, gamma=0.1,
+ clear_state=False, last_epoch=-1):
+ self.milestones = Counter(milestones)
+ self.gamma = gamma
+ self.clear_state = clear_state
+ self.restarts = restarts if restarts else [0]
+ self.restarts = [v + 1 for v in self.restarts]
+ self.restart_weights = weights if weights else [1]
+ assert len(self.restarts) == len(
+ self.restart_weights), 'restarts and their weights do not match.'
+ super(MultiStepLR_Restart, self).__init__(optimizer, last_epoch)
+
+ def get_lr(self):
+ if self.last_epoch in self.restarts:
+ if self.clear_state:
+ self.optimizer.state = defaultdict(dict)
+ weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+ return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+ if self.last_epoch not in self.milestones:
+ return [group['lr'] for group in self.optimizer.param_groups]
+ return [
+ group['lr'] * self.gamma**self.milestones[self.last_epoch]
+ for group in self.optimizer.param_groups
+ ]
+
+
+class CosineAnnealingLR_Restart(_LRScheduler):
+ def __init__(self, optimizer, T_period, restarts=None, weights=None, eta_min=0, last_epoch=-1):
+ self.T_period = T_period
+ self.T_max = self.T_period[0] # current T period
+ self.eta_min = eta_min
+ self.restarts = restarts if restarts else [0]
+ self.restarts = [v + 1 for v in self.restarts]
+ self.restart_weights = weights if weights else [1]
+ self.last_restart = 0
+ assert len(self.restarts) == len(
+ self.restart_weights), 'restarts and their weights do not match.'
+ super(CosineAnnealingLR_Restart, self).__init__(optimizer, last_epoch)
+
+ def get_lr(self):
+ if self.last_epoch == 0:
+ return self.base_lrs
+ elif self.last_epoch in self.restarts:
+ self.last_restart = self.last_epoch
+ self.T_max = self.T_period[self.restarts.index(self.last_epoch) + 1]
+ weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+ return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+ elif (self.last_epoch - self.last_restart - 1 - self.T_max) % (2 * self.T_max) == 0:
+ return [
+ group['lr'] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2
+ for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+ ]
+ return [(1 + math.cos(math.pi * (self.last_epoch - self.last_restart) / self.T_max)) /
+ (1 + math.cos(math.pi * ((self.last_epoch - self.last_restart) - 1) / self.T_max)) *
+ (group['lr'] - self.eta_min) + self.eta_min
+ for group in self.optimizer.param_groups]
+
+
+if __name__ == "__main__":
+ optimizer = torch.optim.Adam([torch.zeros(3, 64, 3, 3)], lr=2e-4, weight_decay=0,
+ betas=(0.9, 0.99))
+ ##############################
+ # MultiStepLR_Restart
+ ##############################
+ ## Original
+ lr_steps = [200000, 400000, 600000, 800000]
+ restarts = None
+ restart_weights = None
+
+ ## two
+ lr_steps = [100000, 200000, 300000, 400000, 490000, 600000, 700000, 800000, 900000, 990000]
+ restarts = [500000]
+ restart_weights = [1]
+
+ ## four
+ lr_steps = [
+ 50000, 100000, 150000, 200000, 240000, 300000, 350000, 400000, 450000, 490000, 550000,
+ 600000, 650000, 700000, 740000, 800000, 850000, 900000, 950000, 990000
+ ]
+ restarts = [250000, 500000, 750000]
+ restart_weights = [1, 1, 1]
+
+ scheduler = MultiStepLR_Restart(optimizer, lr_steps, restarts, restart_weights, gamma=0.5,
+ clear_state=False)
+
+ ##############################
+ # Cosine Annealing Restart
+ ##############################
+ ## two
+ T_period = [500000, 500000]
+ restarts = [500000]
+ restart_weights = [1]
+
+ ## four
+ T_period = [250000, 250000, 250000, 250000]
+ restarts = [250000, 500000, 750000]
+ restart_weights = [1,1, 1]
+
+ scheduler = CosineAnnealingLR_Restart(optimizer, T_period, eta_min=1e-7, restarts=restarts,
+ weights=restart_weights)
+
+ ##############################
+ # Draw figure
+ ##############################
+ N_iter = 1000000
+ lr_l = list(range(N_iter))
+ for i in range(N_iter):
+ scheduler.step()
+ current_lr = optimizer.param_groups[0]['lr']
+ lr_l[i] = current_lr
+
+ import matplotlib as mpl
+ from matplotlib import pyplot as plt
+ import matplotlib.ticker as mtick
+ mpl.style.use('default')
+ import seaborn
+ seaborn.set(style='whitegrid')
+ seaborn.set_context('paper')
+
+ plt.figure(1)
+ plt.subplot(111)
+ plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0))
+ plt.title('Title', fontsize=16, color='k')
+ plt.plot(list(range(N_iter)), lr_l, linewidth=1.5, label='learning rate scheme')
+ legend = plt.legend(loc='upper right', shadow=False)
+ ax = plt.gca()
+ labels = ax.get_xticks().tolist()
+ for k, v in enumerate(labels):
+ labels[k] = str(int(v / 1000)) + 'K'
+ ax.set_xticklabels(labels)
+ ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
+
+ ax.set_ylabel('Learning rate')
+ ax.set_xlabel('Iteration')
+ fig = plt.gcf()
+ plt.show()
diff --git a/codes/models/modules/ActNorms.py b/codes/models/modules/ActNorms.py
new file mode 100644
index 0000000..aa81f14
--- /dev/null
+++ b/codes/models/modules/ActNorms.py
@@ -0,0 +1,122 @@
+import torch
+from torch import nn as nn
+
+from models.modules import thops
+
+
+class _ActNorm(nn.Module):
+ """
+ Activation Normalization
+ Initialize the bias and scale with a given minibatch,
+ so that the output per-channel have zero mean and unit variance for that.
+
+ After initialization, `bias` and `logs` will be trained as parameters.
+ """
+
+ def __init__(self, num_features, scale=1.):
+ super().__init__()
+ # register mean and scale
+ size = [1, num_features, 1, 1]
+ self.register_parameter("bias", nn.Parameter(torch.zeros(*size)))
+ self.register_parameter("logs", nn.Parameter(torch.zeros(*size)))
+ self.num_features = num_features
+ self.scale = float(scale)
+ self.inited = False
+
+ def _check_input_dim(self, input):
+ return NotImplemented
+
+ def initialize_parameters(self, input):
+ self._check_input_dim(input)
+ if not self.training:
+ return
+ if (self.bias != 0).any():
+ self.inited = True
+ return
+ assert input.device == self.bias.device, (input.device, self.bias.device)
+ with torch.no_grad():
+ bias = thops.mean(input.clone(), dim=[0, 2, 3], keepdim=True) * -1.0
+ vars = thops.mean((input.clone() + bias) ** 2, dim=[0, 2, 3], keepdim=True)
+ logs = torch.log(self.scale / (torch.sqrt(vars) + 1e-6))
+ self.bias.data.copy_(bias.data)
+ self.logs.data.copy_(logs.data)
+ self.inited = True
+
+ def _center(self, input, reverse=False, offset=None):
+ bias = self.bias
+
+ if offset is not None:
+ bias = bias + offset
+
+ if not reverse:
+ return input + bias
+ else:
+ return input - bias
+
+ def _scale(self, input, logdet=None, reverse=False, offset=None):
+ logs = self.logs
+
+ if offset is not None:
+ logs = logs + offset
+
+ if not reverse:
+ input = input * torch.exp(logs) # should have shape batchsize, n_channels, 1, 1
+ # input = input * torch.exp(logs+logs_offset)
+ else:
+ input = input * torch.exp(-logs)
+ if logdet is not None:
+ """
+ logs is log_std of `mean of channels`
+ so we need to multiply pixels
+ """
+ dlogdet = thops.sum(logs) * thops.pixels(input)
+ if reverse:
+ dlogdet *= -1
+ logdet = logdet + dlogdet
+ return input, logdet
+
+ def forward(self, input, logdet=None, reverse=False, offset_mask=None, logs_offset=None, bias_offset=None):
+ if not self.inited:
+ self.initialize_parameters(input)
+
+ if offset_mask is not None:
+ logs_offset *= offset_mask
+ bias_offset *= offset_mask
+ # no need to permute dims as old version
+ if not reverse:
+ # center and scale
+ input = self._center(input, reverse, bias_offset)
+ input, logdet = self._scale(input, logdet, reverse, logs_offset)
+ else:
+ # scale and center
+ input, logdet = self._scale(input, logdet, reverse, logs_offset)
+ input = self._center(input, reverse, bias_offset)
+ return input, logdet
+
+
+class ActNorm2d(_ActNorm):
+ def __init__(self, num_features, scale=1.):
+ super().__init__(num_features, scale)
+
+ def _check_input_dim(self, input):
+ assert len(input.size()) == 4
+ assert input.size(1) == self.num_features, (
+ "[ActNorm]: input should be in shape as `BCHW`,"
+ " channels should be {} rather than {}".format(
+ self.num_features, input.size()))
+
+
+class MaskedActNorm2d(ActNorm2d):
+ def __init__(self, num_features, scale=1.):
+ super().__init__(num_features, scale)
+
+ def forward(self, input, mask, logdet=None, reverse=False):
+
+ assert mask.dtype == torch.bool
+ output, logdet_out = super().forward(input, logdet, reverse)
+
+ input[mask] = output[mask]
+ logdet[mask] = logdet_out[mask]
+
+ return input, logdet
+
diff --git a/codes/models/modules/AffineCouplings.py b/codes/models/modules/AffineCouplings.py
new file mode 100644
index 0000000..2c9b840
--- /dev/null
+++ b/codes/models/modules/AffineCouplings.py
@@ -0,0 +1,224 @@
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from models.modules import thops
+from models.modules.Basic import Conv2d, Conv2dZeros, DenseBlock, FCN, RDN
+from utils.util import opt_get, register_hook, trunc_normal_
+
+
+class AffineCoupling(nn.Module):
+ def __init__(self, in_channels, cond_channels=None, opt=None):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_channels = opt_get(opt, ['hidden_channels'], 64)
+ self.n_hidden_layers = 1
+ self.kernel_hidden = 1
+ self.cond_channels = cond_channels
+ f_in_channels = self.in_channels//2 if cond_channels is None else self.in_channels//2 + cond_channels
+ f_out_channels = (self.in_channels - self.in_channels//2) * 2
+ nn_module = opt_get(opt, ['nn_module'], 'FCN')
+ if nn_module == 'DenseBlock':
+ self.f = DenseBlock(in_channels=f_in_channels, out_channels=f_out_channels, gc=self.hidden_channels)
+ elif nn_module == 'FCN':
+ self.f = FCN(in_channels=f_in_channels, out_channels=f_out_channels, hidden_channels=self.hidden_channels,
+ kernel_hidden=self.kernel_hidden, n_hidden_layers=self.n_hidden_layers)
+
+
+ def forward(self, z, u=None, y=None, logdet=None, reverse=False):
+ if not reverse:
+ return self.normal_flow(z, u, y, logdet)
+ else:
+ return self.reverse_flow(z, u, y, logdet)
+
+ def normal_flow(self, z, u=None, y=None, logdet=None):
+ z1, z2 = thops.split_feature(z, "split")
+
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+ # adding 1e-4 is crucial for torch.slogdet(), as used in Glow (leads to black rect in experiments).
+ # see https://github.com/didriknielsen/survae_flows/issues/5 for discussion.
+ # or use `torch.exp(2. * torch.tanh(s / 2.)) as in SurVAE (more unstable in practice).
+
+ # version 1, srflow (use FCN)
+ # scale = torch.sigmoid(scale + 2.) + 1e-4
+ # z2 = (z2 + shift) * scale
+ # logdet += thops.sum(torch.log(scale), dim=[1, 2, 3])
+
+ # version2, survae
+ # logscale = 2. * torch.tanh(scale / 2.)
+ # z2 = (z2+shift) * torch.exp(logscale) # as in glow, it's shift+scale!
+ # logdet += thops.sum(logscale, dim=[1, 2, 3])
+
+ # version3, FrEIA, now have problem with FCN, but densenet is ok. (use FCN2/Denseblock)
+ # logscale = 0.5 * 0.636 * torch.atan(scale / 0.5) # clamp it to be between [-0.5,0.5]
+ logscale = 0.318 * torch.atan(2 * scale)
+ # logscale = 1.0 * 0.636 * torch.atan(scale / 1.0)
+ z2 = (z2 + shift) * torch.exp(logscale)
+ if logdet is not None:
+ logdet += thops.sum(logscale, dim=[1, 2, 3])
+
+ z = thops.cat_feature(z1, z2)
+
+ return z, logdet
+
+ def reverse_flow(self, z, u=None, y=None, logdet=None):
+ z1, z2 = thops.split_feature(z, "split")
+
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+
+ # version1, srflow
+ # scale = torch.sigmoid(scale + 2.) + 1e-4
+ # z2 = (z2 / scale) -shift
+
+ # version2, survae
+ # logscale = 2. * torch.tanh(scale / 2.)
+ # z2 = z2 * torch.exp(-logscale) - shift
+
+ # version3, FrEIA
+ # logscale = 0.5 * 0.636 * torch.atan(scale / 0.5)
+ logscale = 0.318 * torch.atan(2 * scale)
+ # logscale = 1 * 0.636 * torch.atan(scale / 1.0)
+ z2 = z2 * torch.exp(-logscale) - shift
+
+ z = thops.cat_feature(z1, z2)
+
+ return z, logdet
+
+
+'''3 channel conditional on the rest channels, or vice versa. only shift LR.
+ used in image rescaling to divide the low-frequencies and the high-frequencies apart from early flow layers.'''
+class AffineCoupling3shift(nn.Module):
+ def __init__(self, in_channels, cond_channels=None, LRvsothers=True, opt=None):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_channels = opt_get(opt, ['hidden_channels'], 64)
+ self.n_hidden_layers = 1
+ self.kernel_hidden = 1
+ self.cond_channels = cond_channels
+ self.LRvsothers = LRvsothers
+ if LRvsothers:
+ f_in_channels = 3 if cond_channels is None else 3 + cond_channels
+ f_out_channels = (self.in_channels - 3) * 2
+ else:
+ f_in_channels = self.in_channels - 3 if cond_channels is None else self.in_channels - 3 + cond_channels
+ f_out_channels = 3
+ nn_module = opt_get(opt, ['nn_module'], 'FCN')
+
+ if nn_module == 'DenseBlock':
+ self.f = DenseBlock(in_channels=f_in_channels, out_channels=f_out_channels, gc=self.hidden_channels)
+ elif nn_module == 'FCN':
+ self.f = FCN(in_channels=f_in_channels, out_channels=f_out_channels, hidden_channels=self.hidden_channels,
+ kernel_hidden=self.kernel_hidden, n_hidden_layers=self.n_hidden_layers)
+
+
+ def forward(self, z, u=None, y=None, logdet=None, reverse=False):
+ if not reverse:
+ return self.normal_flow(z, u, y, logdet)
+ else:
+ return self.reverse_flow(z, u, y, logdet)
+
+ def normal_flow(self, z, u=None, y=None, logdet=None):
+ if self.LRvsothers:
+ z1, z2 = z[:, :3, ...], z[:, 3:, ...]
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z2 = (z2 + shift) * torch.exp(logscale)
+ if logdet is not None:
+ logdet += thops.sum(logscale, dim=[1, 2, 3])
+ else:
+ z2, z1 = z[:, :3, ...], z[:, 3:, ...]
+ shift = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ z2 = z2 + shift
+
+ if self.LRvsothers:
+ z = thops.cat_feature(z1, z2)
+ else:
+ z = thops.cat_feature(z2, z1)
+
+ return z, logdet
+
+ def reverse_flow(self, z, u=None, y=None, logdet=None):
+ if self.LRvsothers:
+ z1, z2 = z[:, :3, ...], z[:, 3:, ...]
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z2 = z2 * torch.exp(-logscale) - shift
+ else:
+ z2, z1 = z[:, :3, ...], z[:, 3:, ...]
+ shift = self.f(z1)
+ z2 = z2 - shift
+
+ if self.LRvsothers:
+ z = thops.cat_feature(z1, z2)
+ else:
+ z = thops.cat_feature(z2, z1)
+
+ return z, logdet
+
+
+''' srflow's affine injector + original affine coupling, not used in this project'''
+class AffineCouplingInjector(nn.Module):
+ def __init__(self, in_channels, cond_channels=None, opt=None):
+ super().__init__()
+ self.in_channels = in_channels
+ self.hidden_channels = opt_get(opt, ['hidden_channels'], 64)
+ self.n_hidden_layers = 1
+ self.kernel_hidden = 1
+ self.cond_channels = cond_channels
+ f_in_channels = self.in_channels//2 if cond_channels is None else self.in_channels//2 + cond_channels
+ f_out_channels = (self.in_channels - self.in_channels//2) * 2
+ nn_module = opt_get(opt, ['nn_module'], 'FCN')
+ if nn_module == 'DenseBlock':
+ self.f = DenseBlock(in_channels=f_in_channels, out_channels=f_out_channels, gc=self.hidden_channels)
+ self.f_injector = DenseBlock(in_channels=cond_channels, out_channels=self.in_channels*2, gc=self.hidden_channels)
+ elif nn_module == 'FCN':
+ self.f = FCN(in_channels=f_in_channels, out_channels=f_out_channels, hidden_channels=self.hidden_channels,
+ kernel_hidden=self.kernel_hidden, n_hidden_layers=self.n_hidden_layers)
+ self.f_injector = FCN(in_channels=cond_channels, out_channels=self.in_channels*2, hidden_channels=self.hidden_channels,
+ kernel_hidden=self.kernel_hidden, n_hidden_layers=self.n_hidden_layers)
+
+ def forward(self, z, u=None, y=None, logdet=None, reverse=False):
+ if not reverse:
+ return self.normal_flow(z, u, y, logdet)
+ else:
+ return self.reverse_flow(z, u, y, logdet)
+
+ def normal_flow(self, z, u=None, y=None, logdet=None):
+ # overall-conditional
+ h = self.f_injector(u)
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale) # clamp it to be between [-5,5]
+ z = (z + shift) * torch.exp(logscale)
+ logdet += thops.sum(logscale, dim=[1, 2, 3])
+
+ # self-conditional
+ z1, z2 = thops.split_feature(z, "split")
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale) # clamp it to be between [-5,5]
+ z2 = (z2 + shift) * torch.exp(logscale)
+ logdet += thops.sum(logscale, dim=[1, 2, 3])
+ z = thops.cat_feature(z1, z2)
+
+ return z, logdet
+
+ def reverse_flow(self, z, u=None, y=None, logdet=None):
+ # self-conditional
+ z1, z2 = thops.split_feature(z, "split")
+ h = self.f(z1) if self.cond_channels is None else self.f(thops.cat_feature(z1, u))
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z2 = z2 * torch.exp(-logscale) - shift
+ z = thops.cat_feature(z1, z2)
+
+ # overall-conditional
+ h = self.f_injector(u)
+ shift, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z = z * torch.exp(-logscale) - shift
+
+ return z, logdet
diff --git a/codes/models/modules/Basic.py b/codes/models/modules/Basic.py
new file mode 100644
index 0000000..aa5897b
--- /dev/null
+++ b/codes/models/modules/Basic.py
@@ -0,0 +1,500 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import functools
+
+from models.modules.ActNorms import ActNorm2d
+from . import thops
+
+from utils.util import opt_get
+import models.modules.module_util as mutil
+
+
+class Conv2d(nn.Conv2d):
+ pad_dict = {
+ "same": lambda kernel, stride: [((k - 1) * s + 1) // 2 for k, s in zip(kernel, stride)],
+ "valid": lambda kernel, stride: [0 for _ in kernel]
+ }
+
+ @staticmethod
+ def get_padding(padding, kernel_size, stride):
+ # make paddding
+ if isinstance(padding, str):
+ if isinstance(kernel_size, int):
+ kernel_size = [kernel_size, kernel_size]
+ if isinstance(stride, int):
+ stride = [stride, stride]
+ padding = padding.lower()
+ try:
+ padding = Conv2d.pad_dict[padding](kernel_size, stride)
+ except KeyError:
+ raise ValueError("{} is not supported".format(padding))
+ return padding
+
+ def __init__(self, in_channels, out_channels,
+ kernel_size=[3, 3], stride=[1, 1], padding="same", groups=1,
+ do_actnorm=True, weight_std=0.05):
+ padding = Conv2d.get_padding(padding, kernel_size, stride)
+ super().__init__(in_channels, out_channels, kernel_size, stride,
+ padding, groups=groups, bias=(not do_actnorm))
+ # init weight with std
+ self.weight.data.normal_(mean=0.0, std=weight_std)
+ if not do_actnorm:
+ self.bias.data.zero_()
+ else:
+ self.actnorm = ActNorm2d(out_channels)
+ self.do_actnorm = do_actnorm
+
+ def forward(self, input):
+ x = super().forward(input)
+ if self.do_actnorm:
+ x, _ = self.actnorm(x)
+ return x
+
+# Zero initialization. We initialize the last convolution of each NN() with zeros, such that each affine
+# coupling layer initially performs an identity function; we found that this helps training very deep networks.
+class Conv2dZeros(nn.Conv2d):
+ def __init__(self, in_channels, out_channels,
+ kernel_size=[3, 3], stride=[1, 1],
+ padding="same", logscale_factor=3):
+ padding = Conv2d.get_padding(padding, kernel_size, stride)
+ super().__init__(in_channels, out_channels, kernel_size, stride, padding)
+ # logscale_factor
+ self.logscale_factor = logscale_factor
+ self.register_parameter("logs", nn.Parameter(torch.zeros(out_channels, 1, 1)))
+ # init
+ self.weight.data.zero_()
+ self.bias.data.zero_()
+
+ def forward(self, input):
+ output = super().forward(input)
+ return output * torch.exp(self.logs * self.logscale_factor)
+
+
+class GaussianDiag:
+ Log2PI = float(np.log(2 * np.pi))
+
+ @staticmethod
+ def likelihood(mean, logs, x): # logs: log(sigma)
+ """
+ lnL = -1/2 * { ln|Var| + ((X - Mu)^T)(Var^-1)(X - Mu) + kln(2*PI) }
+ k = 1 (Independent)
+ Var = logs ** 2
+ """
+ if mean is None and logs is None:
+ return -0.5 * (x ** 2 + GaussianDiag.Log2PI)
+ else:
+ return -0.5 * (logs * 2. + ((x - mean) ** 2) / torch.exp(logs * 2.) + GaussianDiag.Log2PI)
+
+ @staticmethod
+ def logp(mean, logs, x):
+ likelihood = GaussianDiag.likelihood(mean, logs, x)
+ return thops.sum(likelihood, dim=[1, 2, 3])
+
+ @staticmethod
+ def sample(mean, logs, eps_std=None):
+ # eps_std = eps_std or 1 # may cause problem when eps_std is 0
+ eps = torch.normal(mean=torch.zeros_like(mean),
+ std=torch.ones_like(logs) * eps_std)
+ return mean + torch.exp(logs) * eps
+
+ @staticmethod
+ def sample_eps(shape, eps_std, seed=None):
+ if seed is not None:
+ torch.manual_seed(seed)
+ eps = torch.normal(mean=torch.zeros(shape),
+ std=torch.ones(shape) * eps_std)
+ return eps
+
+
+class LaplaceDiag:
+ Log2= float(np.log(2))
+
+ @staticmethod
+ def likelihood(mean, logs, x): # logs: log(sigma)
+ if mean is None and logs is None:
+ return - (torch.abs(x) + LaplaceDiag.Log2)
+ else:
+ return - (logs + (torch.abs(x - mean)) / torch.exp(logs) + LaplaceDiag.Log2)
+
+ @staticmethod
+ def logp(mean, logs, x):
+ likelihood = LaplaceDiag.likelihood(mean, logs, x)
+ return thops.sum(likelihood, dim=[1, 2, 3])
+
+
+def squeeze2d(input, factor=2):
+ assert factor >= 1 and isinstance(factor, int)
+ if factor == 1:
+ return input
+ size = input.size()
+ B = size[0]
+ C = size[1]
+ H = size[2]
+ W = size[3]
+ assert H % factor == 0 and W % factor == 0, "{}".format((H, W, factor))
+ x = input.view(B, C, H // factor, factor, W // factor, factor)
+ x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
+ x = x.view(B, C * factor * factor, H // factor, W // factor)
+ return x
+
+
+def unsqueeze2d(input, factor=2):
+ assert factor >= 1 and isinstance(factor, int)
+ factor2 = factor ** 2
+ if factor == 1:
+ return input
+ size = input.size()
+ B = size[0]
+ C = size[1]
+ H = size[2]
+ W = size[3]
+ assert C % (factor2) == 0, "{}".format(C)
+ x = input.view(B, C // factor2, factor, factor, H, W)
+ x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
+ x = x.view(B, C // (factor2), H * factor, W * factor)
+ return x
+
+
+class SqueezeLayer(nn.Module):
+ def __init__(self, factor):
+ super().__init__()
+ self.factor = factor
+
+ def forward(self, input, logdet=None, reverse=False):
+ if not reverse:
+ output = squeeze2d(input, self.factor)
+ return output, logdet
+ else:
+ output = unsqueeze2d(input, self.factor)
+ return output, logdet
+
+class UnSqueezeLayer(nn.Module):
+ def __init__(self, factor):
+ super().__init__()
+ self.factor = factor
+
+ def forward(self, input, logdet=None, reverse=False):
+ if not reverse:
+ output = unsqueeze2d(input, self.factor)
+ return output, logdet
+ else:
+ output = squeeze2d(input, self.factor)
+ return output, logdet
+
+class Quant(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, input):
+ input = torch.clamp(input, 0, 1)
+ output = (input * 255.).round() / 255.
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ return grad_output
+
+class Quantization(nn.Module):
+ def __init__(self):
+ super(Quantization, self).__init__()
+
+ def forward(self, input):
+ return Quant.apply(input)
+
+class Sigmoid(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, input, logdet=None, reverse=False):
+ if not reverse:
+ output = torch.sigmoid(input)
+ logdet += -thops.sum(F.softplus(input)+F.softplus(-input), dim=[1, 2, 3])
+ return output, logdet
+ else:
+ output = -torch.log(torch.reciprocal(input) - 1.)
+ logdet += -thops.sum(torch.log(input) + torch.log(1.-input), dim=[1, 2, 3])
+ return output, logdet
+
+# used in SRFlow
+class Split2d_conditional(nn.Module):
+ def __init__(self, num_channels, logs_eps=0, cond_channels=0, position=None, consume_ratio=0.5, opt=None):
+ super().__init__()
+
+ self.num_channels_consume = int(round(num_channels * consume_ratio))
+ self.num_channels_pass = num_channels - self.num_channels_consume
+
+ self.conv = Conv2dZeros(in_channels=self.num_channels_pass + cond_channels,
+ out_channels=self.num_channels_consume * 2)
+ self.logs_eps = logs_eps
+ self.position = position
+ self.opt = opt
+
+ def split2d_prior(self, z, ft):
+ if ft is not None:
+ z = torch.cat([z, ft], dim=1)
+ h = self.conv(z)
+ return thops.split_feature(h, "cross")
+
+ def exp_eps(self, logs):
+ return torch.exp(logs) + self.logs_eps
+
+ def forward(self, input, logdet=0., reverse=False, eps_std=None, eps=None, ft=None, y_onehot=None):
+ if not reverse:
+ # self.input = input
+ z1, z2 = self.split_ratio(input)
+ mean, logs = self.split2d_prior(z1, ft)
+
+ eps = (z2 - mean) / self.exp_eps(logs)
+
+ logdet = logdet + self.get_logdet(logs, mean, z2)
+
+ # print(logs.shape, mean.shape, z2.shape)
+ # self.eps = eps
+ # print('split, enc eps:', eps)
+ return z1, logdet, eps
+ else:
+ z1 = input
+ mean, logs = self.split2d_prior(z1, ft)
+
+ if eps is None:
+ #print("WARNING: eps is None, generating eps untested functionality!")
+ eps = GaussianDiag.sample_eps(mean.shape, eps_std)
+
+ eps = eps.to(mean.device)
+ z2 = mean + self.exp_eps(logs) * eps
+
+ z = thops.cat_feature(z1, z2)
+ logdet = logdet - self.get_logdet(logs, mean, z2)
+
+ return z, logdet
+ # return z, logdet, eps
+
+ def get_logdet(self, logs, mean, z2):
+ logdet_diff = GaussianDiag.logp(mean, logs, z2)
+ # print("Split2D: logdet diff", logdet_diff.item())
+ return logdet_diff
+
+ def split_ratio(self, input):
+ z1, z2 = input[:, :self.num_channels_pass, ...], input[:, self.num_channels_pass:, ...]
+ return z1, z2
+
+
+''' Not used anymore '''
+class Split2d(nn.Module):
+ def __init__(self, num_channels):
+ super().__init__()
+ self.conv = Conv2dZeros(num_channels // 2, num_channels)
+
+ def split2d_prior(self, z):
+ h = self.conv(z)
+ return thops.split_feature(h, "cross")
+
+ def forward(self, input, logdet=0., reverse=False, eps_std=None):
+ if not reverse:
+ z1, z2 = thops.split_feature(input, "split")
+ mean, logs = self.split2d_prior(z1)
+ logdet = GaussianDiag.logp(mean, logs, z2) + logdet
+ return z1, logdet
+ else:
+ z1 = input
+ mean, logs = self.split2d_prior(z1)
+ z2 = GaussianDiag.sample(mean, logs, eps_std)
+ z = thops.cat_feature(z1, z2)
+ return z, logdet
+
+class Split2d_LR(nn.Module):
+ def __init__(self, num_channels, num_channels_split):
+ super().__init__()
+ self.num_channels_split = num_channels_split
+ self.conv = Conv2dZeros(num_channels_split, (num_channels-num_channels_split)*2)
+
+ def split2d_prior(self, z):
+ h = self.conv(z)
+ return thops.split_feature(h, "cross")
+
+ def forward(self, input, eps_std=None, logdet=0., reverse=False):
+ if not reverse:
+ z1, z2 = input[:, :self.num_channels_split, ...], input[:, self.num_channels_split:, ...]
+ mean, logs = self.split2d_prior(z1)
+ logdet += GaussianDiag.logp(mean, logs, z2)
+ return z1, logdet
+ else:
+ z1 = input
+ mean, logs = self.split2d_prior(z1)
+ z2 = GaussianDiag.sample(mean, logs, eps_std)
+ z = torch.cat((z1, z2), dim=1)
+ return z, logdet
+
+# DenseBlock for affine coupling (flow)
+class DenseBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, gc=32, bias=True, init='xavier', for_flow=True):
+ super(DenseBlock, self).__init__()
+ self.conv1 = nn.Conv2d(in_channels, gc, 3, 1, 1, bias=bias)
+ self.conv2 = nn.Conv2d(in_channels + gc, gc, 3, 1, 1, bias=bias)
+ self.conv3 = nn.Conv2d(in_channels + 2 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv4 = nn.Conv2d(in_channels + 3 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv5 = nn.Conv2d(in_channels + 4 * gc, out_channels, 3, 1, 1, bias=bias)
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ # init as 'xavier', following the practice in https://github.com/VLL-HD/FrEIA/blob/c5fe1af0de8ce9122b5b61924ad75a19b9dc2473/README.rst#useful-tips--engineering-heuristics
+ if init == 'xavier':
+ mutil.initialize_weights_xavier([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+ else:
+ mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+ # initialiize input to all zeros to have zero mean and unit variance
+ if for_flow:
+ mutil.initialize_weights(self.conv5, 0)
+
+ 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))
+
+ return x5
+
+
+# ResidualDenseBlock for multi-layer feature extraction
+class ResidualDenseBlock(nn.Module):
+ def __init__(self, nf=64, gc=32, bias=True, init='xavier'):
+ super(ResidualDenseBlock, self).__init__()
+ # gc: growth channel, i.e. intermediate channels
+ self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+ self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+ self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ # init as 'xavier', following the practice in https://github.com/VLL-HD/FrEIA/blob/c5fe1af0de8ce9122b5b61924ad75a19b9dc2473/README.rst#useful-tips--engineering-heuristics
+ if init == 'xavier':
+ mutil.initialize_weights_xavier([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+ else:
+ mutil.initialize_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))
+ return x5 * 0.2 + x # residual scaling are helpful
+
+class RRDB(nn.Module):
+ '''Residual in Residual Dense Block'''
+
+ def __init__(self, nf, gc=32):
+ super(RRDB, self).__init__()
+ self.RDB1 = ResidualDenseBlock(nf, gc)
+ self.RDB2 = ResidualDenseBlock(nf, gc)
+ self.RDB3 = ResidualDenseBlock(nf, gc)
+
+ def forward(self, x):
+ out = self.RDB1(x)
+ out = self.RDB2(out)
+ out = self.RDB3(out)
+ return out * 0.2 + x # residual scaling are helpful
+
+class RDN(nn.Module):
+ '''composed of rrdb blocks'''
+
+ def __init__(self, in_channels, out_channels, nb=3, nf=64, gc=32, init='xavier', for_flow=True):
+ super(RDN, self).__init__()
+
+ RRDB_f = functools.partial(RRDB, nf=nf, gc=gc)
+ self.conv_first = nn.Conv2d(in_channels, nf, 3, 1, 1, bias=True)
+ self.RRDB_trunk = mutil.make_layer(RRDB_f, nb)
+ self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ self.conv_last = nn.Conv2d(nf, out_channels, 3, 1, 1, bias=True)
+
+ if init == 'xavier':
+ mutil.initialize_weights_xavier([self.conv_first, self.trunk_conv, self.conv_last], 0.1)
+ else:
+ mutil.initialize_weights([self.conv_first, self.trunk_conv, self.conv_last], 0.1)
+
+ if for_flow:
+ mutil.initialize_weights(self.conv_last, 0)
+
+ def forward(self, x):
+ x = self.conv_first(x)
+ x = self.trunk_conv(self.RRDB_trunk(x)) + x
+ return self.conv_last(x)
+
+
+class FCN(nn.Module):
+ def __init__(self, in_channels, out_channels, hidden_channels, kernel_hidden=1, n_hidden_layers=1, init='xavier', for_flow=True):
+ super(FCN, self).__init__()
+ self.conv1 = Conv2d(in_channels, hidden_channels, kernel_size=[3, 3], stride=[1, 1])
+ self.conv2 = Conv2d(hidden_channels, hidden_channels, kernel_size=[kernel_hidden, kernel_hidden])
+ self.conv3 = Conv2dZeros(hidden_channels, out_channels, kernel_size=[3, 3], stride=[1, 1])
+ self.relu = nn.ReLU(inplace=False)
+
+ if init == 'xavier':
+ mutil.initialize_weights_xavier([self.conv1, self.conv2, self.conv3], 0.1)
+ else:
+ mutil.initialize_weights([self.conv1, self.conv2, self.conv3], 0.1)
+
+ if for_flow:
+ mutil.initialize_weights(self.conv3, 0)
+
+ def forward(self, x):
+ x = self.relu(self.conv1(x))
+ x = self.relu(self.conv2(x))
+ x = self.conv3(x)
+
+ return x
+
+
+class HaarDownsampling(nn.Module):
+ def __init__(self, channel_in):
+ super(HaarDownsampling, self).__init__()
+ self.channel_in = channel_in
+
+ self.haar_weights = torch.ones(4, 1, 2, 2)
+
+ self.haar_weights[1, 0, 0, 1] = -1
+ self.haar_weights[1, 0, 1, 1] = -1
+
+ self.haar_weights[2, 0, 1, 0] = -1
+ self.haar_weights[2, 0, 1, 1] = -1
+
+ self.haar_weights[3, 0, 1, 0] = -1
+ self.haar_weights[3, 0, 0, 1] = -1
+
+ self.haar_weights = torch.cat([self.haar_weights] * self.channel_in, 0)
+ self.haar_weights = nn.Parameter(self.haar_weights)
+ self.haar_weights.requires_grad = False
+
+ def forward(self, x, logdet=None, reverse=False):
+ if not reverse:
+ self.elements = x.shape[1] * x.shape[2] * x.shape[3]
+ self.last_jac = self.elements / 4 * np.log(1/16.)
+
+ out = F.conv2d(x, self.haar_weights, bias=None, stride=2, groups=self.channel_in) / 4.0
+ out = out.reshape([x.shape[0], self.channel_in, 4, x.shape[2] // 2, x.shape[3] // 2])
+ out = torch.transpose(out, 1, 2)
+ out = out.reshape([x.shape[0], self.channel_in * 4, x.shape[2] // 2, x.shape[3] // 2])
+ return out, logdet
+ else:
+ self.elements = x.shape[1] * x.shape[2] * x.shape[3]
+ self.last_jac = self.elements / 4 * np.log(16.)
+
+ out = x.reshape([x.shape[0], 4, self.channel_in, x.shape[2], x.shape[3]])
+ out = torch.transpose(out, 1, 2)
+ out = out.reshape([x.shape[0], self.channel_in * 4, x.shape[2], x.shape[3]])
+ return F.conv_transpose2d(out, self.haar_weights, bias=None, stride=2, groups = self.channel_in), logdet
+
+class Split(nn.Module):
+ def __init__(self, num_channels_split, level):
+ super().__init__()
+ self.num_channels_split = num_channels_split
+ self.level = level
+
+ def forward(self, z, z2=None, reverse=False):
+ if not reverse:
+ return z[:, :self.num_channels_split, ...], z[:, self.num_channels_split:, ...]
+ else:
+ return torch.cat((z, z2), dim=1)
+
diff --git a/codes/models/modules/ConditionalFlow.py b/codes/models/modules/ConditionalFlow.py
new file mode 100644
index 0000000..0581484
--- /dev/null
+++ b/codes/models/modules/ConditionalFlow.py
@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from . import thops
+from utils.util import opt_get
+from models.modules.Basic import Conv2d, Conv2dZeros, GaussianDiag, DenseBlock, RRDB, FCN
+from models.modules.FlowStep import FlowStep
+
+import functools
+import models.modules.module_util as mutil
+
+
+class ConditionalFlow(nn.Module):
+ def __init__(self, num_channels, num_channels_split, n_flow_step=0, opt=None, num_levels_condition=0, SR=True):
+ super().__init__()
+ self.SR = SR
+
+ # number of levels of RRDB features. One level of conditional feature is enough for image rescaling
+ num_features_condition = 2 if self.SR else 1
+
+ # feature extraction
+ RRDB_nb = opt_get(opt, ['RRDB_nb'], [5, 5])
+ RRDB_nf = opt_get(opt, ['RRDB_nf'], 64)
+ RRDB_gc = opt_get(opt, ['RRDB_gc'], 32)
+ RRDB_f = functools.partial(RRDB, nf=RRDB_nf, gc=RRDB_gc)
+ self.conv_first = nn.Conv2d(num_channels_split + RRDB_nf*num_features_condition*num_levels_condition, RRDB_nf, 3, 1, 1, bias=True)
+ self.RRDB_trunk0 = mutil.make_layer(RRDB_f, RRDB_nb[0])
+ self.RRDB_trunk1 = mutil.make_layer(RRDB_f, RRDB_nb[1])
+ self.trunk_conv1 = nn.Conv2d(RRDB_nf, RRDB_nf, 3, 1, 1, bias=True)
+
+ # conditional flow
+ self.additional_flow_steps = nn.ModuleList()
+ for k in range(n_flow_step):
+ self.additional_flow_steps.append(FlowStep(in_channels=num_channels-num_channels_split,
+ cond_channels=RRDB_nf*num_features_condition,
+ flow_permutation=opt['flow_permutation'],
+ flow_coupling=opt['flow_coupling'], opt=opt))
+
+ self.f = Conv2dZeros(RRDB_nf*num_features_condition, (num_channels-num_channels_split)*2)
+
+
+ def forward(self, z, u, eps_std=None, logdet=0., reverse=False, training=True):
+ # for image SR
+ if self.SR:
+ if not reverse:
+ conditional_feature = self.get_conditional_feature_SR(u)
+
+ for layer in self.additional_flow_steps:
+ z, logdet = layer(z, u=conditional_feature, logdet=logdet, reverse=False)
+
+ h = self.f(conditional_feature)
+ mean, logs = thops.split_feature(h, "cross")
+ logdet += GaussianDiag.logp(mean, logs, z)
+
+ return logdet, conditional_feature
+
+ else:
+ conditional_feature = self.get_conditional_feature_SR(u)
+
+ h = self.f(conditional_feature)
+ mean, logs = thops.split_feature(h, "cross")
+ z = GaussianDiag.sample(mean, logs, eps_std)
+
+ for layer in reversed(self.additional_flow_steps):
+ z, _ = layer(z, u=conditional_feature, reverse=True)
+
+ return z, logdet, conditional_feature
+ else:
+ # for image rescaling
+ if not reverse:
+ conditional_feature = self.get_conditional_feature_Rescaling(u)
+
+ for layer in self.additional_flow_steps:
+ z, logdet = layer(z, u=conditional_feature, logdet=logdet, reverse=False)
+
+ h = self.f(conditional_feature)
+ mean, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z = (z - mean) * torch.exp(-logscale)
+
+ return z, conditional_feature
+
+ else:
+ conditional_feature = self.get_conditional_feature_Rescaling(u)
+
+ h = self.f(conditional_feature)
+ mean, scale = thops.split_feature(h, "cross")
+ logscale = 0.318 * torch.atan(2 * scale)
+ z = GaussianDiag.sample(mean, logscale, eps_std)
+
+ for layer in reversed(self.additional_flow_steps):
+ z, _ = layer(z, u=conditional_feature, reverse=True)
+
+ return z, conditional_feature
+
+
+ def get_conditional_feature_SR(self, u):
+ u_feature_first = self.conv_first(u)
+ u_feature1 = self.RRDB_trunk0(u_feature_first)
+ u_feature2 = self.trunk_conv1(self.RRDB_trunk1(u_feature1)) + u_feature_first
+
+ return torch.cat([u_feature1, u_feature2], 1)
+
+ def get_conditional_feature_Rescaling(self, u):
+ u_feature_first = self.conv_first(u)
+ u_feature = self.trunk_conv1(self.RRDB_trunk1(self.RRDB_trunk0(u_feature_first))) + u_feature_first
+
+ return u_feature
+
+
diff --git a/codes/models/modules/FlowNet_Rescaling_x4.py b/codes/models/modules/FlowNet_Rescaling_x4.py
new file mode 100644
index 0000000..de8f4b3
--- /dev/null
+++ b/codes/models/modules/FlowNet_Rescaling_x4.py
@@ -0,0 +1,129 @@
+import numpy as np
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from utils.util import opt_get
+from models.modules import Basic
+from models.modules.FlowStep import FlowStep
+from models.modules.ConditionalFlow import ConditionalFlow
+
+class FlowNet(nn.Module):
+ def __init__(self, image_shape, opt=None):
+ assert image_shape[2] == 1 or image_shape[2] == 3
+ super().__init__()
+ H, W, self.C = image_shape
+ self.opt = opt
+ self.L = opt_get(opt, ['network_G', 'flowDownsampler', 'L'])
+ self.K = opt_get(opt, ['network_G', 'flowDownsampler', 'K'])
+ if isinstance(self.K, int): self.K = [self.K] * (self.L + 1)
+
+ squeeze = opt_get(self.opt, ['network_G', 'flowDownsampler', 'squeeze'], 'checkerboard')
+ n_additionalFlowNoAffine = opt_get(self.opt, ['network_G', 'flowDownsampler', 'additionalFlowNoAffine'], 0)
+ flow_permutation = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_permutation'], 'invconv')
+ flow_coupling = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_coupling'], 'Affine')
+ cond_channels = opt_get(self.opt, ['network_G', 'flowDownsampler', 'cond_channels'], None)
+ enable_splitOff = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'enable'], False)
+ after_splitOff_flowStep = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'after_flowstep'], 0)
+ if isinstance(after_splitOff_flowStep, int): after_splitOff_flowStep = [after_splitOff_flowStep] * (self.L + 1)
+
+ # construct flow
+ self.layers = nn.ModuleList()
+ self.output_shapes = []
+
+ for level in range(self.L):
+ # 1. Squeeze
+ if squeeze == 'checkerboard':
+ self.layers.append(Basic.SqueezeLayer(factor=2)) # may need a better way for squeezing
+ elif squeeze == 'haar':
+ self.layers.append(Basic.HaarDownsampling(channel_in=self.C))
+
+ self.C, H, W = self.C * 4, H // 2, W // 2
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 2. main FlowSteps (uncodnitional flow)
+ for k in range(self.K[level]-after_splitOff_flowStep[level]):
+ self.layers.append(FlowStep(in_channels=self.C, cond_channels=cond_channels,
+ flow_permutation=flow_permutation,
+ flow_coupling=flow_coupling,
+ LRvsothers=True if k%2==0 else False,
+ opt=opt['network_G']['flowDownsampler']))
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 3. additional FlowSteps (split + conditional flow)
+ if enable_splitOff:
+ if level == 0:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level0_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=1, SR=False)
+ elif level == 1:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level1_condFlow = (ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=0, SR=False))
+
+ self.C = self.C // 2 if level < self.L-1 else 3
+ self.output_shapes.append([-1, self.C, H, W])
+
+
+ self.H = H
+ self.W = W
+ self.scaleH = image_shape[0] / H
+ self.scaleW = image_shape[1] / W
+ print('shapes:', self.output_shapes)
+
+ def forward(self, hr=None, z=None, u=None, eps_std=None, logdet=None, reverse=False, training=True):
+ if not reverse:
+ return self.normal_flow(hr, u=u, logdet=logdet, training=training)
+ else:
+ return self.reverse_flow(z, u=u, eps_std=eps_std, training=training)
+
+
+ '''
+ hr->y1+z1->y2+z2
+ '''
+ def normal_flow(self, z, u=None, logdet=None, training=True):
+ for layer, shape in zip(self.layers, self.output_shapes):
+ if isinstance(layer, FlowStep):
+ z, _ = layer(z, u, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.SqueezeLayer) or isinstance(layer, Basic.HaarDownsampling):
+ z, _ = layer(z, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 0:
+ z, a1 = layer(z, reverse=False)
+ y1 = z.clone()
+ elif layer.level == 1:
+ z, a2 = layer(z, reverse=False)
+ fake_z2, conditional_feature2 = self.level1_condFlow(a2, z, logdet=logdet, reverse=False, training=training)
+
+ conditional_feature1 = torch.cat([y1, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest')],1)
+ fake_z1, _ = self.level0_condFlow(a1, conditional_feature1, logdet=logdet, reverse=False, training=training)
+
+ return z, fake_z1, fake_z2
+
+ '''
+ y2+z2->y1+z1->hr
+ '''
+ def reverse_flow(self, z, u=None, eps_std=None, training=True):
+ for layer, shape in zip(reversed(self.layers), reversed(self.output_shapes)):
+ if isinstance(layer, FlowStep):
+ z, _ = layer(z, u, reverse=True)
+ elif isinstance(layer, Basic.SqueezeLayer) or isinstance(layer, Basic.HaarDownsampling):
+ z, _ = layer(z, reverse=True)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 1:
+ a2, conditional_feature2 = self.level1_condFlow(None, z, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a2, reverse=True)
+ elif layer.level == 0:
+ conditional_feature1 = torch.cat([z, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest')],1)
+ a1, _ = self.level0_condFlow(None, conditional_feature1, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a1, reverse=True)
+
+
+ return z
+
diff --git a/codes/models/modules/FlowNet_SR_x4.py b/codes/models/modules/FlowNet_SR_x4.py
new file mode 100644
index 0000000..17e6924
--- /dev/null
+++ b/codes/models/modules/FlowNet_SR_x4.py
@@ -0,0 +1,124 @@
+import numpy as np
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from utils.util import opt_get
+from models.modules import Basic
+from models.modules.FlowStep import FlowStep
+from models.modules.ConditionalFlow import ConditionalFlow
+
+class FlowNet(nn.Module):
+ def __init__(self, image_shape, opt=None):
+ assert image_shape[2] == 1 or image_shape[2] == 3
+ super().__init__()
+ H, W, self.C = image_shape
+ self.opt = opt
+ self.L = opt_get(opt, ['network_G', 'flowDownsampler', 'L'])
+ self.K = opt_get(opt, ['network_G', 'flowDownsampler', 'K'])
+ if isinstance(self.K, int): self.K = [self.K] * (self.L + 1)
+
+ n_additionalFlowNoAffine = opt_get(self.opt, ['network_G', 'flowDownsampler', 'additionalFlowNoAffine'], 0)
+ flow_permutation = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_permutation'], 'invconv')
+ flow_coupling = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_coupling'], 'Affine')
+ cond_channels = opt_get(self.opt, ['network_G', 'flowDownsampler', 'cond_channels'], None)
+ enable_splitOff = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'enable'], False)
+ after_splitOff_flowStep = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'after_flowstep'], 0)
+ if isinstance(after_splitOff_flowStep, int): after_splitOff_flowStep = [after_splitOff_flowStep] * (self.L + 1)
+
+ # construct flow
+ self.layers = nn.ModuleList()
+ self.output_shapes = []
+
+ for level in range(self.L):
+ # 1. Squeeze
+ self.layers.append(Basic.SqueezeLayer(factor=2)) # may need a better way for squeezing
+ self.C, H, W = self.C * 4, H // 2, W // 2
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 2. main FlowSteps (unconditional flow)
+ for k in range(self.K[level]-after_splitOff_flowStep[level]):
+ self.layers.append(FlowStep(in_channels=self.C, cond_channels=cond_channels,
+ flow_permutation=flow_permutation,
+ flow_coupling=flow_coupling,
+ opt=opt['network_G']['flowDownsampler']))
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 3. additional FlowSteps (split + conditional flow)
+ if enable_splitOff:
+ if level == 0:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level0_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=1, SR=True)
+ elif level == 1:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level1_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=0, SR=True)
+ self.C = self.C // 2 if level < self.L-1 else 3
+ self.output_shapes.append([-1, self.C, H, W])
+
+
+ self.H = H
+ self.W = W
+ self.scaleH = image_shape[0] / H
+ self.scaleW = image_shape[1] / W
+ print('shapes:', self.output_shapes)
+
+ # nodetach version; 0.05 better than detach version, 0.30 better when using only nll loss
+ def forward(self, hr=None, z=None, u=None, eps_std=None, logdet=None, reverse=False, training=True):
+ if not reverse:
+ return self.normal_flow(hr, u=u, logdet=logdet, training=training)
+ else:
+ return self.reverse_flow(z, u=u, eps_std=eps_std, training=training)
+
+
+ '''
+ hr->y1+z1->y2+z2
+ '''
+ def normal_flow(self, z, u=None, logdet=None, training=True):
+ for layer, shape in zip(self.layers, self.output_shapes):
+ if isinstance(layer, FlowStep):
+ z, logdet = layer(z, u, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.SqueezeLayer):
+ z, logdet = layer(z, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 0:
+ z, a1 = layer(z, reverse=False)
+ y1 = z.clone()
+ elif layer.level == 1:
+ z, a2 = layer(z, reverse=False)
+ logdet, conditional_feature2 = self.level1_condFlow(a2, z, logdet=logdet, reverse=False, training=training)
+
+ conditional_feature1 = torch.cat([y1, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest')],1)
+ logdet, _ = self.level0_condFlow(a1, conditional_feature1, logdet=logdet, reverse=False, training=training)
+
+ return z, logdet
+
+ '''
+ y2+z2->y1+z1->hr
+ '''
+ def reverse_flow(self, z, u=None, eps_std=None, training=True):
+ for layer, shape in zip(reversed(self.layers), reversed(self.output_shapes)):
+ if isinstance(layer, FlowStep):
+ z, _ = layer(z, u, reverse=True)
+ elif isinstance(layer, Basic.SqueezeLayer):
+ z, _ = layer(z, reverse=True)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 1:
+ a2, _, conditional_feature2 = self.level1_condFlow(None, z, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a2, reverse=True)
+ elif layer.level == 0:
+ conditional_feature1 = torch.cat([z, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest')],1)
+ a1, _, _ = self.level0_condFlow(None, conditional_feature1, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a1, reverse=True)
+
+
+
+ return z
+
diff --git a/codes/models/modules/FlowNet_SR_x8.py b/codes/models/modules/FlowNet_SR_x8.py
new file mode 100644
index 0000000..b07800d
--- /dev/null
+++ b/codes/models/modules/FlowNet_SR_x8.py
@@ -0,0 +1,145 @@
+import numpy as np
+import torch
+from torch import nn as nn
+import torch.nn.functional as F
+
+from utils.util import opt_get
+from models.modules import Basic
+from models.modules.FlowStep import FlowStep
+from models.modules.ConditionalFlow import ConditionalFlow
+
+class FlowNet(nn.Module):
+ def __init__(self, image_shape, opt=None):
+ assert image_shape[2] == 1 or image_shape[2] == 3
+ super().__init__()
+ H, W, self.C = image_shape
+ self.opt = opt
+ self.L = opt_get(opt, ['network_G', 'flowDownsampler', 'L'])
+ self.K = opt_get(opt, ['network_G', 'flowDownsampler', 'K'])
+ if isinstance(self.K, int): self.K = [self.K] * (self.L + 1)
+
+ n_additionalFlowNoAffine = opt_get(self.opt, ['network_G', 'flowDownsampler', 'additionalFlowNoAffine'], 0)
+ flow_permutation = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_permutation'], 'invconv')
+ flow_coupling = opt_get(self.opt, ['network_G', 'flowDownsampler', 'flow_coupling'], 'Affine')
+ cond_channels = opt_get(self.opt, ['network_G', 'flowDownsampler', 'cond_channels'], None)
+ enable_splitOff = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'enable'], False)
+ after_splitOff_flowStep = opt_get(opt, ['network_G', 'flowDownsampler', 'splitOff', 'after_flowstep'], 0)
+ if isinstance(after_splitOff_flowStep, int): after_splitOff_flowStep = [after_splitOff_flowStep] * (self.L + 1)
+
+ # construct flow
+ self.layers = nn.ModuleList()
+ self.output_shapes = []
+
+ for level in range(self.L):
+ # 1. Squeeze
+ self.layers.append(Basic.SqueezeLayer(factor=2)) # may need a better way for squeezing
+ self.C, H, W = self.C * 4, H // 2, W // 2
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 2. main FlowSteps (unconditional flow)
+ for k in range(self.K[level]-after_splitOff_flowStep[level]):
+ self.layers.append(FlowStep(in_channels=self.C, cond_channels=cond_channels,
+ flow_permutation=flow_permutation,
+ flow_coupling=flow_coupling,
+ opt=opt['network_G']['flowDownsampler']))
+ self.output_shapes.append([-1, self.C, H, W])
+
+ # 3. additional FlowSteps (split + conditional flow)
+ if enable_splitOff:
+ if level == 0:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level0_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=2, SR=True)
+ elif level == 1:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level1_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=1, SR=True)
+ elif level == 2:
+ self.layers.append(Basic.Split(num_channels_split=self.C // 2 if level < self.L-1 else 3, level=level))
+ self.level2_condFlow = ConditionalFlow(num_channels=self.C,
+ num_channels_split=self.C // 2 if level < self.L-1 else 3,
+ n_flow_step=after_splitOff_flowStep[level],
+ opt=opt['network_G']['flowDownsampler']['splitOff'],
+ num_levels_condition=0, SR=True)
+
+ self.C = self.C // 2 if level < self.L-1 else 3
+ self.output_shapes.append([-1, self.C, H, W])
+
+ self.H = H
+ self.W = W
+ self.scaleH = image_shape[0] / H
+ self.scaleW = image_shape[1] / W
+ print('shapes:', self.output_shapes)
+
+
+ # nodetach version: # nodetach version; 0.05 better than detach version, 0.30 better when using only nll loss
+ def forward(self, hr=None, z=None, u=None, eps_std=None, logdet=None, reverse=False, training=True):
+ if not reverse:
+ return self.normal_flow(hr, u=u, logdet=logdet, training=training)
+ else:
+ return self.reverse_flow(z, u=u, eps_std=eps_std, training=training)
+
+ '''
+ hr->y1+a1(z1)->y2+a2(z2)->y3+z3
+ '''
+ def normal_flow(self, z, u=None, logdet=None, training=True):
+ for layer, shape in zip(self.layers, self.output_shapes):
+ if isinstance(layer, FlowStep):
+ z, logdet = layer(z, u, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.SqueezeLayer):
+ z, logdet = layer(z, logdet=logdet, reverse=False)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 0:
+ z, a1 = layer(z, reverse=False)
+ y1 = z.clone()
+ elif layer.level == 1:
+ z, a2 = layer(z, reverse=False)
+ y2 = z.clone()
+ elif layer.level == 2:
+ z, a3 = layer(z, reverse=False)
+
+ logdet, conditional_feature3 = self.level2_condFlow(a3, z, logdet=logdet, reverse=False, training=training)
+
+ conditional_feature2 = torch.cat([y2, F.interpolate(conditional_feature3, scale_factor=2, mode='nearest')],1)
+ logdet, conditional_feature2 = self.level1_condFlow(a2, conditional_feature2, logdet=logdet, reverse=False, training=training)
+
+ conditional_feature1 = torch.cat([y1, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest'),
+ F.interpolate(conditional_feature3, scale_factor=4, mode='nearest')],1)
+ logdet, _ = self.level0_condFlow(a1, conditional_feature1, logdet=logdet, reverse=False, training=training)
+
+ return z, logdet
+
+ '''
+ y3+z3(a3)->y2+z2(a2)->y1+z1(a1)->hr
+ '''
+ def reverse_flow(self, z, u=None, eps_std=None, training=True):
+ for layer, shape in zip(reversed(self.layers), reversed(self.output_shapes)):
+ if isinstance(layer, FlowStep):
+ z, _ = layer(z, u, reverse=True)
+ elif isinstance(layer, Basic.SqueezeLayer):
+ z, _ = layer(z, reverse=True)
+ elif isinstance(layer, Basic.Split):
+ if layer.level == 2:
+ a3, _, conditional_feature3 = self.level2_condFlow(None, z, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a3, reverse=True)
+ elif layer.level == 1:
+ conditional_feature2 = torch.cat([z, F.interpolate(conditional_feature3, scale_factor=2, mode='nearest')],1)
+ a2, _, conditional_feature2 = self.level1_condFlow(None, conditional_feature2, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a2, reverse=True)
+ elif layer.level == 0:
+ conditional_feature1 = torch.cat([z, F.interpolate(conditional_feature2, scale_factor=2, mode='nearest'),
+ F.interpolate(conditional_feature3, scale_factor=4, mode='nearest')],1)
+ a1, _, _ = self.level0_condFlow(None, conditional_feature1, eps_std=eps_std, reverse=True, training=training)
+ z = layer(z, a1, reverse=True)
+
+
+
+
+ return z
+
diff --git a/codes/models/modules/FlowStep.py b/codes/models/modules/FlowStep.py
new file mode 100644
index 0000000..a6bc07d
--- /dev/null
+++ b/codes/models/modules/FlowStep.py
@@ -0,0 +1,65 @@
+import torch
+from torch import nn as nn
+
+from utils.util import opt_get
+from models.modules import ActNorms, Permutations, AffineCouplings
+
+
+class FlowStep(nn.Module):
+ def __init__(self, in_channels, cond_channels=None, flow_permutation='invconv', flow_coupling='Affine', LRvsothers=True,
+ actnorm_scale=1.0, LU_decomposed=False, opt=None):
+ super().__init__()
+ self.flow_permutation = flow_permutation
+ self.flow_coupling = flow_coupling
+
+ # 1. actnorm
+ self.actnorm = ActNorms.ActNorm2d(in_channels, actnorm_scale)
+
+ # 2. permute # todo: maybe hurtful for downsampling; presever the structure of downsampling
+ if self.flow_permutation == "invconv":
+ self.permute = Permutations.InvertibleConv1x1(in_channels, LU_decomposed=LU_decomposed)
+ elif self.flow_permutation == "none":
+ self.permute = None
+
+ # 3. coupling
+ if self.flow_coupling == "AffineInjector":
+ self.affine = AffineCouplings.AffineCouplingInjector(in_channels=in_channels, cond_channels=cond_channels, opt=opt)
+ elif self.flow_coupling == "noCoupling":
+ pass
+ elif self.flow_coupling == "Affine":
+ self.affine = AffineCouplings.AffineCoupling(in_channels=in_channels, cond_channels=cond_channels, opt=opt)
+ elif self.flow_coupling == "Affine3shift":
+ self.affine = AffineCouplings.AffineCoupling3shift(in_channels=in_channels, cond_channels=cond_channels, LRvsothers=LRvsothers, opt=opt)
+
+ def forward(self, z, u=None, logdet=None, reverse=False):
+ if not reverse:
+ return self.normal_flow(z, u, logdet)
+ else:
+ return self.reverse_flow(z, u)
+
+ def normal_flow(self, z, u=None, logdet=None):
+ # 1. actnorm
+ z, logdet = self.actnorm(z, logdet=logdet, reverse=False)
+
+ # 2. permute
+ if self.permute is not None:
+ z, logdet = self.permute( z, logdet=logdet, reverse=False)
+
+ # 3. coupling
+ z, logdet = self.affine(z, u=u, logdet=logdet, reverse=False)
+
+ return z, logdet
+
+ def reverse_flow(self, z, u=None, logdet=None):
+ # 1.coupling
+ z, _ = self.affine(z, u=u, reverse=True)
+
+ # 2. permute
+ if self.permute is not None:
+ z, _ = self.permute(z, reverse=True)
+
+ # 3. actnorm
+ z, _ = self.actnorm(z, reverse=True)
+
+ return z, logdet
+
diff --git a/codes/models/modules/HCFlowNet_Rescaling_arch.py b/codes/models/modules/HCFlowNet_Rescaling_arch.py
new file mode 100644
index 0000000..c4839c7
--- /dev/null
+++ b/codes/models/modules/HCFlowNet_Rescaling_arch.py
@@ -0,0 +1,60 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from utils.util import opt_get
+from models.modules.FlowNet_Rescaling_x4 import FlowNet
+from models.modules import Basic, thops
+
+
+
+class HCFlowNet_Rescaling(nn.Module):
+ def __init__(self, opt, step=None):
+ super(HCFlowNet_Rescaling, self).__init__()
+ self.opt = opt
+ self.quant = opt_get(opt, ['datasets', 'train', 'quant'], 256)
+
+ hr_size = opt_get(opt, ['datasets', 'train', 'GT_size'], 160)
+ hr_channel = opt_get(opt, ['network_G', 'in_nc'], 3)
+
+ # hr->lr+z
+ self.flow = FlowNet((hr_size, hr_size, hr_channel), opt=opt)
+
+ # hr: HR image, lr: LR image, z: latent variable, u: conditional variable
+ def forward(self, hr=None, lr=None, z=None, u=None, eps_std=None,
+ add_gt_noise=False, step=None, reverse=False, training=True):
+
+ # hr->z
+ if not reverse:
+ return self.normal_flow_diracLR(hr, lr, u, step=step, training=training)
+ # z->hr
+ else: # setting z to lr!!!
+ return self.reverse_flow_diracLR(lr, z, u, eps_std=eps_std, training=training)
+
+
+ #########################################diracLR
+ # hr->lr+z, diracLR
+ def normal_flow_diracLR(self, hr, lr, u=None, step=None, training=True):
+ # 1. quantitize HR
+ # hr = hr + (torch.rand(hr.shape, device=hr.device)) / self.quant # no quantization is better
+
+ # 2. hr->lr+z
+ fake_lr_from_hr, fake_z1, fake_z2 = self.flow(hr=hr, u=u, logdet=None, reverse=False, training=training)
+
+ return torch.clamp(fake_lr_from_hr, 0, 1), fake_z1, fake_z2
+
+ # lr+z->hr
+ def reverse_flow_diracLR(self, lr, z, u, eps_std, training=True):
+
+ # lr+z->hr
+ fake_hr = self.flow(z=lr, u=u, eps_std=eps_std, reverse=True, training=training)
+
+ return torch.clamp(fake_hr, 0, 1)
+
+
+ def get_score(self, disc_loss_sigma, z):
+ score_real = 0.5 * (1 - 1 / (disc_loss_sigma ** 2)) * thops.sum(z ** 2, dim=[1, 2, 3]) - \
+ z.shape[1] * z.shape[2] * z.shape[3] * math.log(disc_loss_sigma)
+ return -score_real
\ No newline at end of file
diff --git a/codes/models/modules/HCFlowNet_SR_arch.py b/codes/models/modules/HCFlowNet_SR_arch.py
new file mode 100644
index 0000000..12214f7
--- /dev/null
+++ b/codes/models/modules/HCFlowNet_SR_arch.py
@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from utils.util import opt_get
+from models.modules import Basic, thops
+
+
+
+class HCFlowNet_SR(nn.Module):
+ def __init__(self, opt, step=None):
+ super(HCFlowNet_SR, self).__init__()
+ self.opt = opt
+ self.quant = opt_get(opt, ['quant'], 256)
+
+ hr_size = opt_get(opt, ['datasets', 'train', 'GT_size'], 160)
+ hr_channel = opt_get(opt, ['network_G', 'in_nc'], 3)
+ scale = opt_get(opt, ['scale'])
+
+ if scale == 4:
+ from models.modules.FlowNet_SR_x4 import FlowNet
+ elif scale == 8:
+ from models.modules.FlowNet_SR_x8 import FlowNet
+ else:
+ raise NotImplementedError('Scale {} is not implemented'.format(scale))
+
+ # hr->lr+z
+ self.flow = FlowNet((hr_size, hr_size, hr_channel), opt=opt)
+
+ self.quantization = Basic.Quantization()
+
+ # hr: HR image, lr: LR image, z: latent variable, u: conditional variable
+ def forward(self, hr=None, lr=None, z=None, u=None, eps_std=None,
+ add_gt_noise=False, step=None, reverse=False, training=True):
+
+ # hr->z
+ if not reverse:
+ return self.normal_flow_diracLR(hr, lr, u, step=step, training=training)
+ # z->hr
+ else:
+ return self.reverse_flow_diracLR(lr, z, u, eps_std=eps_std, training=training)
+
+
+ #########################################diracLR
+ # hr->lr+z, diracLR
+ def normal_flow_diracLR(self, hr, lr, u=None, step=None, training=True):
+ # 1. quantitize HR
+ pixels = thops.pixels(hr)
+
+ # according to Glow and ours, it should be u~U(0,a) (0.06 better in practice), not u~U(-0.5,0.5) (though better in theory)
+ hr = hr + (torch.rand(hr.shape, device=hr.device)) / self.quant
+ logdet = torch.zeros_like(hr[:, 0, 0, 0]) + float(-np.log(self.quant) * pixels)
+
+ # 2. hr->lr+z
+ fake_lr_from_hr, logdet = self.flow(hr=hr, u=u, logdet=logdet, reverse=False, training=training)
+
+ # note in rescaling, we use LR for LR loss before quantization
+ fake_lr_from_hr = self.quantization(fake_lr_from_hr)
+
+ # 3. loss, Gaussian with small variance to approximate Dirac delta function of LR.
+ # for the second term, using small log-variance may lead to svd problem, for both exp and tanh version
+ objective = logdet + Basic.GaussianDiag.logp(lr, -torch.ones_like(lr)*6, fake_lr_from_hr)
+
+ nll = ((-objective) / float(np.log(2.) * pixels)).mean()
+
+ return torch.clamp(fake_lr_from_hr, 0, 1), nll
+
+ # lr+z->hr
+ def reverse_flow_diracLR(self, lr, z, u, eps_std, training=True):
+
+ # lr+z->hr
+ fake_hr = self.flow(z=lr, u=u, eps_std=eps_std, reverse=True, training=training)
+
+ return torch.clamp(fake_hr, 0, 1)
diff --git a/codes/models/modules/Permutations.py b/codes/models/modules/Permutations.py
new file mode 100644
index 0000000..a58773c
--- /dev/null
+++ b/codes/models/modules/Permutations.py
@@ -0,0 +1,108 @@
+import numpy as np
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+import scipy.linalg
+
+from models.modules import thops
+
+
+class Permute2d(nn.Module):
+ def __init__(self, num_channels, shuffle):
+ super().__init__()
+ self.num_channels = num_channels
+ self.indices = np.arange(self.num_channels - 1, -1, -1).astype(np.long)
+ self.indices_inverse = np.zeros((self.num_channels), dtype=np.long)
+ for i in range(self.num_channels):
+ self.indices_inverse[self.indices[i]] = i
+ if shuffle:
+ self.reset_indices()
+
+ def reset_indices(self):
+ np.random.shuffle(self.indices)
+ for i in range(self.num_channels):
+ self.indices_inverse[self.indices[i]] = i
+
+ def forward(self, input, logdet=None, reverse=False):
+ if not reverse:
+ return input[:, self.indices, :, :], logdet
+ else:
+ return input[:, self.indices_inverse, :, :], logdet
+
+
+class InvertibleConv1x1(nn.Module):
+ def __init__(self, num_channels, LU_decomposed=False):
+ super().__init__()
+ w_shape = [num_channels, num_channels]
+ w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32)
+ if not LU_decomposed:
+ # Sample a random orthogonal matrix:
+ self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))
+ else:
+ # W = PL(U+diag(s))
+ np_p, np_l, np_u = scipy.linalg.lu(w_init)
+ np_s = np.diag(np_u)
+ np_sign_s = np.sign(np_s)
+ np_log_s = np.log(np.abs(np_s))
+ np_u = np.triu(np_u, k=1)
+ l_mask = np.tril(np.ones(w_shape, dtype=np.float32), -1)
+ eye = np.eye(*w_shape, dtype=np.float32)
+
+ self.register_buffer('p', torch.Tensor(np_p.astype(np.float32))) # remains fixed
+ self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(np.float32))) # the sign is fixed
+ self.l = nn.Parameter(torch.Tensor(np_l.astype(np.float32))) # optimized except diagonal 1
+ self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(np.float32)))
+ self.u = nn.Parameter(torch.Tensor(np_u.astype(np.float32))) # optimized
+ self.l_mask = torch.Tensor(l_mask)
+ self.eye = torch.Tensor(eye)
+ self.w_shape = w_shape
+ self.LU = LU_decomposed
+
+ def get_weight(self, input, reverse):
+ # The difference in computational cost will become significant for large c, although for the networks in
+ # our experiments we did not measure a large difference in wallclock computation time.
+ if not self.LU:
+ if not reverse:
+ # pixels = thops.pixels(input)
+ # GPU version
+ # dlogdet = torch.slogdet(self.weight)[1] * pixels
+ # CPU version is 2x faster, https://github.com/didriknielsen/survae_flows/issues/5.
+ dlogdet = (torch.slogdet(self.weight.to('cpu'))[1] * thops.pixels(input)).to(self.weight.device)
+ weight = self.weight.view(self.w_shape[0], self.w_shape[1], 1, 1)
+ else:
+ dlogdet = 0
+ weight = torch.inverse(self.weight.double()).float().view(self.w_shape[0], self.w_shape[1], 1, 1)
+
+
+ return weight, dlogdet
+ else:
+ self.p = self.p.to(input.device)
+ self.sign_s = self.sign_s.to(input.device)
+ self.l_mask = self.l_mask.to(input.device)
+ self.eye = self.eye.to(input.device)
+ l = self.l * self.l_mask + self.eye
+ u = self.u * self.l_mask.transpose(0, 1).contiguous() + torch.diag(self.sign_s * torch.exp(self.log_s))
+ dlogdet = thops.sum(self.log_s) * thops.pixels(input)
+ if not reverse:
+ w = torch.matmul(self.p, torch.matmul(l, u))
+ else:
+ l = torch.inverse(l.double()).float()
+ u = torch.inverse(u.double()).float()
+ w = torch.matmul(u, torch.matmul(l, self.p.inverse()))
+ return w.view(self.w_shape[0], self.w_shape[1], 1, 1), dlogdet
+
+ def forward(self, input, logdet=None, reverse=False):
+ """
+ log-det = log|abs(|W|)| * pixels
+ """
+ weight, dlogdet = self.get_weight(input, reverse)
+ if not reverse:
+ z = F.conv2d(input, weight) # fc layer, ie, permute channel
+ if logdet is not None:
+ logdet = logdet + dlogdet
+ return z, logdet
+ else:
+ z = F.conv2d(input, weight)
+ if logdet is not None:
+ logdet = logdet - dlogdet
+ return z, logdet
diff --git a/codes/models/modules/discriminator_vgg_arch.py b/codes/models/modules/discriminator_vgg_arch.py
new file mode 100644
index 0000000..d029266
--- /dev/null
+++ b/codes/models/modules/discriminator_vgg_arch.py
@@ -0,0 +1,189 @@
+import torch
+import torch.nn as nn
+import torchvision
+
+
+class Discriminator_VGG_128(nn.Module):
+ def __init__(self, in_nc, nf):
+ super(Discriminator_VGG_128, self).__init__()
+ # [64, 128, 128]
+ self.conv0_0 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+ self.conv0_1 = nn.Conv2d(nf, nf, 4, 2, 1, bias=False)
+ self.bn0_1 = nn.BatchNorm2d(nf, affine=True)
+ # [64, 64, 64]
+ self.conv1_0 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=False)
+ self.bn1_0 = nn.BatchNorm2d(nf * 2, affine=True)
+ self.conv1_1 = nn.Conv2d(nf * 2, nf * 2, 4, 2, 1, bias=False)
+ self.bn1_1 = nn.BatchNorm2d(nf * 2, affine=True)
+ # [128, 32, 32]
+ self.conv2_0 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=False)
+ self.bn2_0 = nn.BatchNorm2d(nf * 4, affine=True)
+ self.conv2_1 = nn.Conv2d(nf * 4, nf * 4, 4, 2, 1, bias=False)
+ self.bn2_1 = nn.BatchNorm2d(nf * 4, affine=True)
+ # [256, 16, 16]
+ self.conv3_0 = nn.Conv2d(nf * 4, nf * 8, 3, 1, 1, bias=False)
+ self.bn3_0 = nn.BatchNorm2d(nf * 8, affine=True)
+ self.conv3_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
+ self.bn3_1 = nn.BatchNorm2d(nf * 8, affine=True)
+ # [512, 8, 8]
+ self.conv4_0 = nn.Conv2d(nf * 8, nf * 8, 3, 1, 1, bias=False)
+ self.bn4_0 = nn.BatchNorm2d(nf * 8, affine=True)
+ self.conv4_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
+ self.bn4_1 = nn.BatchNorm2d(nf * 8, affine=True)
+
+ self.linear1 = nn.Linear(512 * 4 * 4, 100)
+ self.linear2 = nn.Linear(100, 1)
+
+ # activation function
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ def forward(self, x):
+ fea = self.lrelu(self.conv0_0(x))
+ fea = self.lrelu(self.bn0_1(self.conv0_1(fea)))
+
+ fea = self.lrelu(self.bn1_0(self.conv1_0(fea)))
+ fea = self.lrelu(self.bn1_1(self.conv1_1(fea)))
+
+ fea = self.lrelu(self.bn2_0(self.conv2_0(fea)))
+ fea = self.lrelu(self.bn2_1(self.conv2_1(fea)))
+
+ fea = self.lrelu(self.bn3_0(self.conv3_0(fea)))
+ fea = self.lrelu(self.bn3_1(self.conv3_1(fea)))
+
+ fea = self.lrelu(self.bn4_0(self.conv4_0(fea)))
+ fea = self.lrelu(self.bn4_1(self.conv4_1(fea)))
+
+ fea = fea.view(fea.size(0), -1)
+ fea = self.lrelu(self.linear1(fea))
+ out = self.linear2(fea)
+ return out
+
+ def reset_parameters(self):
+ for layer in self.children():
+ if hasattr(layer, 'reset_parameters'):
+ layer.reset_parameters()
+
+
+
+class Discriminator_VGG_160(nn.Module):
+ def __init__(self, in_nc, nf):
+ super(Discriminator_VGG_160, self).__init__()
+ # [64, 160, 160]
+ self.conv0_0 = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+ self.conv0_1 = nn.Conv2d(nf, nf, 4, 2, 1, bias=False)
+ self.bn0_1 = nn.BatchNorm2d(nf, affine=True)
+ # [64, 80, 80]
+ self.conv1_0 = nn.Conv2d(nf, nf * 2, 3, 1, 1, bias=False)
+ self.bn1_0 = nn.BatchNorm2d(nf * 2, affine=True)
+ self.conv1_1 = nn.Conv2d(nf * 2, nf * 2, 4, 2, 1, bias=False)
+ self.bn1_1 = nn.BatchNorm2d(nf * 2, affine=True)
+ # [128, 40, 40]
+ self.conv2_0 = nn.Conv2d(nf * 2, nf * 4, 3, 1, 1, bias=False)
+ self.bn2_0 = nn.BatchNorm2d(nf * 4, affine=True)
+ self.conv2_1 = nn.Conv2d(nf * 4, nf * 4, 4, 2, 1, bias=False)
+ self.bn2_1 = nn.BatchNorm2d(nf * 4, affine=True)
+ # [256, 20, 20]
+ self.conv3_0 = nn.Conv2d(nf * 4, nf * 8, 3, 1, 1, bias=False)
+ self.bn3_0 = nn.BatchNorm2d(nf * 8, affine=True)
+ self.conv3_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
+ self.bn3_1 = nn.BatchNorm2d(nf * 8, affine=True)
+ # [512, 10, 10]
+ self.conv4_0 = nn.Conv2d(nf * 8, nf * 8, 3, 1, 1, bias=False)
+ self.bn4_0 = nn.BatchNorm2d(nf * 8, affine=True)
+ self.conv4_1 = nn.Conv2d(nf * 8, nf * 8, 4, 2, 1, bias=False)
+ self.bn4_1 = nn.BatchNorm2d(nf * 8, affine=True)
+ # [512, 5, 5]
+
+ self.linear1 = nn.Linear(512 * 5 * 5, 100)
+ self.linear2 = nn.Linear(100, 1)
+
+ # activation function
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ def forward(self, x):
+ fea = self.lrelu(self.conv0_0(x))
+ fea = self.lrelu(self.bn0_1(self.conv0_1(fea)))
+
+ fea = self.lrelu(self.bn1_0(self.conv1_0(fea)))
+ fea = self.lrelu(self.bn1_1(self.conv1_1(fea)))
+
+ fea = self.lrelu(self.bn2_0(self.conv2_0(fea)))
+ fea = self.lrelu(self.bn2_1(self.conv2_1(fea)))
+
+ fea = self.lrelu(self.bn3_0(self.conv3_0(fea)))
+ fea = self.lrelu(self.bn3_1(self.conv3_1(fea)))
+
+ fea = self.lrelu(self.bn4_0(self.conv4_0(fea)))
+ fea = self.lrelu(self.bn4_1(self.conv4_1(fea)))
+
+ fea = fea.view(fea.size(0), -1)
+ fea = self.lrelu(self.linear1(fea))
+ out = self.linear2(fea)
+ return out
+
+ def reset_parameters(self):
+ for layer in self.children():
+ if hasattr(layer, 'reset_parameters'):
+ layer.reset_parameters()
+
+
+class VGGFeatureExtractor(nn.Module):
+ def __init__(self, feature_layer=34, use_bn=False, use_input_norm=True,
+ device=torch.device('cpu')):
+ super(VGGFeatureExtractor, self).__init__()
+ self.use_input_norm = use_input_norm
+ if use_bn:
+ model = torchvision.models.vgg19_bn(pretrained=True)
+ else:
+ model = torchvision.models.vgg19(pretrained=True)
+ if self.use_input_norm:
+ mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+ # [0.485 - 1, 0.456 - 1, 0.406 - 1] if input in range [-1, 1]
+ std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+ # [0.229 * 2, 0.224 * 2, 0.225 * 2] if input in range [-1, 1]
+ self.register_buffer('mean', mean)
+ self.register_buffer('std', std)
+ self.features = nn.Sequential(*list(model.features.children())[:(feature_layer + 1)])
+ # No need to BP to variable
+ for k, v in self.features.named_parameters():
+ v.requires_grad = False
+
+ def forward(self, x):
+ # Assume input range is [0, 1]
+ if self.use_input_norm:
+ x = (x - self.mean) / self.std
+ output = self.features(x)
+ return output
+
+
+class PatchGANDiscriminator(nn.Module):
+ """Defines a PatchGAN discriminator"""
+
+ def __init__(self, in_nc=3, ndf=64, n_layers=35, norm_layer=nn.BatchNorm2d):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(PatchGANDiscriminator, self).__init__()
+ use_bias = False
+ kw = 3
+ padw = 0
+ sequence = [nn.Conv2d(in_nc, ndf, kernel_size=kw, stride=1, padding=padw), nn.LeakyReLU(0.2, True)]
+
+ for i in range(0, n_layers):
+ sequence += [
+ nn.Conv2d(ndf, ndf, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+ norm_layer(ndf),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ sequence += [nn.Conv2d(ndf, 1, kernel_size=kw, stride=1, padding=padw, bias=use_bias)] # output 1 channel prediction map
+ # TODO
+ self.model = nn.Sequential(*sequence)
+
+ def forward(self, x):
+ """Standard forward."""
+ return self.model(x)
diff --git a/codes/models/modules/loss.py b/codes/models/modules/loss.py
new file mode 100644
index 0000000..41b528b
--- /dev/null
+++ b/codes/models/modules/loss.py
@@ -0,0 +1,91 @@
+import torch
+import torch.nn as nn
+
+
+class CharbonnierLoss(nn.Module):
+ """Charbonnier Loss (L1)"""
+
+ def __init__(self, eps=1e-6):
+ super(CharbonnierLoss, self).__init__()
+ self.eps = eps
+
+ def forward(self, x, y):
+ diff = x - y
+ loss = torch.sum(torch.sqrt(diff * diff + self.eps))
+ return loss
+
+
+# Define GAN loss: [gan(vanilla) | lsgan | wgan-gp | ragan]
+class GANLoss(nn.Module):
+ def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0):
+ super(GANLoss, self).__init__()
+ self.gan_type = gan_type.lower()
+ self.real_label_val = real_label_val
+ self.fake_label_val = fake_label_val
+
+ if self.gan_type == 'gan' or self.gan_type == 'ragan':
+ self.loss = nn.BCEWithLogitsLoss()
+ elif self.gan_type == 'lsgan':
+ self.loss = nn.MSELoss()
+ elif self.gan_type == 'wgan-gp':
+
+ def wgan_loss(input, target):
+ # target is boolean
+ return -1 * input.mean() if target else input.mean()
+
+ self.loss = wgan_loss
+ else:
+ raise NotImplementedError('GAN type [{:s}] is not found'.format(self.gan_type))
+
+ def get_target_label(self, input, target_is_real):
+ if self.gan_type == 'wgan-gp':
+ return target_is_real
+ if target_is_real:
+ return torch.empty_like(input).fill_(self.real_label_val)
+ else:
+ return torch.empty_like(input).fill_(self.fake_label_val)
+
+ def forward(self, input, target_is_real):
+ target_label = self.get_target_label(input, target_is_real)
+ loss = self.loss(input, target_label)
+ return loss
+
+
+class GradientPenaltyLoss(nn.Module):
+ def __init__(self, device=torch.device('cpu')):
+ super(GradientPenaltyLoss, self).__init__()
+ self.register_buffer('grad_outputs', torch.Tensor())
+ self.grad_outputs = self.grad_outputs.to(device)
+
+ def get_grad_outputs(self, input):
+ if self.grad_outputs.size() != input.size():
+ self.grad_outputs.resize_(input.size()).fill_(1.0)
+ return self.grad_outputs
+
+ def forward(self, interp, interp_crit):
+ grad_outputs = self.get_grad_outputs(interp_crit)
+ grad_interp = torch.autograd.grad(outputs=interp_crit, inputs=interp,
+ grad_outputs=grad_outputs, create_graph=True,
+ retain_graph=True, only_inputs=True)[0]
+ grad_interp = grad_interp.view(grad_interp.size(0), -1)
+ grad_interp_norm = grad_interp.norm(2, dim=1)
+
+ loss = ((grad_interp_norm - 1)**2).mean()
+ return loss
+
+class ReconstructionLoss(nn.Module):
+ def __init__(self, losstype='l2', eps=1e-6):
+ super(ReconstructionLoss, self).__init__()
+ self.losstype = losstype
+ self.eps = eps
+
+ def forward(self, x, target):
+ if self.losstype == 'l2':
+ return torch.mean(torch.sum((x - target)**2, (1, 2, 3)))
+ elif self.losstype == 'l1':
+ diff = x - target
+ return torch.mean(torch.sum(torch.sqrt(diff * diff + self.eps), (1, 2, 3)))
+ else:
+ print("reconstruction loss type error!")
+ return 0
+
diff --git a/codes/models/modules/module_util.py b/codes/models/modules/module_util.py
new file mode 100644
index 0000000..166355f
--- /dev/null
+++ b/codes/models/modules/module_util.py
@@ -0,0 +1,98 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+
+
+def initialize_weights(net_l, scale=1):
+ if not isinstance(net_l, list):
+ net_l = [net_l]
+ for net in net_l:
+ for m in net.modules():
+ if isinstance(m, nn.Conv2d):
+ init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+ m.weight.data *= scale # for residual block
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.Linear):
+ init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+ m.weight.data *= scale
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.BatchNorm2d):
+ init.constant_(m.weight, 1)
+ init.constant_(m.bias.data, 0.0)
+
+def initialize_weights_xavier(net_l, scale=1):
+ if not isinstance(net_l, list):
+ net_l = [net_l]
+ for net in net_l:
+ for m in net.modules():
+ if isinstance(m, nn.Conv2d):
+ init.xavier_normal_(m.weight)
+ m.weight.data *= scale # for residual block
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.Linear):
+ init.xavier_normal_(m.weight)
+ m.weight.data *= scale
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.BatchNorm2d):
+ init.constant_(m.weight, 1)
+ init.constant_(m.bias.data, 0.0)
+
+
+def make_layer(block, n_layers):
+ layers = []
+ for _ in range(n_layers):
+ layers.append(block())
+ return nn.Sequential(*layers)
+
+
+class ResidualBlock_noBN(nn.Module):
+ '''Residual block w/o BN
+ ---Conv-ReLU-Conv-+-
+ |________________|
+ '''
+
+ def __init__(self, nf=64):
+ super(ResidualBlock_noBN, self).__init__()
+ self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+
+ # initialization
+ initialize_weights([self.conv1, self.conv2], 0.1)
+
+ def forward(self, x):
+ identity = x
+ out = F.relu(self.conv1(x), inplace=True)
+ out = self.conv2(out)
+ return identity + out
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros'):
+ """Warp an image or feature map with optical flow
+ Args:
+ x (Tensor): size (N, C, H, W)
+ flow (Tensor): size (N, H, W, 2), normal value
+ interp_mode (str): 'nearest' or 'bilinear'
+ padding_mode (str): 'zeros' or 'border' or 'reflection'
+
+ Returns:
+ Tensor: warped image or feature map
+ """
+ assert x.size()[-2:] == flow.size()[1:3]
+ B, C, H, W = x.size()
+ # mesh grid
+ grid_y, grid_x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W))
+ grid = torch.stack((grid_x, grid_y), 2).float() # W(x), H(y), 2
+ grid.requires_grad = False
+ grid = grid.type_as(x)
+ vgrid = grid + flow
+ # scale grid to [-1,1]
+ vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(W - 1, 1) - 1.0
+ vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(H - 1, 1) - 1.0
+ vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+ output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode)
+ return output
diff --git a/codes/models/modules/thops.py b/codes/models/modules/thops.py
new file mode 100644
index 0000000..777d40a
--- /dev/null
+++ b/codes/models/modules/thops.py
@@ -0,0 +1,53 @@
+import torch
+
+
+def sum(tensor, dim=None, keepdim=False):
+ if dim is None:
+ # sum up all dim
+ return torch.sum(tensor)
+ else:
+ if isinstance(dim, int):
+ dim = [dim]
+ dim = sorted(dim)
+ for d in dim:
+ tensor = tensor.sum(dim=d, keepdim=True)
+ if not keepdim:
+ for i, d in enumerate(dim):
+ tensor.squeeze_(d-i)
+ return tensor
+
+
+def mean(tensor, dim=None, keepdim=False):
+ if dim is None:
+ # mean all dim
+ return torch.mean(tensor)
+ else:
+ if isinstance(dim, int):
+ dim = [dim]
+ dim = sorted(dim)
+ for d in dim:
+ tensor = tensor.mean(dim=d, keepdim=True)
+ if not keepdim:
+ for i, d in enumerate(dim):
+ tensor.squeeze_(d-i)
+ return tensor
+
+
+
+def split_feature(tensor, type="split"):
+ """
+ type = ["split", "cross"]
+ """
+ C = tensor.size(1)
+ if type == "split":
+ return tensor[:, :C // 2, ...], tensor[:, C // 2:, ...]
+ elif type == "cross":
+ return tensor[:, 0::2, ...], tensor[:, 1::2, ...]
+
+
+def cat_feature(tensor_a, tensor_b):
+ return torch.cat((tensor_a, tensor_b), dim=1)
+
+
+def pixels(tensor):
+ return int(tensor.size(2) * tensor.size(3))
\ No newline at end of file
diff --git a/codes/models/networks.py b/codes/models/networks.py
new file mode 100644
index 0000000..e605266
--- /dev/null
+++ b/codes/models/networks.py
@@ -0,0 +1,71 @@
+import importlib
+import logging
+import torch
+import models.modules.discriminator_vgg_arch as SRGAN_arch
+
+logger = logging.getLogger('base')
+
+
+def find_model_using_name(model_name):
+ model_filename = "models.modules." + model_name + "_arch"
+ modellib = importlib.import_module(model_filename)
+
+ model = None
+ target_model_name = model_name.replace('_Net', '')
+ for name, cls in modellib.__dict__.items():
+ if name.lower() == target_model_name.lower():
+ model = cls
+
+ if model is None:
+ print(
+ "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s." % (
+ model_filename, target_model_name))
+ exit(0)
+
+ return model
+
+def define_Flow(opt, step):
+ opt_net = opt['network_G']
+ which_model = opt_net['which_model_G']
+
+ Arch = find_model_using_name(which_model)
+ netG = Arch(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
+ nf=opt_net['nf'], nb=opt_net['nb'], scale=opt['scale'], K=opt_net['flow']['K'], opt=opt, step=step)
+ return netG
+
+def define_G(opt, step):
+ which_model = opt['network_G']['which_model_G']
+
+ Arch = find_model_using_name(which_model)
+ netG = Arch(opt=opt, step=step)
+ return netG
+
+#### Discriminator
+def define_D(opt):
+ opt_net = opt['network_D']
+ which_model = opt_net['which_model_D']
+
+ if which_model == 'discriminator_vgg_128':
+ netD = SRGAN_arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], nf=opt_net['nf'])
+ elif which_model == 'discriminator_vgg_160':
+ netD = SRGAN_arch.Discriminator_VGG_160(in_nc=opt_net['in_nc'], nf=opt_net['nf'])
+ elif which_model == 'PatchGANDiscriminator':
+ netD = SRGAN_arch.PatchGANDiscriminator(in_nc=opt_net['in_nc'], ndf=opt_net['ndf'], n_layers=opt_net['n_layers'],)
+ else:
+ raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
+ return netD
+
+
+#### Define Network used for Perceptual Loss
+def define_F(opt, use_bn=False):
+ gpu_ids = opt['gpu_ids']
+ device = torch.device('cuda' if gpu_ids else 'cpu')
+ # PyTorch pretrained VGG19-54, before ReLU.
+ if use_bn:
+ feature_layer = 49
+ else:
+ feature_layer = 34
+ netF = SRGAN_arch.VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn,
+ use_input_norm=True, device=device)
+ netF.eval() # No need to train
+ return netF
diff --git a/codes/options/options.py b/codes/options/options.py
new file mode 100644
index 0000000..bc2b07f
--- /dev/null
+++ b/codes/options/options.py
@@ -0,0 +1,138 @@
+import os
+import os.path as osp
+import logging
+import yaml
+from utils.util import OrderedYaml
+
+Loader, Dumper = OrderedYaml()
+
+
+def parse(opt_path, gpu_ids=None, is_train=True):
+ with open(opt_path, mode='r') as f:
+ opt = yaml.load(f, Loader=Loader)
+ # export CUDA_VISIBLE_DEVICES
+ if gpu_ids is not None: opt['gpu_ids'] = [int(x) for x in gpu_ids.split(',')]
+ gpu_list = ','.join(str(x) for x in opt['gpu_ids'])
+ os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+ print('exporting CUDA_VISIBLE_DEVICES=' + gpu_list)
+
+ opt['is_train'] = is_train
+ if opt['distortion'] == 'sr':
+ scale = opt['scale']
+
+ if 'datasets' in opt:
+ # datasets
+ for phase, dataset in opt['datasets'].items():
+ phase = phase.split('_')[0]
+ print(dataset)
+ dataset['phase'] = phase
+ if opt['distortion'] == 'sr':
+ dataset['scale'] = scale
+ is_lmdb = False
+ if dataset.get('dataroot_GT', None) is not None:
+ dataset['dataroot_GT'] = osp.expanduser(dataset['dataroot_GT'])
+ if dataset['dataroot_GT'].endswith('lmdb'):
+ is_lmdb = True
+ # if dataset.get('dataroot_GT_bg', None) is not None:
+ # dataset['dataroot_GT_bg'] = osp.expanduser(dataset['dataroot_GT_bg'])
+ if dataset.get('dataroot_LQ', None) is not None:
+ dataset['dataroot_LQ'] = osp.expanduser(dataset['dataroot_LQ'])
+ if dataset['dataroot_LQ'].endswith('lmdb'):
+ is_lmdb = True
+ dataset['data_type'] = 'lmdb' if is_lmdb else 'img'
+ if dataset['mode'].endswith('mc'): # for memcached
+ dataset['data_type'] = 'mc'
+ dataset['mode'] = dataset['mode'].replace('_mc', '')
+
+ # path
+ for key, path in opt['path'].items():
+ if path and key in opt['path'] and key != 'strict_load':
+ opt['path'][key] = osp.expanduser(path)
+ opt['path']['root'] = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir, osp.pardir))
+
+ if is_train:
+ experiments_root = osp.join(opt['path']['root'], 'experiments', opt['name'])
+ opt['path']['experiments_root'] = experiments_root
+ opt['path']['models'] = osp.join(experiments_root, 'models')
+ opt['path']['training_state'] = osp.join(experiments_root, 'training_state')
+ opt['path']['log'] = experiments_root
+ opt['path']['val_images'] = osp.join(experiments_root, 'val_images')
+
+ # change some options for debug mode
+ if 'debug' in opt['name']:
+ opt['train']['val_freq'] = 8
+ opt['logger']['print_freq'] = 1
+ opt['logger']['save_checkpoint_freq'] = 8
+ else: # test
+ results_root = osp.join(opt['path']['root'], 'results', opt['name'])
+ opt['path']['results_root'] = results_root
+ opt['path']['log'] = results_root
+
+
+ # network
+ if opt['distortion'] == 'sr':
+ opt['network_G']['scale'] = scale
+
+ # relative learning rate
+ if 'train' in opt:
+ niter = opt['train']['niter']
+ if 'T_period_rel' in opt['train']:
+ opt['train']['T_period'] = [int(x * niter) for x in opt['train']['T_period_rel']]
+ if 'restarts_rel' in opt['train']:
+ opt['train']['restarts'] = [int(x * niter) for x in opt['train']['restarts_rel']]
+ if 'lr_steps_rel' in opt['train']:
+ opt['train']['lr_steps'] = [int(x * niter) for x in opt['train']['lr_steps_rel']]
+ if 'lr_steps_inverse_rel' in opt['train']:
+ opt['train']['lr_steps_inverse'] = [int(x * niter) for x in opt['train']['lr_steps_inverse_rel']]
+ print(opt['train'])
+
+
+ return opt
+
+
+def dict2str(opt, indent_l=1):
+ '''dict to string for logger'''
+ msg = ''
+ for k, v in opt.items():
+ if isinstance(v, dict):
+ msg += ' ' * (indent_l * 2) + k + ':[\n'
+ msg += dict2str(v, indent_l + 1)
+ msg += ' ' * (indent_l * 2) + ']\n'
+ else:
+ msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
+ return msg
+
+
+class NoneDict(dict):
+ def __missing__(self, key):
+ return None
+
+
+# convert to NoneDict, which return None for missing key.
+def dict_to_nonedict(opt):
+ if isinstance(opt, dict):
+ new_opt = dict()
+ for key, sub_opt in opt.items():
+ new_opt[key] = dict_to_nonedict(sub_opt)
+ return NoneDict(**new_opt)
+ elif isinstance(opt, list):
+ return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+ else:
+ return opt
+
+
+def check_resume(opt, resume_iter):
+ '''Check resume states and pretrain_model paths (overriding pretrain_paths)'''
+ logger = logging.getLogger('base')
+ if opt['path']['resume_state']:
+ if opt['path'].get('pretrain_model_G', None) is not None or opt['path'].get(
+ 'pretrain_model_D', None) is not None:
+ logger.warning('pretrain_model path will be ignored when resuming training.')
+
+ opt['path']['pretrain_model_G'] = osp.join(opt['path']['models'], '{}_G.pth'.format(resume_iter))
+ logger.info('Set [pretrain_model_G] to ' + opt['path']['pretrain_model_G'])
+
+ if opt['train']['gan_weight'] > 0:
+ opt['path']['pretrain_model_D'] = osp.join(opt['path']['models'],
+ '{}_D.pth'.format(resume_iter))
+ logger.info('Set [pretrain_model_D] to ' + opt['path']['pretrain_model_D'])
diff --git a/codes/options/test/test_Rescaling_DF2K_4X_HCFlow.yml b/codes/options/test/test_Rescaling_DF2K_4X_HCFlow.yml
new file mode 100644
index 0000000..178a05c
--- /dev/null
+++ b/codes/options/test/test_Rescaling_DF2K_4X_HCFlow.yml
@@ -0,0 +1,93 @@
+#### general settings
+name: 003_HCFlow_DF2K_x4_rescaling_test
+suffix: ~
+use_tb_logger: true
+model: HCFlow_Rescaling
+distortion: sr
+scale: 4
+gpu_ids: [0]
+
+
+datasets:
+ test0:
+ name: example
+ mode: GTLQ
+ dataroot_GT: ../datasets/example_general_4X/HR
+ dataroot_LQ: ../datasets/example_general_4X/LR
+
+# test_1:
+# name: Set5
+# mode: GTLQx
+# dataroot_GT: ../datasets/Set5/HR
+# dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+#
+# test_2:
+# name: Set14
+# mode: GTLQx
+# dataroot_GT: ../datasets/Set14/HR
+# dataroot_LQ: ../datasets/Set14/LR_bicubic/X4
+#
+# test_3:
+# name: BSD100
+# mode: GTLQx
+# dataroot_GT: ../datasets/BSD100/HR
+# dataroot_LQ: ../datasets/BSD100/LR_bicubic/X4
+#
+# test_4:
+# name: Urban100
+# mode: GTLQx
+# dataroot_GT: ../datasets/Urban100/HR
+# dataroot_LQ: ../datasets/Urban100/LR_bicubic/X4
+#
+# test_5:
+# name: DIV2K-validation
+# mode: GTLQx
+# dataroot_GT: ../datasets/DIV2K/HR
+# dataroot_LQ: ../datasets/DIV2K/LR_bicubic/X4
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_Rescaling
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 14
+ L: 2
+ squeeze: haar # better than squeeze2d
+ flow_permutation: none # bettter than invconv
+ flow_coupling: Affine3shift # better than affine
+ nn_module: DenseBlock # better than FCN
+ hidden_channels: 32
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [6, 6]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ stage1: True
+ feature_extractor: RRDB
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [2,1]
+ RRDB_nf: 64
+ RRDB_gc: 16
+
+
+
+#### validation settings
+val:
+ heats: [1.0]
+ n_sample: 1
+
+
+path:
+ strict_load: true
+ load_submodule: ~
+ pretrain_model_G: ../experiments/pretrained_models/Rescaling_DF2K_X4_HCFlow.pth
+
+
+
diff --git a/codes/options/test/test_SR_CelebA_8X_HCFlow.yml b/codes/options/test/test_SR_CelebA_8X_HCFlow.yml
new file mode 100644
index 0000000..2daf537
--- /dev/null
+++ b/codes/options/test/test_SR_CelebA_8X_HCFlow.yml
@@ -0,0 +1,80 @@
+#### general settings
+name: 002_HCFlow_CelebA_x8_bicSR_test
+suffix: ~
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 8
+quant: 256
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ test0:
+ name: example
+ mode: GTLQ
+ dataroot_GT: ../datasets/example_face_4X/HR
+ dataroot_LQ: ../datasets/example_face_4X/LR
+
+# val:
+# name: SR_CelebA_8X_160_val
+# mode: LRHR_PKL
+# dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va.pklv4
+# dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va_X8.pklv4
+# n_max: 20
+#
+# test:
+# name: SR_CelebA_8X_160_test
+# mode: LRHR_PKL
+# dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_te.pklv4
+# dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_te_X8.pklv4
+# n_max: 5000
+
+
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 3
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ stage1: True
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [5, 5]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+
+
+#### validation settings
+val:
+ heats: [0, 0.8]
+ n_sample: 1
+
+
+path:
+ strict_load: true
+ load_submodule: ~
+# pretrain_model_G: ../experiments/pretrained_models/SR_CelebA_X8_HCFlow.pth
+# pretrain_model_G: ../experiments/pretrained_models/SR_CelebA_X8_HCFlow+.pth
+ pretrain_model_G: ../experiments/pretrained_models/SR_CelebA_X8_HCFlow++.pth
+
+
diff --git a/codes/options/test/test_SR_DF2K_4X_HCFlow.yml b/codes/options/test/test_SR_DF2K_4X_HCFlow.yml
new file mode 100644
index 0000000..c67e5d2
--- /dev/null
+++ b/codes/options/test/test_SR_DF2K_4X_HCFlow.yml
@@ -0,0 +1,93 @@
+#### general settings
+name: 001_HCFlow_DF2K_x4_bicSR_test
+suffix: ~
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 4
+quant: 64
+gpu_ids: [0]
+
+
+
+datasets:
+ test0:
+ name: example
+ mode: GTLQ
+ dataroot_GT: ../datasets/example_general_4X/HR
+ dataroot_LQ: ../datasets/example_general_4X/LR
+
+# test_1:
+# name: Set5
+# mode: GTLQx
+# dataroot_GT: ../datasets/Set5/HR
+# dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+
+# test_2:
+# name: Set14
+# mode: GTLQ
+# dataroot_GT: ../datasets/Set14/HR
+# dataroot_LQ: ../datasets/Set14/LR_bicubic/X4
+#
+# test_3:
+# name: BSD100
+# mode: GTLQ
+# dataroot_GT: ../datasets/BSD100/HR
+# dataroot_LQ: ../datasets/BSD100/LR_bicubic/X4
+#
+# test_4:
+# name: Urban100
+# mode: GTLQ
+# dataroot_GT: ../datasets/Urban100/HR
+# dataroot_LQ: ../datasets/Urban100/LR_bicubic/X4
+#
+# test_5:
+# name: DIV2K-va-4X
+# mode: GTLQ
+# dataroot_GT: ../datasets/srflow_datasets/div2k-validation-modcrop8-gt
+# dataroot_LQ: ../datasets/srflow_datasets/div2k-validation-modcrop8-x4
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 2
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [7, 7]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+
+#### validation settings
+val:
+ heats: [0,0, 0.9]
+ n_sample: 1
+
+
+path:
+ strict_load: true
+ load_submodule: ~
+# pretrain_model_G: ../experiments/pretrained_models/SR_DF2K_X4_HCFlow.pth
+# pretrain_model_G: ../experiments/pretrained_models/SR_DF2K_X4_HCFlow+.pth
+ pretrain_model_G: ../experiments/pretrained_models/SR_DF2K_X4_HCFlow++.pth
+
+
+
diff --git a/codes/options/train/train_Rescaling_DF2K_4X_HCFlow.yml b/codes/options/train/train_Rescaling_DF2K_4X_HCFlow.yml
new file mode 100644
index 0000000..c1bbb92
--- /dev/null
+++ b/codes/options/train/train_Rescaling_DF2K_4X_HCFlow.yml
@@ -0,0 +1,127 @@
+#### general settings
+name: 003_DF2K_x4_rescaling_HCFlow
+use_tb_logger: true
+model: HCFlow_Rescaling
+distortion: sr
+scale: 4
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: DF2K_tr
+ mode: GTLQnpy
+ dataroot_GT: /cluster/work/cvl/jinliang_dataset/DIV2K+Flickr2K_decoded/DIV2K+Flickr2K_HR
+ dataroot_LQ: /cluster/work/cvl/jinliang_dataset/DIV2K+Flickr2K_decoded/DIV2K+Flickr2K_LR_bicubic/X4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ use_rot: true
+ color: RGB
+
+ val:
+ name: Set5
+ mode: GTLQx
+ dataroot_GT: ../datasets/Set5/HR
+ dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+
+
+# The optimization may not be stable for rescaling (+-0.1dB). A simple trick: for each stage of learning rate,
+# resume training from the best model of the previous stage of learning rate.
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_Rescaling
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 14
+ L: 2
+ squeeze: haar # better than squeeze2d
+ flow_permutation: none # bettter than invconv
+ flow_coupling: Affine3shift # better than affine
+ nn_module: DenseBlock # better than FCN
+ hidden_channels: 32
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [6, 6]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ stage1: True
+ feature_extractor: RRDB
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [2,1]
+ RRDB_nf: 64
+ RRDB_gc: 16
+
+
+#### path
+path:
+ pretrain_model_G: ~
+ strict_load: true
+ resume_state: auto
+
+#### training settings: learning rate scheme, loss
+train:
+ two_stage_opt: True
+
+ lr_G: !!float 2.5e-4
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 500000
+ warmup_iter: -1 # no warm up
+ lr_steps: [100000, 200000, 300000, 400000, 450000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ weight_z: !!float 1e-5
+
+ pixel_criterion_lr: l2
+ pixel_weight_lr: !!float 5e-2
+
+ eps_std_reverse: 1.0
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 1.0
+
+ # perceptual loss
+ feature_criterion: l1
+ feature_weight: 0
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: 0
+
+ lr_D: 0
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 1.0]
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_CelebA_8X_HCFlow++.yml b/codes/options/train/train_SR_CelebA_8X_HCFlow++.yml
new file mode 100644
index 0000000..8a48a96
--- /dev/null
+++ b/codes/options/train/train_SR_CelebA_8X_HCFlow++.yml
@@ -0,0 +1,126 @@
+#### general settings
+name: 002_CelebA_x8_bicSR_HCFlow++
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 8
+quant: 256
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: CelebA_160_tr
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr_X8.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: CelebA_160_va
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va_X8.pklv4
+ n_max: 20
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 3
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64 # 64 and 128 are similar
+ cond_channels: ~ # affine coupling in the main trunk, testo 3 or None
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ stage1: True
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [5, 5]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ../experiments/pretrained_models/SR_CelebA_X8_HCFlow.pth
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 1.25e-5
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 50000
+ warmup_iter: -1 # no warm up
+ lr_steps: [20000, 40000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: !!float 2e-3
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 1.0
+
+ # perceptual loss
+ eps_std_reverse: 0.8
+ feature_criterion: l1
+ feature_weight: !!float 5e-2 # balance diversity and lpips
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: !!float 5e-1
+
+ lr_D: !!float 5e-5
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.8] # 0.8 has best visual quality for face SR
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_CelebA_8X_HCFlow+.yml b/codes/options/train/train_SR_CelebA_8X_HCFlow+.yml
new file mode 100644
index 0000000..2953518
--- /dev/null
+++ b/codes/options/train/train_SR_CelebA_8X_HCFlow+.yml
@@ -0,0 +1,126 @@
+#### general settings
+name: 002_CelebA_x8_bicSR_HCFlow+
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 8
+quant: 256
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: CelebA_160_tr
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr_X8.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: CelebA_160_va
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va_X8.pklv4
+ n_max: 20
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 3
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64 # 64 and 128 are similar
+ cond_channels: ~ # affine coupling in the main trunk, testo 3 or None
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ stage1: True
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [5, 5]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ../experiments/pretrained_models/SR_CelebA_X8_HCFlow.pth
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 5e-5
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 50000
+ warmup_iter: -1 # no warm up
+ lr_steps: [20000, 40000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: !!float 2e-3
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 1.0
+
+ # perceptual loss
+ eps_std_reverse: 0.8
+ feature_criterion: l1
+ feature_weight: 0
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: 0
+
+ lr_D: 0
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.8] # 0.8 has best visual quality for face SR
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_CelebA_8X_HCFlow.yml b/codes/options/train/train_SR_CelebA_8X_HCFlow.yml
new file mode 100644
index 0000000..db4ca58
--- /dev/null
+++ b/codes/options/train/train_SR_CelebA_8X_HCFlow.yml
@@ -0,0 +1,125 @@
+#### general settings
+name: 002_CelebA_x8_bicSR_HCFlow
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 8
+quant: 256
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: CelebA_160_tr
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_tr_X8.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: CelebA_160_va
+ mode: LRHR_PKL
+ dataroot_GT: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va.pklv4
+ dataroot_LQ: /cluster/work/cvl/sr_datasets/srflow/celebA/CelebAHq_160_MBic_va_X8.pklv4
+ n_max: 20
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 3
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [5, 5]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ~
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 2.5e-4
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 350000
+ warmup_iter: -1 # no warm up
+ lr_steps: [200000, 250000, 280000, 310000, 340000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: 1
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 0
+
+ # perceptual loss
+ eps_std_reverse: 0.8
+ feature_criterion: l1
+ feature_weight: 0
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: 0
+
+ lr_D: 0
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.8]
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_DF2K_4X_HCFlow++.yml b/codes/options/train/train_SR_DF2K_4X_HCFlow++.yml
new file mode 100644
index 0000000..32b6f80
--- /dev/null
+++ b/codes/options/train/train_SR_DF2K_4X_HCFlow++.yml
@@ -0,0 +1,124 @@
+#### general settings
+name: 001_DF2K_x4_bicSR_HCFlow++
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 4
+quant: 64
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: DF2K_tr
+ mode: LRHR_PKL
+ dataroot_GT: ../datasets/srflow_datasets/DF2K-tr.pklv4
+ dataroot_LQ: ../datasets/srflow_datasets/DF2K-tr_X4.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: Set5
+ mode: GTLQx
+ dataroot_GT: ../datasets/Set5/HR
+ dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 2
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [7, 7]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ../experiments/pretrained_models/SR_DF2K_X4_HCFlow.pth
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 1.25e-5
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 50000
+ warmup_iter: -1 # no warm up
+ lr_steps: [20000, 40000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: !!float 2e-3
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 1.0
+
+ # perceptual loss
+ eps_std_reverse: 0.9
+ feature_criterion: l1
+ feature_weight: !!float 5e-2 # balance diversity and lpips
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: !!float 5e-1
+
+ lr_D: !!float 5e-5
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.9] # 0.9 has best visual quality for general SR
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_DF2K_4X_HCFlow+.yml b/codes/options/train/train_SR_DF2K_4X_HCFlow+.yml
new file mode 100644
index 0000000..4e918b6
--- /dev/null
+++ b/codes/options/train/train_SR_DF2K_4X_HCFlow+.yml
@@ -0,0 +1,124 @@
+#### general settings
+name: 001_DF2K_x4_bicSR_HCFlow+
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 4
+quant: 64
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: DF2K_tr
+ mode: LRHR_PKL
+ dataroot_GT: ../datasets/srflow_datasets/DF2K-tr.pklv4
+ dataroot_LQ: ../datasets/srflow_datasets/DF2K-tr_X4.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: Set5
+ mode: GTLQx
+ dataroot_GT: ../datasets/Set5/HR
+ dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 2
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [7, 7]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ../experiments/pretrained_models/SR_DF2K_X4_HCFlow.pth
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 5e-5
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 50000
+ warmup_iter: -1 # no warm up
+ lr_steps: [20000, 40000]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: !!float 2e-3
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 1.0
+
+ # perceptual loss
+ eps_std_reverse: 0.9
+ feature_criterion: l1
+ feature_weight: 0
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: 0
+
+ lr_D: 0
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.9] # 0.9 has best visual quality for general SR
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/options/train/train_SR_DF2K_4X_HCFlow.yml b/codes/options/train/train_SR_DF2K_4X_HCFlow.yml
new file mode 100644
index 0000000..c202e93
--- /dev/null
+++ b/codes/options/train/train_SR_DF2K_4X_HCFlow.yml
@@ -0,0 +1,124 @@
+#### general settings
+name: 001_DF2K_x4_bicSR_HCFlow
+use_tb_logger: true
+model: HCFlow_SR
+distortion: sr
+scale: 4
+quant: 64
+gpu_ids: [0]
+
+
+#### datasets
+datasets:
+ train:
+ name: DF2K_tr
+ mode: LRHR_PKL
+ dataroot_GT: ../datasets/srflow_datasets/DF2K-tr.pklv4
+ dataroot_LQ: ../datasets/srflow_datasets/DF2K-tr_X4.pklv4
+
+ use_shuffle: true
+ n_workers: 16
+ batch_size: 16
+ GT_size: 160
+ use_flip: true
+ color: RGB
+ val:
+ name: Set5
+ mode: GTLQx
+ dataroot_GT: ../datasets/Set5/HR
+ dataroot_LQ: ../datasets/Set5/LR_bicubic/X4
+
+
+#### network structures
+network_G:
+ which_model_G: HCFlowNet_SR
+ in_nc: 3
+ out_nc: 3
+ act_norm_start_step: 100
+
+ flowDownsampler:
+ K: 26
+ L: 2
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ hidden_channels: 64
+ cond_channels: ~
+ splitOff:
+ enable: true
+ after_flowstep: [13, 13]
+ flow_permutation: invconv
+ flow_coupling: Affine
+ nn_module: FCN
+ nn_module_last: Conv2dZeros
+ hidden_channels: 64
+ RRDB_nb: [7, 7]
+ RRDB_nf: 64
+ RRDB_gc: 32
+
+network_D:
+ which_model_D: discriminator_vgg_160
+ in_nc: 3
+ nf: 64
+
+
+#### path
+path:
+ pretrain_model_G: ~
+ strict_load: true
+ resume_state: auto
+
+
+#### training settings: learning rate scheme, loss
+train:
+ lr_G: !!float 2.5e-4
+ lr_scheme: MultiStepLR
+ weight_decay_G: 0
+ max_grad_clip: 5
+ max_grad_norm: 100
+ beta1: 0.9
+ beta2: 0.99
+ niter: 300000
+ warmup_iter: -1 # no warm up
+ lr_steps_rel: [0.5,0.75,0.9,0.95]
+ lr_gamma: 0.5
+ restarts: ~
+ restart_weights: ~
+ eta_min: !!float 1e-8
+
+ nll_weight: 1
+
+ # pixel loss
+ pixel_criterion_hr: l1
+ pixel_weight_hr: 0
+
+ # perceptual loss
+ eps_std_reverse: 0.9
+ feature_criterion: l1
+ feature_weight: 0
+
+ # gan loss
+ gan_type: gan # gan | lsgan | wgangp | ragan (patchgan uses lsgan)
+ gan_weight: 0
+
+ lr_D: 0
+ beta1_D: 0.9
+ beta2_D: 0.99
+ D_update_ratio: 1
+ D_init_iters: 1500
+
+ manual_seed: 0
+ val_freq: !!float 5e3
+
+
+#### validation settings
+val:
+ heats: [0.0, 0.9]
+ n_sample: 3
+
+
+#### logger
+logger:
+ print_freq: 200
+ save_checkpoint_freq: !!float 5e3
+
diff --git a/codes/scripts/png2npy.py b/codes/scripts/png2npy.py
new file mode 100644
index 0000000..93b6024
--- /dev/null
+++ b/codes/scripts/png2npy.py
@@ -0,0 +1,42 @@
+import os
+import argparse
+import skimage.io as sio
+import numpy as np
+
+# usage: python scripts/png2npy.py --pathFrom ../datasets/DIV2K+Flickr2K/DIV2K+Flickr2K_HR/ --pathTo ../datasets/DIV2K+Flickr2K_decoded/DIV2K+Flickr2K_HR/
+# python scripts/png2npy.py --pathFrom ../datasets/DIV2K+Flickr2K/DIV2K+Flickr2K_LR_bicubic/ --pathTo ../datasets/DIV2K+Flickr2K_decoded/DIV2K+Flickr2K_LR_bicubic/
+
+
+parser = argparse.ArgumentParser(description='Pre-processing .png images')
+parser.add_argument('--pathFrom', default='',
+ help='directory of images to convert')
+parser.add_argument('--pathTo', default='',
+ help='directory of images to save')
+parser.add_argument('--split', default=True,
+ help='save individual images')
+parser.add_argument('--select', default='',
+ help='select certain path')
+
+args = parser.parse_args()
+
+for (path, dirs, files) in os.walk(args.pathFrom):
+ print(path)
+ targetDir = os.path.join(args.pathTo, path[len(args.pathFrom) + 1:])
+ if len(args.select) > 0 and path.find(args.select) == -1:
+ continue
+
+ if not os.path.exists(targetDir):
+ os.mkdir(targetDir)
+
+ if len(dirs) == 0:
+ pack = {}
+ n = 0
+ for fileName in files:
+ (idx, ext) = os.path.splitext(fileName)
+ if ext == '.png':
+ image = sio.imread(os.path.join(path, fileName))
+ if args.split:
+ np.save(os.path.join(targetDir, idx + '.npy'), image)
+ n += 1
+ if n % 100 == 0:
+ print('Converted ' + str(n) + ' images.')
diff --git a/codes/scripts/prepare_data_pkl.py b/codes/scripts/prepare_data_pkl.py
new file mode 100644
index 0000000..b056452
--- /dev/null
+++ b/codes/scripts/prepare_data_pkl.py
@@ -0,0 +1,119 @@
+# Copyright (c) 2020 Huawei Technologies Co., Ltd.
+# Licensed under CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International) (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+#
+# The code is released for academic research use only. For commercial use, please contact Huawei Technologies Co., Ltd.
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import os
+import sys
+
+import numpy as np
+import random
+import imageio
+
+from natsort import natsort
+from tqdm import tqdm
+
+def get_img_paths(dir_path, wildcard='*.png'):
+ return natsort.natsorted(glob.glob(dir_path + '/' + wildcard))
+
+def create_all_dirs(path):
+ if "." in path.split("/")[-1]:
+ dirs = os.path.dirname(path)
+ else:
+ dirs = path
+ os.makedirs(dirs, exist_ok=True)
+
+def to_pklv4(obj, path, vebose=False):
+ create_all_dirs(path)
+ with open(path, 'wb') as f:
+ pickle.dump(obj, f, protocol=4)
+ if vebose:
+ print("Wrote {}".format(path))
+
+
+from imresize import imresize
+
+def random_crop(img, size):
+ h, w, c = img.shape
+
+ h_start = np.random.randint(0, h - size)
+ h_end = h_start + size
+
+ w_start = np.random.randint(0, w - size)
+ w_end = w_start + size
+
+ return img[h_start:h_end, w_start:w_end]
+
+
+def imread(img_path):
+ img = imageio.imread(img_path)
+ if len(img.shape) == 2:
+ img = np.stack([img, ] * 3, axis=2)
+ return img
+
+
+def to_pklv4_1pct(obj, path, vebose):
+ n = int(round(len(obj) * 0.01))
+ path = path.replace(".", "_1pct.")
+ to_pklv4(obj[:n], path, vebose=True)
+
+
+def main(dir_path):
+ hrs = []
+ lqs = []
+
+ img_paths = get_img_paths(dir_path)
+ for img_path in tqdm(img_paths):
+ img = imread(img_path)
+
+ for i in range(47):
+ crop = random_crop(img, 160)
+ cropX4 = imresize(crop, scalar_scale=0.25)
+ hrs.append(crop)
+ lqs.append(cropX4)
+
+ shuffle_combined(hrs, lqs)
+
+ hrs_path = get_hrs_path(dir_path)
+ to_pklv4(hrs, hrs_path, vebose=True)
+ to_pklv4_1pct(hrs, hrs_path, vebose=True)
+
+ lqs_path = get_lqs_path(dir_path)
+ to_pklv4(lqs, lqs_path, vebose=True)
+ to_pklv4_1pct(lqs, lqs_path, vebose=True)
+
+
+def get_hrs_path(dir_path):
+ base_dir = os.path.dirname(dir_path)
+ name = os.path.basename(dir_path)
+ hrs_path = os.path.join(base_dir, 'pkls', name + '.pklv4')
+ return hrs_path
+
+
+def get_lqs_path(dir_path):
+ base_dir = os.path.dirname(dir_path)
+ name = os.path.basename(dir_path)
+ hrs_path = os.path.join(base_dir, 'pkls', name + '_X4.pklv4')
+ return hrs_path
+
+
+def shuffle_combined(hrs, lqs):
+ combined = list(zip(hrs, lqs))
+ random.shuffle(combined)
+ hrs[:], lqs[:] = zip(*combined)
+
+
+if __name__ == "__main__":
+ dir_path = sys.argv[1]
+ assert os.path.isdir(dir_path)
+ main(dir_path)
diff --git a/codes/test_HCFlow.py b/codes/test_HCFlow.py
new file mode 100644
index 0000000..69b4696
--- /dev/null
+++ b/codes/test_HCFlow.py
@@ -0,0 +1,252 @@
+import os.path
+import logging
+import time
+import argparse
+from collections import OrderedDict
+import numpy as np
+import torch
+import options.options as option
+import utils.util as util
+from utils.imresize import imresize
+from data.util import bgr2ycbcr
+from data import create_dataset, create_dataloader
+from models import create_model
+import lpips
+
+os.environ['CUDA_HOME'] = '/scratch_net/rind/cuda-11.0'
+os.environ['LD_LIBRARY_PATH'] = '/scratch_net/rind/cuda-11.0/lib64'
+os.environ['PATH'] = '/scratch_net/rind/cuda-11.0/bin'
+
+#### options
+parser = argparse.ArgumentParser() # test_SR_CelebA_8X_HCFlow test_SR_DF2K_4X_HCFlow test_Rescaling_DF2K_4X_HCFlow
+parser.add_argument('--opt', type=str, default='options/test/test_SR_CelebA_8X_HCFlow.yml', help='Path to options YMAL file.')
+parser.add_argument('--save_kernel', action='store_true', default=False, help='Save Kernel Esimtation.')
+args = parser.parse_args()
+opt = option.parse(args.opt, is_train=False)
+opt = option.dict_to_nonedict(opt)
+device_id = torch.cuda.current_device()
+
+#### mkdir and logger
+util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
+ and 'pretrain_model' not in key and 'resume' not in key and 'load_submodule' not in key))
+util.setup_logger('base', opt['path']['log'], 'test_' + opt['name'], level=logging.INFO,
+ screen=True, tofile=True)
+logger = logging.getLogger('base')
+logger.info(option.dict2str(opt))
+
+# set random seed
+util.set_random_seed(0)
+
+#### Create test dataset and dataloader
+test_loaders = []
+for phase, dataset_opt in sorted(opt['datasets'].items()):
+ test_set = create_dataset(dataset_opt)
+ test_loader = create_dataloader(test_set, dataset_opt)
+ logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))
+ test_loaders.append(test_loader)
+
+# load pretrained model by default
+model = create_model(opt)
+loss_fn_alex = lpips.LPIPS(net='alex').to('cuda')
+crop_border = opt['crop_border'] if opt['crop_border'] else opt['scale']
+
+for test_loader in test_loaders:
+ test_set_name = test_loader.dataset.opt['name']
+ logger.info('\n\nTesting [{:s}]...'.format(test_set_name))
+ test_start_time = time.time()
+ dataset_dir = os.path.join(opt['path']['results_root'], test_set_name)
+ util.mkdir(dataset_dir)
+
+ idx = 0
+ psnr_dict={} # for HR image
+ ssim_dict={}
+ psnr_y_dict = {}
+ ssim_y_dict = {}
+ bic_hr_psnr_dict={} # for bic(HR)
+ bic_hr_ssim_dict={}
+ bic_hr_psnr_y_dict = {}
+ bic_hr_ssim_y_dict = {}
+ lpips_dict = {}
+ diversity_dict = {} # pixel-wise variance
+ avg_lr_psnr = 0.0 # for generated LR image
+ avg_lr_ssim = 0.0
+ avg_lr_psnr_y = 0.0
+ avg_lr_ssim_y = 0.0
+ avg_nll = 0.0
+
+ for test_data in test_loader:
+ idx += 1
+
+ real_image = True if test_loader.dataset.opt['dataroot_GT'] is None else False
+ generate_online = True if test_loader.dataset.opt['dataroot_GT'] is not None and test_loader.dataset.opt[
+ 'dataroot_LQ'] is None else False
+ img_path = test_data['LQ_path'][0] if real_image else test_data['GT_path'][0]
+ img_name = os.path.splitext(os.path.basename(img_path))[0]
+
+ model.feed_data(test_data)
+ nll = model.test()
+ avg_nll += nll
+ visuals = model.get_current_visuals()
+
+ # calculate PSNR for LR
+ gt_img_lr = util.tensor2img(visuals['LQ'])
+ sr_img_lr = util.tensor2img(visuals['LQ_fromH'])
+ # save_img_path = os.path.join(dataset_dir, 'LR_{:s}_{:.1f}_{:d}.png'.format(img_name, 1.0, 0))
+ # util.save_img(sr_img_lr, save_img_path)
+ gt_img_lr = gt_img_lr / 255.
+ sr_img_lr = sr_img_lr / 255.
+
+ lr_psnr, lr_ssim, lr_psnr_y, lr_ssim_y = util.calculate_psnr_ssim(gt_img_lr, sr_img_lr, 0)
+ avg_lr_psnr += lr_psnr
+ avg_lr_ssim += lr_ssim
+ avg_lr_psnr_y += lr_psnr_y
+ avg_lr_ssim_y += lr_ssim_y
+
+ # deal with real-world data (just save)
+ if real_image:
+ for heat in opt['val']['heats']:
+ for sample in range(opt['val']['n_sample']):
+ sr_img = util.tensor2img(visuals['SR', heat, sample])
+
+ # deal with the image margins for real images
+ if opt['scale'] == 4:
+ real_crop = 3
+ elif opt['scale'] == 2:
+ real_crop = 6
+ elif opt['scale'] == 1:
+ real_crop = 11
+ assert real_crop * opt['scale'] * 2 > opt['kernel_size']
+ sr_img = sr_img[real_crop * opt['scale']:-real_crop * opt['scale'],
+ real_crop * opt['scale']:-real_crop * opt['scale'], :]
+ save_img_path = os.path.join(dataset_dir, 'SR_{:s}_{:.1f}_{:d}.png'.
+ format(img_name, heat, sample))
+ util.save_img(sr_img, save_img_path)
+
+ # deal with synthetic data (calculate psnr and save)
+ else:
+ for heat in opt['val']['heats']:
+ psnr = 0.0
+ ssim = 0.0
+ psnr_y = 0.0
+ ssim_y = 0.0
+ lpips_value = 0.0
+ bic_hr_psnr = 0.0
+ bic_hr_ssim = 0.0
+ bic_hr_psnr_y = 0.0
+ bic_hr_ssim_y = 0.0
+
+ sr_img_list =[]
+ for sample in range(opt['val']['n_sample']):
+ gt_img = visuals['GT']
+ sr_img = visuals['SR', heat, sample]
+ sr_img_list.append(sr_img.unsqueeze(0)*255)
+ lpips_dict[(idx, heat, sample)] = float(loss_fn_alex(2 * gt_img.to('cuda') - 1, 2 * sr_img.to('cuda') - 1).cpu())
+ lpips_value += lpips_dict[(idx, heat, sample)]
+
+ gt_img = util.tensor2img(gt_img) # uint8
+ sr_img = util.tensor2img(sr_img) # uint8
+ suffix = opt['suffix']
+ if suffix:
+ save_img_path = os.path.join(dataset_dir, 'SR_{:s}_{:.1f}_{:d}_{:s}.png'.format(img_name, heat, sample, suffix))
+ else:
+ save_img_path = os.path.join(dataset_dir, 'SR_{:s}_{:.1f}_{:d}.png'.format(img_name, heat, sample))
+ util.save_img(sr_img, save_img_path)
+
+ gt_img = gt_img / 255.
+ sr_img = sr_img / 255.
+ bic_hr_gt_img = imresize(gt_img, 1 / opt['scale'])
+ bic_hr_sr_img = imresize(sr_img, 1 / opt['scale'])
+
+ psnr_dict[(idx, heat, sample)], ssim_dict[(idx, heat, sample)], \
+ psnr_y_dict[(idx, heat, sample)], ssim_y_dict[(idx, heat, sample)] = util.calculate_psnr_ssim(gt_img, sr_img, crop_border)
+ psnr += psnr_dict[(idx, heat, sample)]
+ ssim += ssim_dict[(idx, heat, sample)]
+ psnr_y += psnr_y_dict[(idx, heat, sample)]
+ ssim_y += ssim_y_dict[(idx, heat, sample)]
+ bic_hr_psnr_dict[(idx, heat, sample)], bic_hr_ssim_dict[(idx, heat, sample)], \
+ bic_hr_psnr_y_dict[(idx, heat, sample)], bic_hr_ssim_y_dict[(idx, heat, sample)] = util.calculate_psnr_ssim(bic_hr_gt_img, bic_hr_sr_img, 0)
+ bic_hr_psnr += bic_hr_psnr_dict[(idx, heat, sample)]
+ bic_hr_ssim += bic_hr_ssim_dict[(idx, heat, sample)]
+ bic_hr_psnr_y += bic_hr_psnr_y_dict[(idx, heat, sample)]
+ bic_hr_ssim_y += bic_hr_ssim_y_dict[(idx, heat, sample)]
+
+
+ # mean pixel-wise variance
+ psnr /= opt['val']['n_sample']
+ ssim /= opt['val']['n_sample']
+ psnr_y /= opt['val']['n_sample']
+ ssim_y /= opt['val']['n_sample']
+ diversity_dict[(idx, heat)] = float(torch.cat(sr_img_list, 0).std([0]).mean().cpu())
+ lpips_value /= opt['val']['n_sample']
+ bic_hr_psnr /= opt['val']['n_sample']
+ bic_hr_ssim /= opt['val']['n_sample']
+ bic_hr_psnr_y /= opt['val']['n_sample']
+ bic_hr_ssim_y /= opt['val']['n_sample']
+
+
+ logger.info('{:20s} ({}samples),heat:{:.1f}) '
+ 'HR:PSNR/SSIM/PSNR_Y/SSIM_Y/LPIPS/Diversity: {:.2f}/{:.4f}/{:.2f}/{:.4f}/{:.4f}/{:.4f}, '
+ 'bicHR:PSNR/SSIM/PSNR_Y/SSIM_Y: {:.2f}/{:.4f}/{:.2f}/{:.4f}, '
+ 'LR:PSNR/SSIM/PSNR_Y/SSIM_Y: {:.2f}/{:.4f}/{:.2f}/{:.4f}, NLL: {:.4f}'.format(
+ img_name, opt['val']['n_sample'], heat,
+ psnr, ssim, psnr_y, ssim_y, lpips_value, diversity_dict[(idx, heat)],
+ bic_hr_psnr, bic_hr_ssim, bic_hr_psnr_y, bic_hr_ssim_y,
+ lr_psnr, lr_ssim, lr_psnr_y, lr_ssim_y, nll))
+
+ # Average PSNR/SSIM results
+ avg_lr_psnr /= idx
+ avg_lr_ssim /= idx
+ avg_lr_psnr_y /= idx
+ avg_lr_ssim_y /= idx
+ avg_nll = avg_nll / idx
+
+ if real_image:
+ logger.info('----{} ({} images), avg LR PSNR/SSIM/PSNR_K/LR_SSIM_Y: {:.2f}/{:.4f}/{:.2f}/{:.4f}\n'.format(test_set_name, idx, avg_lr_psnr, avg_lr_ssim, avg_lr_psnr_y, avg_lr_ssim_y))
+ else:
+ logger.info('-------------------------------------------------------------------------------------')
+ for heat in opt['val']['heats']:
+ avg_psnr = 0.0
+ avg_ssim = 0.0
+ avg_psnr_y = 0.0
+ avg_ssim_y = 0.0
+ avg_lpips = 0.0
+ avg_diversity = 0.0
+ avg_bic_hr_psnr = 0.0
+ avg_bic_hr_ssim = 0.0
+ avg_bic_hr_psnr_y = 0.0
+ avg_bic_hr_ssim_y = 0.0
+
+ for iidx in range(1, idx+1):
+ for sample in range(opt['val']['n_sample']):
+ avg_psnr += psnr_dict[(iidx, heat, sample)]
+ avg_ssim += ssim_dict[(iidx, heat, sample)]
+ avg_psnr_y += psnr_y_dict[(iidx, heat, sample)]
+ avg_ssim_y += ssim_y_dict[(iidx, heat, sample)]
+ avg_lpips += lpips_dict[(iidx, heat, sample)]
+ avg_bic_hr_psnr += bic_hr_psnr_dict[(iidx, heat, sample)]
+ avg_bic_hr_ssim += bic_hr_ssim_dict[(iidx, heat, sample)]
+ avg_bic_hr_psnr_y += bic_hr_psnr_y_dict[(iidx, heat, sample)]
+ avg_bic_hr_ssim_y += bic_hr_ssim_y_dict[(iidx, heat, sample)]
+ avg_diversity += diversity_dict[(iidx, heat)]
+
+ avg_psnr = avg_psnr / idx / opt['val']['n_sample']
+ avg_ssim = avg_ssim / idx / opt['val']['n_sample']
+ avg_psnr_y = avg_psnr_y / idx / opt['val']['n_sample']
+ avg_ssim_y = avg_ssim_y / idx / opt['val']['n_sample']
+ avg_lpips = avg_lpips / idx / opt['val']['n_sample']
+ avg_diversity = avg_diversity / idx
+ avg_bic_hr_psnr = avg_bic_hr_psnr / idx / opt['val']['n_sample']
+ avg_bic_hr_ssim = avg_bic_hr_ssim / idx / opt['val']['n_sample']
+ avg_bic_hr_psnr_y = avg_bic_hr_psnr_y / idx / opt['val']['n_sample']
+ avg_bic_hr_ssim_y = avg_bic_hr_ssim_y / idx / opt['val']['n_sample']
+
+ # log
+ logger.info(opt['path']['pretrain_model_G'])
+ logger.info('----{} ({}images,{}samples,heat:{:.1f}) '
+ 'average HR:PSNR/SSIM/PSNR_Y/SSIM_Y/LPIPS/Diversity: {:.2f}/{:.4f}/{:.2f}/{:.4f}/{:.4f}/{:.4f}, '
+ 'bicHR:PSNR/SSIM/PSNR_Y/SSIM_Y: {:.2f}/{:.4f}/{:.2f}/{:.4f}, '
+ 'LR:PSNR/SSIM/PSNR_Y/SSIM_Y: {:.2f}/{:.4f}/{:.2f}/{:.4f}, NLL: {:.4f}'.format(
+ test_set_name, idx, opt['val']['n_sample'], heat,
+ avg_psnr, avg_ssim, avg_psnr_y, avg_ssim_y, avg_lpips, avg_diversity,
+ avg_bic_hr_psnr, avg_bic_hr_ssim, avg_bic_hr_psnr_y, avg_bic_hr_ssim_y,
+ avg_lr_psnr, avg_lr_ssim, avg_lr_psnr_y, avg_lr_ssim_y, avg_nll))
diff --git a/codes/train_HCFlow.py b/codes/train_HCFlow.py
new file mode 100644
index 0000000..b64c9db
--- /dev/null
+++ b/codes/train_HCFlow.py
@@ -0,0 +1,314 @@
+import os
+import math
+import argparse
+import random
+import logging
+import numpy as np
+import torch
+from data.data_sampler import DistIterSampler, EnlargedSampler
+from data.util import bgr2ycbcr
+
+import options.options as option
+from utils import util
+from data import create_dataloader, create_dataset
+from models import create_model
+from utils.dist_util import get_dist_info, init_dist
+
+import socket
+import getpass
+import lpips
+
+
+
+def main():
+ #### setup options
+ parser = argparse.ArgumentParser() # train_SR_CelebA_8X_HCFlow train_SR_DF2K_4X_HCFlow train_Rescaling_DF2K_4X_HCFlow
+ parser.add_argument('--opt', type=str, default='options/train/train_SR_CelebA_8X_HCFlow.yml',
+ help='Path to option YMAL file of MANet.')
+ parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
+ help='job launcher')
+ parser.add_argument('--local_rank', type=int, default=0)
+ parser.add_argument('--gpu_ids', type=str, default=None)
+ parser.add_argument('--job_id', type=str, default=0)
+ parser.add_argument('--job_path', type=str, default='')
+ args = parser.parse_args()
+ opt = option.parse(args.opt, args.gpu_ids, is_train=True)
+ device_id = torch.cuda.current_device()
+
+ # convert to NoneDict, which returns None for missing keys
+ opt = option.dict_to_nonedict(opt)
+ print(torch.__version__)
+
+ #### random seed
+ seed = opt['train']['manual_seed']
+ if seed is None:
+ seed = random.randint(1, 10000)
+ util.set_random_seed(seed)
+
+ # distributed settings
+ if args.launcher == 'none':
+ opt['dist'] = False
+ print('Disable distributed.', flush=True)
+ else:
+ opt['dist'] = True
+ if args.launcher == 'slurm' and 'dist_params' in opt:
+ init_dist(args.launcher, **opt['dist_params'])
+ else:
+ init_dist(args.launcher)
+
+ opt['rank'], opt['world_size'] = get_dist_info()
+
+ torch.backends.cudnn.benchmark = True
+ # torch.backends.cudnn.deterministic = True
+
+ #### loading resume state if exists
+ if opt['path'].get('resume_state', None):
+ resume_state_path, _ = util.get_resume_paths(opt)
+ if resume_state_path is None:
+ resume_state = None
+ else:
+ # distributed resuming: all load into default GPU
+ device_id = torch.cuda.current_device()
+ resume_state = torch.load(resume_state_path,
+ map_location=lambda storage, loc: storage.cuda(device_id))
+ option.check_resume(opt, resume_state['iter']) # override model pretrain path with resume path
+ else:
+ resume_state = None
+
+ #### mkdir and loggers
+ # normal training (rank -1) OR distributed training (rank (gpu id) 0-7)
+ if opt['rank'] <= 0:
+ if resume_state is None:
+ util.mkdir_and_rename(
+ opt['path']['experiments_root']) # rename experiment folder if exists
+ util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
+ and 'pretrain_model' not in key and 'resume' not in key))
+
+ # config loggers. Before it, the log will not work
+ util.setup_logger('base', opt['path']['log'], 'train{}_'.format(args.job_id) + opt['name'], level=logging.INFO,
+ screen=True, tofile=True)
+ util.setup_logger('val', opt['path']['log'], 'val{}_'.format(args.job_id) + opt['name'], level=logging.INFO,
+ screen=True, tofile=True)
+ logger = logging.getLogger('base')
+ logger.info('{}@{}, GPU {}, Job_id {}, Job path {}'.format(getpass.getuser(), socket.gethostname(),
+ opt['gpu_ids'], args.job_id, args.job_path))
+ logger.info(option.dict2str(opt))
+ # tensorboard logger
+ if opt['use_tb_logger'] and 'debug' not in opt['name']:
+ version = float(torch.__version__[0:3])
+ if version >= 1.1: # PyTorch 1.1
+ from torch.utils.tensorboard import SummaryWriter
+ else:
+ logger.info(
+ 'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
+ from tensorboardX import SummaryWriter
+ tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
+ else:
+ util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
+ logger = logging.getLogger('base')
+
+ # symlink the code/working dir
+ try:
+ os.symlink(args.job_path.replace('/cluster/home/{}'.format(getpass.getuser()),'/scratch/e_home'),
+ opt['path']['experiments_root']+'/{}'.format(os.path.basename(args.job_path)))
+ except:
+ pass
+ try:
+ os.symlink(args.job_path.replace('/cluster/home/{}'.format(getpass.getuser()),'/scratch/e_home')+'/options/train',
+ opt['path']['experiments_root']+'/{}'.format(args.job_id))
+ except:
+ pass
+
+ #### create train and val dataloader
+ dataset_ratio = 200 # enlarge the size of each epoch
+ for phase, dataset_opt in opt['datasets'].items():
+ if phase == 'train':
+ train_set = create_dataset(dataset_opt)
+ train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
+ total_iters = int(opt['train']['niter'])
+ total_epochs = int(math.ceil(total_iters / train_size))
+ if opt['dist']:
+ train_sampler = DistIterSampler(train_set, opt['world_size'], opt['rank'], dataset_ratio)
+ total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
+ else:
+ train_sampler = None
+ train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
+ if opt['rank'] <= 0:
+ logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
+ len(train_set), train_size))
+ logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
+ total_epochs, total_iters))
+ elif phase == 'val':
+ val_set = create_dataset(dataset_opt)
+ val_loader = create_dataloader(val_set, dataset_opt, opt, None)
+ if opt['rank'] <= 0:
+ logger.info('Number of val images in [{:s}]: {:d}'.format(
+ dataset_opt['name'], len(val_set)))
+ else:
+ raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
+ assert train_loader is not None
+ assert val_loader is not None
+
+ #### create model
+ model = create_model(opt)
+
+ #### resume training
+ if resume_state:
+ logger.info('Resuming training from epoch: {}, iter: {}.'.format(
+ resume_state['epoch'], resume_state['iter']))
+
+ start_epoch = resume_state['epoch']
+ current_step = resume_state['iter']
+ model.resume_training(resume_state) # handle optimizers and schedulers
+ else:
+ current_step = 0
+ start_epoch = 0
+
+ #### training
+
+ # logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
+ for epoch in range(start_epoch, total_epochs + 1):
+ if opt['dist']:
+ train_sampler.set_epoch(epoch)
+
+ for _, train_data in enumerate(train_loader):
+ current_step += 1
+ if current_step > total_iters:
+ break
+
+ #### update learning rate, schedulers
+ model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])
+
+ #### training
+ model.feed_data(train_data)
+ model.optimize_parameters(current_step)
+
+ #### log
+ if current_step % opt['logger']['print_freq'] == 0:
+ logs = model.get_current_log()
+ message = '