diff --git a/codes/data/GTLQ_dataset.py b/codes/data/Downsample_dataset.py
similarity index 56%
rename from codes/data/GTLQ_dataset.py
rename to codes/data/Downsample_dataset.py
index c9a47a4c..efe07242 100644
--- a/codes/data/GTLQ_dataset.py
+++ b/codes/data/Downsample_dataset.py
@@ -7,14 +7,14 @@ import torch.utils.data as data
import data.util as util
-class GTLQDataset(data.Dataset):
+class DownsampleDataset(data.Dataset):
"""
- Reads unpaired high-resolution and low resolution images. Downsampled, LR images matching the provided high res
- images are produced and fed to the downstream model, which can be used in a pixel loss.
+ Reads an unpaired HQ and LQ image. Clips both images to the expected input sizes of the model. Produces a
+ downsampled LQ image from the HQ image and feeds that as well.
"""
def __init__(self, opt):
- super(GTLQDataset, self).__init__()
+ super(DownsampleDataset, self).__init__()
self.opt = opt
self.data_type = self.opt['data_type']
self.paths_LQ, self.paths_GT = None, None
@@ -23,8 +23,11 @@ class GTLQDataset(data.Dataset):
self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT'])
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
+
+ self.data_sz_mismatch_ok = opt['mismatched_Data_OK']
assert self.paths_GT, 'Error: GT path is empty.'
- if self.paths_LQ and self.paths_GT:
+ assert self.paths_LQ, 'LQ is required for downsampling.'
+ if not self.data_sz_mismatch_ok:
assert len(self.paths_LQ) == len(
self.paths_GT
), 'GT and LQ datasets have different number of images - {}, {}.'.format(
@@ -41,9 +44,8 @@ class GTLQDataset(data.Dataset):
def __getitem__(self, index):
if self.data_type == 'lmdb' and (self.GT_env is None or self.LQ_env is None):
self._init_lmdb()
- GT_path, LQ_path = None, None
scale = self.opt['scale']
- GT_size = self.opt['target_size']
+ GT_size = self.opt['target_size'] * scale
# get GT image
GT_path = self.paths_GT[index]
@@ -56,43 +58,19 @@ class GTLQDataset(data.Dataset):
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
# get LQ image
- if self.paths_LQ:
- LQ_path = self.paths_LQ[index]
- resolution = [int(s) for s in self.sizes_LQ[index].split('_')
- ] if self.data_type == 'lmdb' else None
- img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
- else: # down-sampling on-the-fly
- # randomly scale during training
- if self.opt['phase'] == 'train':
- random_scale = random.choice(self.random_scale_list)
- H_s, W_s, _ = img_GT.shape
+ lqind = index % len(self.paths_LQ)
+ LQ_path = self.paths_LQ[index % len(self.paths_LQ)]
+ resolution = [int(s) for s in self.sizes_LQ[index].split('_')
+ ] if self.data_type == 'lmdb' else None
+ img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
- def _mod(n, random_scale, scale, thres):
- rlt = int(n * random_scale)
- rlt = (rlt // scale) * scale
- return thres if rlt < thres else rlt
-
- H_s = _mod(H_s, random_scale, scale, GT_size)
- W_s = _mod(W_s, random_scale, scale, GT_size)
- img_GT = cv2.resize(img_GT, (W_s, H_s), interpolation=cv2.INTER_LINEAR)
- if img_GT.ndim == 2:
- img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
-
- H, W, _ = img_GT.shape
- # using matlab imresize
- img_LQ = util.imresize_np(img_GT, 1 / scale, True)
- if img_LQ.ndim == 2:
- img_LQ = np.expand_dims(img_LQ, axis=2)
+ # Create a downsampled version of the HQ image using matlab imresize.
+ img_Downsampled = util.imresize_np(img_GT, 1 / scale)
+ assert img_Downsampled.ndim == 3
if self.opt['phase'] == 'train':
- # if the image size is too small
H, W, _ = img_GT.shape
- if H < GT_size or W < GT_size:
- img_GT = cv2.resize(img_GT, (GT_size, GT_size), interpolation=cv2.INTER_LINEAR)
- # using matlab imresize
- img_LQ = util.imresize_np(img_GT, 1 / scale, True)
- if img_LQ.ndim == 2:
- img_LQ = np.expand_dims(img_LQ, axis=2)
+ assert H >= GT_size and W >= GT_size
H, W, C = img_LQ.shape
LQ_size = GT_size // scale
@@ -101,27 +79,35 @@ class GTLQDataset(data.Dataset):
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
+ img_Downsampled = img_Downsampled[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
rnd_h_GT, rnd_w_GT = int(rnd_h * scale), int(rnd_w * scale)
img_GT = img_GT[rnd_h_GT:rnd_h_GT + GT_size, rnd_w_GT:rnd_w_GT + GT_size, :]
# augmentation - flip, rotate
- img_LQ, img_GT = util.augment([img_LQ, img_GT], self.opt['use_flip'],
+ img_LQ, img_GT, img_Downsampled = util.augment([img_LQ, img_GT, img_Downsampled], self.opt['use_flip'],
self.opt['use_rot'])
if self.opt['color']: # change color space if necessary
- img_LQ = util.channel_convert(C, self.opt['color'],
- [img_LQ])[0] # TODO during val no definition
+ img_Downsampled = util.channel_convert(C, self.opt['color'],
+ [img_Downsampled])[0] # TODO during val no definition
# BGR to RGB, HWC to CHW, numpy to tensor
if img_GT.shape[2] == 3:
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQ = img_LQ[:, :, [2, 1, 0]]
+ img_Downsampled = img_Downsampled[:, :, [2, 1, 0]]
img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
+ img_Downsampled = torch.from_numpy(np.ascontiguousarray(np.transpose(img_Downsampled, (2, 0, 1)))).float()
img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
- if LQ_path is None:
- LQ_path = GT_path
- return {'LQ': img_LQ, 'GT': img_GT, 'LQ_path': LQ_path, 'GT_path': GT_path}
+ # This may seem really messed up, but let me explain:
+ # The goal is to re-use existing code as much as possible. SRGAN_model was coded to supersample, not downsample,
+ # but it can be retrofitted. To do so, we need to "trick" it. In this case the "input" is the HQ image and the
+ # "output" is the LQ image. SRGAN_model will be using a Generator and a Discriminator which already know this,
+ # we just need to trick its logic into following this rules.
+ # Do this by setting LQ(which is the input into the models)=img_GT and GT(which is the expected outpuut)=img_LQ.
+ # PIX is used as a reference for the pixel loss. Use the manually downsampled image for this.
+ return {'LQ': img_GT, 'GT': img_LQ, 'PIX': img_Downsampled, 'LQ_path': LQ_path, 'GT_path': GT_path}
def __len__(self):
return len(self.paths_GT)
diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py
index a1419a89..30c4f7aa 100644
--- a/codes/models/SRGAN_model.py
+++ b/codes/models/SRGAN_model.py
@@ -55,6 +55,9 @@ class SRGANModel(BaseModel):
else:
raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
self.l_fea_w = train_opt['feature_weight']
+ self.l_fea_w_decay = train_opt['feature_weight_decay']
+ self.l_fea_w_decay_steps = train_opt['feature_weight_decay_steps']
+ self.l_fea_w_minimum = train_opt['feature_weight_minimum']
else:
logger.info('Remove feature loss.')
self.cri_fea = None
@@ -143,13 +146,6 @@ class SRGANModel(BaseModel):
self.pix = data['PIX'].to(self.device)
def optimize_parameters(self, step):
-
- if step % 50 == 0:
- for i in range(self.var_L.shape[0]):
- utils.save_image(self.var_H[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\hr", "%05i_%02i.png" % (step, i)))
- utils.save_image(self.var_L[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\\lr", "%05i_%02i.png" % (step, i)))
- utils.save_image(self.pix[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\\pix", "%05i_%02i.png" % (step, i)))
-
# G
for p in self.netD.parameters():
p.requires_grad = False
@@ -157,6 +153,13 @@ class SRGANModel(BaseModel):
self.optimizer_G.zero_grad()
self.fake_H = self.netG(self.var_L)
+ if step % 50 == 0:
+ for i in range(self.var_L.shape[0]):
+ utils.save_image(self.var_H[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\hr", "%05i_%02i.png" % (step, i)))
+ utils.save_image(self.var_L[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\\lr", "%05i_%02i.png" % (step, i)))
+ utils.save_image(self.pix[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\\pix", "%05i_%02i.png" % (step, i)))
+ utils.save_image(self.fake_H[i].cpu().detach(), os.path.join("E:\\4k6k\\temp\\gen", "%05i_%02i.png" % (step, i)))
+
l_g_total = 0
if step % self.D_update_ratio == 0 and step > self.D_init_iters:
if self.cri_pix: # pixel loss
@@ -168,6 +171,11 @@ class SRGANModel(BaseModel):
l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
l_g_total += l_g_fea
+ # Decay the influence of the feature loss. As the model trains, the GAN will play a stronger role
+ # in the resultant image.
+ if step % self.l_fea_w_decay_steps == 0:
+ self.l_fea_w = max(self.l_fea_w_minimum, self.l_fea_w * self.l_fea_w_decay)
+
if self.opt['train']['gan_type'] == 'gan':
pred_g_fake = self.netD(self.fake_H)
l_g_gan = self.l_gan_w * self.cri_gan(pred_g_fake, True)
@@ -193,7 +201,8 @@ class SRGANModel(BaseModel):
# real
pred_d_real = self.netD(self.var_ref)
l_d_real = self.cri_gan(pred_d_real, True)
- l_d_real.backward()
+ with amp.scale_loss(l_d_real, self.optimizer_D, loss_id=2) as l_d_real_scaled:
+ l_d_real_scaled.backward()
# fake
pred_d_fake = self.netD(self.fake_H.detach()) # detach to avoid BP to G
l_d_fake = self.cri_gan(pred_d_fake, False)
@@ -222,12 +231,13 @@ class SRGANModel(BaseModel):
if self.cri_pix:
self.log_dict['l_g_pix'] = l_g_pix.item()
if self.cri_fea:
+ self.log_dict['feature_weight'] = self.l_fea_w
self.log_dict['l_g_fea'] = l_g_fea.item()
self.log_dict['l_g_gan'] = l_g_gan.item()
- self.log_dict['l_g_total'] = l_g_total.item()
- self.log_dict['l_d_real'] = l_d_real.item()
- self.log_dict['l_d_fake'] = l_d_fake.item()
- self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+ self.log_dict['l_g_total'] = l_g_total.item()
+ self.log_dict['l_d_real'] = l_d_real.item()
+ self.log_dict['l_d_fake'] = l_d_fake.item()
+ self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
def test(self):
self.netG.eval()
diff --git a/codes/models/__init__.py b/codes/models/__init__.py
index c95004c9..0767eeb3 100644
--- a/codes/models/__init__.py
+++ b/codes/models/__init__.py
@@ -7,7 +7,7 @@ def create_model(opt):
# image restoration
if model == 'sr': # PSNR-oriented super resolution
from .SR_model import SRModel as M
- elif model == 'srgan': # GAN-based super resolution, SRGAN / ESRGAN
+ elif model == 'srgan' or model == 'corruptgan': # GAN-based super resolution(SRGAN / ESRGAN), or corruption use same logic
from .SRGAN_model import SRGANModel as M
# video restoration
elif model == 'video_base':
diff --git a/codes/models/archs/HighToLowResNet.py b/codes/models/archs/HighToLowResNet.py
new file mode 100644
index 00000000..470359f2
--- /dev/null
+++ b/codes/models/archs/HighToLowResNet.py
@@ -0,0 +1,63 @@
+import functools
+import torch.nn as nn
+import torch.nn.functional as F
+import models.archs.arch_util as arch_util
+import torch
+
+
+class HighToLowResNet(nn.Module):
+ ''' ResNet that applies a noise channel to the input, then downsamples it. Currently only downscale=4 is supported. '''
+
+ def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16, downscale=4):
+ super(HighToLowResNet, self).__init__()
+ self.downscale = downscale
+
+ # We will always apply a noise channel to the inputs, account for that here.
+ in_nc += 1
+
+ self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+ basic_block = functools.partial(arch_util.ResidualBlock_noBN, nf=nf)
+ basic_block2 = functools.partial(arch_util.ResidualBlock_noBN, nf=nf*2)
+ # To keep the total model size down, the residual trunks will be applied across 3 downsampling stages.
+ # The first will be applied against the hi-res inputs and will have only 4 layers.
+ # The second will be applied after half of the downscaling and will also have only 6 layers.
+ # The final will be applied against the final resolution and will have all of the remaining layers.
+ self.trunk_hires = arch_util.make_layer(basic_block, 4)
+ self.trunk_medres = arch_util.make_layer(basic_block, 6)
+ self.trunk_lores = arch_util.make_layer(basic_block2, nb - 10)
+
+ # downsampling
+ if self.downscale == 4:
+ self.downconv1 = nn.Conv2d(nf, nf, 3, stride=2, padding=1, bias=True)
+ self.downconv2 = nn.Conv2d(nf, nf*2, 3, stride=2, padding=1, bias=True)
+ else:
+ raise EnvironmentError("Requested downscale not supported: %i" % (downscale,))
+
+ self.HRconv = nn.Conv2d(nf*2, nf*2, 3, stride=1, padding=1, bias=True)
+ self.conv_last = nn.Conv2d(nf*2, out_nc, 3, stride=1, padding=1, bias=True)
+
+ # activation function
+ self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+ # initialization
+ arch_util.initialize_weights([self.conv_first, self.HRconv, self.conv_last, self.downconv1, self.downconv2],
+ 0.1)
+
+ def forward(self, x):
+ # Noise has the same shape as the input with only one channel.
+ rand_feature = torch.randn((x.shape[0], 1) + x.shape[2:], device=x.device)
+ out = torch.cat([x, rand_feature], dim=1)
+
+ out = self.lrelu(self.conv_first(out))
+ out = self.trunk_hires(out)
+
+ if self.downscale == 4:
+ out = self.lrelu(self.downconv1(out))
+ out = self.trunk_medres(out)
+ out = self.lrelu(self.downconv2(out))
+ out = self.trunk_lores(out)
+
+ out = self.conv_last(self.lrelu(self.HRconv(out)))
+ base = F.interpolate(x, scale_factor=1/self.downscale, mode='bilinear', align_corners=False)
+ out += base
+ return out
diff --git a/codes/models/archs/discriminator_vgg_arch.py b/codes/models/archs/discriminator_vgg_arch.py
index ae51ba16..10a3ccdc 100644
--- a/codes/models/archs/discriminator_vgg_arch.py
+++ b/codes/models/archs/discriminator_vgg_arch.py
@@ -32,7 +32,7 @@ class Discriminator_VGG_128(nn.Module):
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(int(512 * 4 * input_img_factor * 4 * input_img_factor), 100)
+ self.linear1 = nn.Linear(int(nf * 8 * 4 * input_img_factor * 4 * input_img_factor), 100)
self.linear2 = nn.Linear(100, 1)
# activation function
diff --git a/codes/models/networks.py b/codes/models/networks.py
index 2b79249b..1b7563dc 100644
--- a/codes/models/networks.py
+++ b/codes/models/networks.py
@@ -3,6 +3,7 @@ import models.archs.SRResNet_arch as SRResNet_arch
import models.archs.discriminator_vgg_arch as SRGAN_arch
import models.archs.RRDBNet_arch as RRDBNet_arch
import models.archs.EDVR_arch as EDVR_arch
+import models.archs.HighToLowResNet as HighToLowResNet
import math
# Generator
@@ -20,6 +21,10 @@ def define_G(opt):
scale_per_step = math.sqrt(scale)
netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
nf=opt_net['nf'], nb=opt_net['nb'], interpolation_scale_factor=scale_per_step)
+ # image corruption
+ elif which_model == 'HighToLowResNet':
+ netG = HighToLowResNet.HighToLowResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
+ nf=opt_net['nf'], nb=opt_net['nb'], downscale=opt_net['scale'])
# video restoration
elif which_model == 'EDVR':
netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'],
diff --git a/codes/options/options.py b/codes/options/options.py
index 99181b34..5dc34b11 100644
--- a/codes/options/options.py
+++ b/codes/options/options.py
@@ -15,14 +15,14 @@ def parse(opt_path, is_train=True):
print('export CUDA_VISIBLE_DEVICES=' + gpu_list)
opt['is_train'] = is_train
- if opt['distortion'] == 'sr':
+ if opt['distortion'] == 'sr' or opt['distortion'] == 'downsample':
scale = opt['scale']
# datasets
for phase, dataset in opt['datasets'].items():
phase = phase.split('_')[0]
dataset['phase'] = phase
- if opt['distortion'] == 'sr':
+ if opt['distortion'] == 'sr' or opt['distortion'] == 'downsample':
dataset['scale'] = scale
is_lmdb = False
if dataset.get('dataroot_GT', None) is not None:
@@ -62,7 +62,7 @@ def parse(opt_path, is_train=True):
opt['path']['log'] = results_root
# network
- if opt['distortion'] == 'sr':
+ if opt['distortion'] == 'sr' or opt['distortion'] == 'downsample':
opt['network_G']['scale'] = scale
return opt
diff --git a/codes/train.py b/codes/train.py
index 8721d33f..2906463a 100644
--- a/codes/train.py
+++ b/codes/train.py
@@ -29,7 +29,7 @@ def init_dist(backend='nccl', **kwargs):
def main():
#### options
parser = argparse.ArgumentParser()
- parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='options/train/finetune_ESRGAN_blacked.yml')
+ parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='options/train/finetune_corruptGAN_adrianna.yml')
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
@@ -176,7 +176,7 @@ def main():
logger.info(message)
#### validation
if opt['datasets'].get('val', None) and current_step % opt['train']['val_freq'] == 0:
- if opt['model'] in ['sr', 'srgan'] and rank <= 0: # image restoration validation
+ if opt['model'] in ['sr', 'srgan', 'corruptgan'] and rank <= 0: # image restoration validation
# does not support multi-GPU validation
pbar = util.ProgressBar(len(val_loader))
avg_psnr = 0.