import os
import math
import time
import datetime
from functools import reduce
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import scipy.misc as misc
from skimage.restoration import denoise_bilateral
import torch
import torch.optim as optim
import torch.optim.lr_scheduler as lrs
class timer():
def __init__(self):
self.acc = 0
self.tic()
#print ("2-1-1-checkpoint")
def tic(self):
self.t0 = time.time()
#print ("2-1-2-checkpoint")
def toc(self):
return time.time() - self.t0
#print ("2-1-3-checkpoint")
def hold(self):
self.acc += self.toc()
#print ("2-1-4-checkpoint")
def release(self):
ret = self.acc
self.acc = 0
#print ("2-1-5-checkpoint")
return ret
def reset(self):
self.acc = 0
#print ("2-1-6-checkpoint")
class checkpoint():
def __init__(self, args):
self.args = args
self.ok = True
self.log = torch.Tensor()
now = datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S')
if args.load == '.':
if args.save == '.': args.save = now
self.dir = '../experiment/' + args.save
else:
self.dir = '../experiment/' + args.load
if not os.path.exists(self.dir):
args.load = '.'
else:
self.log = torch.load(self.dir + '/psnr_log.pt')
print('Continue from epoch {}...'.format(len(self.log)))
if args.reset:
os.system('rm -rf ' + self.dir)
args.load = '.'
def _make_dir(path):
if not os.path.exists(path): os.makedirs(path)
_make_dir(self.dir)
_make_dir(self.dir + '/model')
_make_dir(self.dir + '/results')
_make_dir(self.dir + '/residuals')
_make_dir(self.dir + '/branches')
open_type = 'a' if os.path.exists(self.dir + '/log.txt') else 'w'
self.log_file = open(self.dir + '/log.txt', open_type)
with open(self.dir + '/config.txt', open_type) as f:
f.write(now + '\n\n')
for arg in vars(args):
f.write('{}: {}\n'.format(arg, getattr(args, arg)))
f.write('\n')
#print ("2-2-1-checkpoint")
def save(self, trainer, epoch, is_best=False):
trainer.model.save(self.dir, epoch, is_best=is_best)
trainer.loss.save(self.dir)
trainer.loss.plot_loss(self.dir, epoch)
self.plot_psnr(epoch)
torch.save(self.log, os.path.join(self.dir, 'psnr_log.pt'))
torch.save(
trainer.optimizer.state_dict(),
os.path.join(self.dir, 'optimizer.pt')
)
#print ("2-2-2-checkpoint")
def add_log(self, log):
self.log = torch.cat([self.log, log])
#print ("2-2-3-checkpoint")
def write_log(self, log, refresh=False):
#print(log)
self.log_file.write(log + '\n')
if refresh:
self.log_file.close()
self.log_file = open(self.dir + '/log.txt', 'a')
#print ("2-2-4-checkpoint")
def done(self):
self.log_file.close()
#print ("2-2-5-checkpoint")
def plot_psnr(self, epoch):
axis = np.linspace(1, epoch, epoch)
label = 'SR on {}'.format(self.args.data_test)
fig = plt.figure()
plt.title(label)
for idx_scale, scale in enumerate(self.args.scale):
plt.plot(
axis,
self.log[:, idx_scale].numpy(),
label='Scale {}'.format(scale)
)
plt.legend()
plt.xlabel('Epochs')
plt.ylabel('PSNR')
plt.grid(True)
plt.savefig('{}/test_{}.pdf'.format(self.dir, self.args.data_test))
plt.close(fig)
#print ("2-2-5-checkpoint")
def save_results(self, filename, save_list, scale):
filename = '{}/results/{}_x{}_'.format(self.dir, filename, scale)
postfix = ('SR', 'LR', 'HR')
for v, p in zip(save_list, postfix):
normalized = v[0].data.mul(255 / self.args.rgb_range)
ndarr = normalized.byte().permute(1, 2, 0).cpu().numpy()
if ndarr.shape[-1] == 1:
ndarr = ndarr[:,:,0]
misc.imsave('{}{}.png'.format(filename, p), ndarr)
#print ("2-2-6-checkpoint")
def save_residuals(self, filename, save_list, scale):
filename = '{}/residuals/{}_x{}'.format(self.dir, filename, scale)
sr, hr = save_list[0], save_list[-1]
def _prepare(x):
normalized = x[0].data.mul(1. / self.args.rgb_range)
out = normalized.permute(1,2,0).cpu().numpy()
if out.shape[-1] == 1:
out = out[:,:,0]
return out
ndarr_sr, ndarr_hr = _prepare(sr), _prepare(hr)
out = np.abs(ndarr_hr - ndarr_sr)
misc.imsave('{}.png'.format(filename), out)
#print ("2-2-7-checkpoint")
def save_branches(self, filename, save_list, scale):
filename = '{}/branches/{}_x{}'.format(self.dir, filename, scale)
def _prepare(x, residual):
normalized = x[0].data.mul(1. / self.args.rgb_range)
if not residual:
out = normalized.permute(1,2,0).cpu().numpy()
else:
out = np.abs(normalized.permute(1,2,0).cpu().numpy())
if out.shape[-1] == 1:
out = out[:,:,0]
return out
for i, branch_output in enumerate(save_list):
ndarr = _prepare(branch_output, not (i==0))
misc.imsave('{}{}.png'.format(filename, '_branch{}'.format(i)), ndarr)
#print ("2-2-8-checkpoint")
return
def get_bilateral(tensor, rgb_range):
tensor = tensor.numpy().transpose(0,2,3,1) / rgb_range
out = np.zeros_like(tensor)
for i, t in enumerate(tensor):
out[i] = denoise_bilateral(t)
#print ("2-3-checkpoint")
return torch.Tensor(out.transpose(0,3,1,2)) * rgb_range
def quantize(img, rgb_range):
pixel_range = 255 / rgb_range
return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)
#print ("2-4-checkpoint")
def calc_psnr(sr, hr, scale, rgb_range, benchmark=False):
diff = (sr - hr).data.div(rgb_range)
if benchmark:
shave = scale
if diff.size(1) > 1:
convert = diff.new(1, 3, 1, 1)
convert[0, 0, 0, 0] = 65.738
convert[0, 1, 0, 0] = 129.057
convert[0, 2, 0, 0] = 25.064
diff.mul_(convert).div_(256)
diff = diff.sum(dim=1, keepdim=True)
else:
shave = scale + 6
valid = diff[:, :, shave:-shave, shave:-shave]
mse = valid.pow(2).mean()
#print ("2-5-checkpoint")
return -10 * math.log10(mse)
def make_optimizer(args, my_model):
trainable = filter(lambda x: x.requires_grad, my_model.parameters())
if args.optimizer == 'SGD':
optimizer_function = optim.SGD
kwargs = {'momentum': args.momentum}
elif args.optimizer == 'ADAM':
optimizer_function = optim.Adam
kwargs = {
'betas': (args.beta1, args.beta2),
'eps': args.epsilon
}
elif args.optimizer == 'RMSprop':
optimizer_function = optim.RMSprop
kwargs = {'eps': args.epsilon}
kwargs['lr'] = args.lr
kwargs['weight_decay'] = args.weight_decay
#print ("2-6-checkpoint")
return optimizer_function(trainable, **kwargs)
def make_scheduler(args, my_optimizer):
if args.decay_type == 'step':
scheduler = lrs.StepLR(
my_optimizer,
step_size=args.lr_decay,
gamma=args.gamma
)
elif args.decay_type.find('step') >= 0:
milestones = args.decay_type.split('_')
milestones.pop(0)
milestones = list(map(lambda x: int(x), milestones))
scheduler = lrs.MultiStepLR(
my_optimizer,
milestones=milestones,
gamma=args.gamma
)
#print ("2-7-checkpoint")
return scheduler
