def getLocalScore(self,SR,HR,size=16,step=8):
bound = step
srimg = util.rgb2ycbcr(SR)
hrimg = util.rgb2ycbcr(HR)
#Sliding window approach to find local psnr and ssim values
srimg = np.pad(srimg,pad_width=((bound,bound),(bound,bound)),mode='symmetric')
hrimg = np.pad(hrimg,pad_width=((bound,bound),(bound,bound)),mode='symmetric')
h,w = hrimg.shape[:2]
psnr_vals = np.zeros(hrimg.shape)
ssim_vals = np.zeros(hrimg.shape)
for i in range(bound,h-bound-1,1):
for j in range(bound,w-bound-1,1):
img1 = srimg[i-bound:i+bound+1,j-bound:j+bound+1]
img2 = hrimg[i-bound:i+bound+1,j-bound:j+bound+1]
psnr_vals[i,j] = util.calc_psnr(img1 * 255,img2 * 255)
ssim_vals[i,j] = util.calc_ssim(img1 * 255,img2 * 255)
psnr_vals = psnr_vals[bound:-bound,bound:-bound]
ssim_vals = ssim_vals[bound:-bound,bound:-bound]
psnr_std = np.std(psnr_vals[psnr_vals > 0])
psnr_mean = np.mean(psnr_vals[psnr_vals > 0])
ssim_std = np.std(ssim_vals[ssim_vals > 0])
ssim_mean = np.mean(ssim_vals[ssim_vals > 0])
info = {'local_psnr':psnr_vals, 'local_ssim': ssim_vals, 'psnr_std': psnr_std,'psnr_mean':psnr_mean,'ssim_std':ssim_std,'ssim_mean':ssim_mean}
return info
def main():
# os.environ['CUDA_VISIBLE_DEVICES']='1' # You can specify your GPU device here.
parser = argparse.ArgumentParser(description='Train Super Resolution Models')
parser.add_argument('-opt', type=str, required=True, help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt)
if opt['train']['resume'] is False:
util.mkdir_and_rename(opt['path']['exp_root']) # rename old experiments if exists
util.mkdirs((path for key, path in opt['path'].items() if not key == 'exp_root' and \
not key == 'pretrain_G' and not key == 'pretrain_D'))
option.save(opt)
opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key.
else:
opt = option.dict_to_nonedict(opt)
if opt['train']['resume_path'] is None:
raise ValueError("The 'resume_path' does not declarate")
if opt['exec_debug']:
NUM_EPOCH = 100
opt['datasets']['train']['dataroot_HR'] = opt['datasets']['train']['dataroot_HR_debug']
opt['datasets']['train']['dataroot_LR'] = opt['datasets']['train']['dataroot_LR_debug']
else:
NUM_EPOCH = int(opt['train']['num_epochs'])
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
print("Random Seed: ", seed)
random.seed(seed)
torch.manual_seed(seed)
# create train and val dataloader
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
train_set = create_dataset(dataset_opt)
train_loader = create_dataloader(train_set, dataset_opt)
print('Number of train images in [%s]: %d' % (dataset_opt['name'], len(train_set)))
elif phase == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt)
print('Number of val images in [%s]: %d' % (dataset_opt['name'], len(val_set)))
elif phase == 'test':
pass
else:
raise NotImplementedError("Phase [%s] is not recognized." % phase)
if train_loader is None:
raise ValueError("The training data does not exist")
if opt['mode'] == 'sr':
solver = SRModel(opt)
else:
assert 'Invalid opt.mode [%s] for SRModel class!'
solver.summary(train_set[0]['LR'].size())
solver.net_init()
print('[Start Training]')
start_time = time.time()
start_epoch = 1
if opt['train']['resume']:
start_epoch = solver.load()
for epoch in range(start_epoch, NUM_EPOCH + 1):
# Initialization
solver.training_loss = 0.0
epoch_loss_log = 0.0
if opt['mode'] == 'sr' :
training_results = {'batch_size': 0, 'training_loss': 0.0}
else:
pass # TODO
train_bar = tqdm(train_loader)
# Train model
for iter, batch in enumerate(train_bar):
solver.feed_data(batch)
iter_loss = solver.train_step()
epoch_loss_log += iter_loss.item()
batch_size = batch['LR'].size(0)
training_results['batch_size'] += batch_size
if opt['mode'] == 'sr':
training_results['training_loss'] += iter_loss * batch_size
train_bar.set_description(desc='[%d/%d] Loss: %.4f ' % (
epoch, NUM_EPOCH, iter_loss))
else:
pass # TODO
solver.last_epoch_loss = epoch_loss_log / (len(train_bar))
train_bar.close()
time_elapse = time.time() - start_time
start_time = time.time()
print('Train Loss: %.4f' % (training_results['training_loss'] / training_results['batch_size']))
# validate
val_results = {'batch_size': 0, 'val_loss': 0.0, 'psnr': 0.0, 'ssim': 0.0}
if epoch % solver.val_step == 0 and epoch != 0:
print('[Validating...]')
start_time = time.time()
solver.val_loss = 0.0
vis_index = 1
for iter, batch in enumerate(val_loader):
visuals_list = []
solver.feed_data(batch)
iter_loss = solver.test(opt['chop'])
batch_size = batch['LR'].size(0)
val_results['batch_size'] += batch_size
visuals = solver.get_current_visual() # float cpu tensor
sr_img = np.transpose(util.quantize(visuals['SR'], opt['rgb_range']).numpy(), (1,2,0)).astype(np.uint8)
gt_img = np.transpose(util.quantize(visuals['HR'], opt['rgb_range']).numpy(), (1,2,0)).astype(np.uint8)
# calculate PSNR
crop_size = opt['scale']
cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :]
cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :]
val_results['val_loss'] += iter_loss * batch_size
val_results['psnr'] += util.calc_psnr(cropped_sr_img, cropped_gt_img)
val_results['ssim'] += util.calc_ssim(cropped_sr_img, cropped_gt_img)
if opt['mode'] == 'srgan':
pass # TODO
visuals_list.extend([util.quantize(visuals['HR'].squeeze(0), opt['rgb_range']),
util.quantize(visuals['SR'].squeeze(0), opt['rgb_range'])])
images = torch.stack(visuals_list)
img = thutil.make_grid(images, nrow=2, padding=5)
ndarr = img.byte().permute(1, 2, 0).numpy()
misc.imsave(os.path.join(solver.vis_dir, 'epoch_%d_%d.png' % (epoch, vis_index)), ndarr)
vis_index += 1
avg_psnr = val_results['psnr']/val_results['batch_size']
avg_ssim = val_results['ssim']/val_results['batch_size']
print('Valid Loss: %.4f | Avg. PSNR: %.4f | Avg. SSIM: %.4f | Learning Rate: %f'%(val_results['val_loss']/val_results['batch_size'], avg_psnr, avg_ssim, solver.current_learning_rate()))
time_elapse = start_time - time.time()
#if epoch%solver.log_step == 0 and epoch != 0:
# tensorboard visualization
solver.training_loss = training_results['training_loss'] / training_results['batch_size']
solver.val_loss = val_results['val_loss'] / val_results['batch_size']
solver.tf_log(epoch)
# statistics
if opt['mode'] == 'sr' :
solver.results['training_loss'].append(solver.training_loss.cpu().data.item())
solver.results['val_loss'].append(solver.val_loss.cpu().data.item())
solver.results['psnr'].append(avg_psnr)
solver.results['ssim'].append(avg_ssim)
else:
pass # TODO
is_best = False
if solver.best_prec < solver.results['psnr'][-1]:
solver.best_prec = solver.results['psnr'][-1]
is_best = True
solver.save(epoch, is_best)
# update lr
solver.update_learning_rate(epoch)
data_frame = pd.DataFrame(
data={'training_loss': solver.results['training_loss']
, 'val_loss': solver.results['val_loss']
, 'psnr': solver.results['psnr']
, 'ssim': solver.results['ssim']
},
index=range(1, NUM_EPOCH+1)
)
data_frame.to_csv(os.path.join(solver.results_dir, 'train_results.csv'),
index_label='Epoch')
