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')