def img2HQimg(self): select_idx = data_util.index_generation(0, 1, 5, padding=self.padding) imgs = data_util.read_img_seq(os.path.dirname(self.input_path)) imgs_in = imgs.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(self.device) output = util.single_forward(self.model, imgs_in) output = util.tensor2img(output.squeeze(0)) util.save_img(img=output, img_path=self.output_path)
def __getitem__(self, index): folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] # print(self.data_info['path_LQ'][index], '\n', self.data_info['path_SLQ'][index]) select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding']) imgs_GT = self.imgs_GT[folder].index_select( 0, torch.LongTensor(select_idx)) if self.opt['degradation_mode'] == 'preset': if self.opt['N_frames'] == 5: imgs_LR = self.imgs_LQ[folder][idx, 1:-1] imgs_SuperLR = self.imgs_SLQ[folder][idx, 1:-1] else: imgs_LR = self.imgs_LQ[folder][idx] imgs_SuperLR = self.imgs_SLQ[folder][idx] else: imgs_LR = self.imgs_LQ[folder].index_select( 0, torch.LongTensor(select_idx)) imgs_SuperLR = self.imgs_SLQ[folder].index_select( 0, torch.LongTensor(select_idx)) return { 'SuperLQs': imgs_SuperLR, 'LQs': imgs_LR, 'GT': imgs_GT, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border }
def __getitem__(self, index): # path_LQ = self.data_info['path_LQ'][index] # path_GT = self.data_info['path_GT'][index] folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] if self.cache_data: select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding']) imgs_LQ = self.imgs_LQ[folder].index_select( 0, torch.LongTensor(select_idx)) img_GT = self.imgs_GT[folder][idx] else: pass # TODO return { 'LQs': imgs_LQ, 'GT': img_GT, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border }
def get_min_in_axis(img, s_length, direction='horizontal'): # img: H W T smooth_img = np.zeros_like(img) h, w = img.shape[-2], img.shape[-1] if direction == 'horizontal': for i in range(w): smooth_img[:, :, i] = np.amin( img[:, :, data_util.index_generation(i, w, s_length, 'replicate')], axis=2) elif direction == 'vertical': for i in range(h): smooth_img[:, i, :] = np.amin( img[:, data_util.index_generation(i, h, s_length, 'replicate' ), :], axis=1) return smooth_img
def __getitem__(self, index): # path_LQ = self.data_info['path_LQ'][index] # path_GT = self.data_info['path_GT'][index] folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding']) if self.data_type == 'lmdb': if self.GT_env is None or self.LQ_env is None: self._init_lmdb() key = self.lmdb_paths_GT[index] name_a, name_b = key.split('_') center_frame_idx = int(name_b) GT_size_tuple = self.opt['GT_shape'] LQ_size_tuple = self.opt['LQ_shape'] img_GT = util.read_img(self.GT_env, key, GT_size_tuple) img_LQ_l = [] for v in select_idx: img_LQ = util.read_img(self.LQ_env, '{}_{:08d}'.format(name_a, v), LQ_size_tuple) img_LQ_l.append(img_LQ) # stack LQ images to NHWC, N is the frame number img_LQs = np.stack(img_LQ_l, axis=0) # BGR to RGB, HWC to CHW, numpy to tensor img_GT = img_GT[:, :, [2, 1, 0]] img_LQs = img_LQs[:, :, :, [2, 1, 0]] img_GT = torch.from_numpy( np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float() imgs_LQ = torch.from_numpy( np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float() elif self.cache_data: imgs_LQ = self.imgs_LQ[folder].index_select( 0, torch.LongTensor(select_idx)) img_GT = self.imgs_GT[folder][idx] else: imgs_LQ = util.read_img_seq(self.imgs_LQ[folder]).index_select( 0, torch.LongTensor(select_idx)) img_GT = util.read_img_seq(self.imgs_GT[folder])[idx] return { 'LQs': imgs_LQ, 'GT': img_GT, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border }
def read_data(self, index): video_name = self.videos_name[index] idx, max_idx = self.idx_maxidx[index].split('/') idx, max_idx = int(idx), int(max_idx) indices = util.index_generation(idx, max_idx, self.n_frames) lr_path = osp.dirname(self.lr_paths[index]) filename = osp.basename(self.lr_paths[index]) lrs = [] for i in indices: lr = util.read_img(None, osp.join(lr_path, '{:03d}.png'.format(i))) lrs.append(lr) gt = util.read_img(None, self.gt_paths[index]) return lrs, gt, video_name, filename
def imgs2HQimgs(self): """create a HQ imgs from a cached imgs and save them in a ./output/ subfolder of curren directory""" util.mkdir(self.output_path) print(100 * len(os.listdir(self.imgs_cache))) for batch_num in range(100, ( 100 * len(os.listdir(self.imgs_cache)) ) + 100, 100): self.read_batch(batch_num) bar = self.progress.ProgressBar(max_value=len(self.imgs_paths)) for img_idx, img_path in enumerate(self.imgs_paths): select_idx = data_util.index_generation(img_idx, len(self.imgs_paths), self.N_in, padding=self.padding) imgs_in = self.imgs.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(self.device) output = util.single_forward(self.model, imgs_in) output = util.tensor2img(output.squeeze(0)) print(output.shape) cv2.imwrite(f'{self.output_path}{img_idx}.png', output) bar.update(img_idx)
def video2HQvideo(self): """create a HQ video from a cached imgs and save it into output path""" fourcc = cv2.VideoWriter_fourcc(*'mp4v') video_writer = cv2.VideoWriter(self.output_path, fourcc, self.fps, self.img_shape) for batch_num in range(100, ( 100 * len(os.listdir(self.imgs_cache)) ) + 100, 100): self.read_batch(batch_num) bar = self.progress.ProgressBar(max_value=len(self.imgs_paths)) for img_idx, img_path in enumerate(self.imgs_paths): select_idx = data_util.index_generation(img_idx, len(self.imgs_paths), self.N_in, padding=self.padding) imgs_in = self.imgs.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(self.device) output = util.single_forward(self.model, imgs_in) output = util.tensor2img(output.squeeze(0)) video_writer.write(output) bar.update(img_idx) video_writer.release()
def __getitem__(self, index): folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding']) imgs = self.imgs[folder].index_select(0, torch.LongTensor(select_idx)) return { 'LQs': imgs, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border }
def __getitem__(self, index): scale = self.scale if self.data_name.lower() in ['vid4', 'reds4']: folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding']) if self.cache_data: imgs = self.imgs_GT[folder][select_idx] N, H, W, C = imgs.shape img_GT = imgs[self.half_N_frames] / 255. imgs_LQ = [ np.array( Image.fromarray(img).resize( (LQ_W, LQ_H), Image.BICUBIC)) / 255. for img in imgs ] else: pass # TODO out = { 'LQs': imgs_LQ, 'GT': img_GT, 'folder': folder, 'scale': scale, 'idx': self.data_info['idx'][index], 'border': border, } else: key = self.data_info['path_GT'][index]
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '4' data_mode = 'youku' split_test = True if split_test: forward_fn = PatchedForward(3, 8, 4, 4) else: forward_fn = PaddingForward(4, 4) ############################################################################ #### model model_path = '../experiments/pretrained_models/EDVR_YOUKU_M_woTSA.pth' N_in = 5 predeblur, HR_in = False, False back_RBs = 10 model = EDVR_arch.EDVR(64, N_in, 8, 5, 10, predeblur=predeblur, HR_in=HR_in, w_TSA=False, block_type='rcab', non_local=False) ssim_fn = SSIM() #### dataset test_dataset_folder = '/media/tclwh2/public/youku/val/lq' GT_dataset_folder = '/media/tclwh2/public/youku/val/gt' #### evaluation border_frame = N_in // 2 # border frames when evaluate # temporal padding mode padding = 'new_info' save_imgs = False save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Split test: {}'.format(split_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) ssim_fn = ssim_fn.to(device) ssim_fn.eval() avg_psnr_l, avg_ssim_l = [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr = 0 avg_ssim = 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) with torch.no_grad(): model_output = forward_fn(model, imgs_in) GT = np.copy(img_GT_l[img_idx][:, :, [2, 1, 0]]) GT = torch.from_numpy(GT.transpose(2, 0, 1)).unsqueeze_(0).to(device) crt_ssim = ssim_fn(model_output, GT).data.cpu().item() output = model_output.data.float().cpu() output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) crt_psnr = util.calculate_psnr(output * 255, GT * 255) logger.info('{:3d} - {:25} \tPSNR/SSIM: {:.6f}/{:.6f}'.format( img_idx + 1, img_name, crt_psnr, crt_ssim)) avg_psnr += crt_psnr avg_ssim += crt_ssim avg_psnr /= max_idx avg_ssim /= max_idx avg_psnr_l.append(avg_psnr) avg_ssim_l.append(avg_ssim) logger.info( 'Folder {} - Average PSNR/SSIM: {:.6f}/{:.6f} for {} frames. '. format(subfolder_name, avg_psnr, avg_ssim, max_idx)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, ssim in zip(subfolder_name_l, avg_psnr_l, avg_ssim_l): logger.info('Folder {} - Average PSNR/SSIM: {:.6f}/{:.6f}. '.format( subfolder_name, psnr, ssim)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Total Average PSNR/SSIM: {:.6f}/{:.6f} for {} clips. '.format( sum(avg_psnr_l) / len(avg_psnr_l), sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l)))
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '1' stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False #### model data_mode = 'sharp' if stage == 1: model_path = '../experiments/001_EDVRwoTSA_scratch_lr4e-4_600k_SR4K_LrCAR4S_64_20_5/models/600000_G.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' N_in = 5 # use N_in images to restore one HR image predeblur, HR_in = False, False back_RBs = 20 if stage == 2: HR_in = True back_RBs = 20 model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=True) #### dataset if stage == 1: test_dataset_folder = '/home/mcc/4khdr/image/540p_test' else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'sharp': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) # for each subfolder for subfolder in subfolder_l: subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output, [int(cv2.IMWRITE_PNG_COMPRESSION), 1]) logger.info('{:3d} - {:25}'.format(img_idx + 1, img_name))
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' data_mode = 'SDR_4bit' stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'SDR_4bit': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError # use N_in images to restore one high bitdepth image N_in = 5 # predeblur: predeblur for blurry input # HR_in: downsample high resolution input predeblur, HR_in = False, False back_RBs = 40 predeblur = True HR_in = True if data_mode == 'SDR_4bit': # predeblur, HR_in = False, True pass if stage == 2: HR_in = True back_RBs = 20 # EDVR(num_feature_map, num_input_frames, deformable_groups?, front_RBs, # back_RBs, predeblur, HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### dataset if stage == 1: test_dataset_folder = '../datasets/{}'.format(data_mode) else: test_dataset_folder = '../' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/SDR_10bit/' #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LBD and GT images #### resize to avoid cuda out of memory, 2160x3840->720x1280 imgs_LBD = data_util.read_img_seq(subfolder, scale=65535., zoomout=(1280, 720)) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append( data_util.read_img(None, img_GT_path, scale=65535., zoomout=True)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] # generate frame index select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LBD.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: # self ensemble with fipping input at four different directions output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0), out_type=np.uint16) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR # output = output / 255. output = output / 65535. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 65535, GT * 65535) logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format( img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip( subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
print('Done processing') break prep_frame =prepare_frame(frame) currentFrame += 1 if currentFrame< 5: prep_frame_lst.append(prep_frame) else: prep_frame_lst.append(prep_frame) # stack to Torch tensor imgs_LQ = get_Torch_tensor_from_prep_frame_lst(prep_frame_lst) # process each image count +=1 img_name = "frame_%05i" % count select_idx = data_util.index_generation(3, 5, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) # To stop duplicate images # When everything done, release the capture cap.release()
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '6' test_set = 'AI4K_val' # Vid4 | YouKu10 | REDS4 | AI4K_val | zhibo | AI4K_val_bic test_name = 'PCD_Vis_Test_35_ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW_A01xxx_900000_AI4K_5000' # 'AI4K_val_Denoise_A02_420000' data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic N_in = 5 # load test set if test_set == 'Vid4': test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' elif test_set == 'YouKu10': test_dataset_folder = '../datasets/YouKu10/LR' GT_dataset_folder = '../datasets/YouKu10/HR' elif test_set == 'YouKu_val': test_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_lr_bmp' GT_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_hr_bmp' elif test_set == 'REDS4': test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) GT_dataset_folder = '../datasets/REDS4/GT' elif test_set == 'AI4K_val': test_dataset_folder = '/home/yhliu/AI4K/contest2/val2_LR_png/' GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png/' elif test_set == 'AI4K_val_bic': test_dataset_folder = '/home/yhliu/AI4K/contest1/val1_LR_png_bic/' GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png_bic/' elif test_set == 'zhibo': test_dataset_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/Test_video_LR/' GT_dataset_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/Test_video_HR/' flip_test = False #model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' #model_path = '../experiments/A01b/models/250000_G.pth' #model_path = '../experiments/A02_predenoise/models/415000_G.pth' model_path = '../experiments/A37_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd/models/5000_G.pth' predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #model = my_EDVR_arch.MYEDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #model = my_EDVR_arch.MYEDVR_RES(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True #True | False save_folder = '../results/{}'.format(test_name) if test_set == 'zhibo': save_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/SR_png_sample_150' util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) model = nn.DataParallel(model) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) print(subfolder_l) print(subfolder_GT_l) #exit() # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) print(img_path_l) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): #print(img_GT_path) img_GT_l.append(data_util.read_img(None, img_GT_path)) #print(img_GT_l[0].shape) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).cpu() #to(device) print(imgs_in.size()) if flip_test: output = util.flipx4_forward(model, imgs_in) else: start_time = time.time() output = util.single_forward(model, imgs_in) end_time = time.time() print('Forward One image:', end_time - start_time) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel ''' if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) ''' output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format( img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip( subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' data_mode = 'HDR' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = True ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' else: raise ValueError('Vid4 does not support stage 2.') elif data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' elif data_mode == 'HDR': # model_path = '../experiments/001_EDVR_scratch_lr4e-4_600k_HDR_LrCAR4S/models/20000_G.pth' # model_path = '../experiments/005_EDVRwoTSA_scratch_lr4e-4_600k_HDR_LrCAR4S/models/490000_G.pth' model_path = '../experiments/pretrained_models/50000_G.pth' else: raise NotImplementedError if data_mode == 'Vid4': N_in = 7 # use N_in images to restore one HR image else: N_in = 5 predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if data_mode == 'HDR': predeblur = True if stage == 2: HR_in = True back_RBs = 20 model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=True) #### dataset GT_dataset_folder = None if data_mode == 'Vid4': test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' elif data_mode == 'HDR': test_dataset_folder = '../datasets/HDR/valid/new_method/540p' GT_dataset_folder = '../datasets/HDR/valid/new_method/4k' # test_dataset_folder = '../datasets/HDR/valid/sequences_540' # GT_dataset_folder = '../datasets/HDR/valid/sequences_4k' else: if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}_50000'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): # print(subfolder, subfolde if '10675978' not in subfolder: print('pass') continue subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR # output = output / 255. # GT = np.copy(img_GT_l[img_idx]) # # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel # if data_mode == 'Vid4': # bgr2y, [0, 1] # GT = data_util.bgr2ycbcr(GT, only_y=True) # output = data_util.bgr2ycbcr(output, only_y=True) # output, GT = util.crop_border([output, GT], crop_border) # crt_psnr = util.calculate_psnr(output * 255, GT * 255) # # crt_ssim = util.calculate_ssim(output * 255, GT * 255) logger.info('{:3d} - {:25} '.format(img_idx + 1, img_name))
def main(opts): ################## configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpus cache_all_imgs = opts.cache > 0 n_gpus = len(opts.gpus.split(',')) flip_test, save_imgs = False, False scale = 4 N_in, nf = 5, 64 back_RBs = 10 w_TSA = False predeblur, HR_in = False, False crop_border = 0 border_frame = N_in // 2 padding = 'new_info' ################## model files #################### model_dir = opts.model_dir if osp.isfile(model_dir): model_names = [osp.basename(model_dir)] model_dir = osp.dirname(model_dir) elif osp.isdir(model_dir): model_names = [ x for x in os.listdir(model_dir) if str.isdigit(x.split('_')[0]) ] model_names = sorted(model_names, key=lambda x: int(x.split("_")[0])) else: raise IOError('Invalid model_dir: {}'.format(model_dir)) ################## dataset ######################## test_subs = sorted(os.listdir(opts.test_dir)) gt_subs = os.listdir(opts.gt_dir) valid_test_subs = [sub in gt_subs for sub in test_subs] assert (all(valid_test_subs)), 'Invalid sub folders exists in {}'.format( opts.test_dir) scale = float(os.path.basename(os.path.dirname(opts.test_dir))[1:]) if cache_all_imgs: print('Cacheing all testing images ...') all_imgs = {} for sub in test_subs: print('Reading sub-folder: {} ...'.format(sub)) test_sub_dir = osp.join(opts.test_dir, sub) gt_sub_dir = osp.join(opts.gt_dir, sub) all_imgs[sub] = {'test': [], 'gt': []} im_names = sorted(os.listdir(test_sub_dir)) for i, name in enumerate(im_names): test_im_path = osp.join(test_sub_dir, name) gt_im_path = osp.join(gt_sub_dir, name) test_im = cv2.imread(test_im_path, cv2.IMREAD_UNCHANGED)[:, :, (2, 1, 0)] test_im = test_im.astype(np.float32).transpose( (2, 0, 1)) / 255. all_imgs[sub]['test'].append(test_im) gt_im = cv2.imread(gt_im_path, cv2.IMREAD_UNCHANGED).astype(np.float32) all_imgs[sub]['gt'].append(gt_im) all_psnrs = [] for model_name in model_names: model_path = osp.join(model_dir, model_name) exp_name = model_name.split('_')[0] if 'meta' in opts.mode.lower(): model = EDVR_arch.MetaEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=w_TSA) elif opts.mode.lower() == 'edvr': model = EDVR_arch.EDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=w_TSA) elif opts.mode.lower() == 'upedvr': model = EDVR_arch.UPEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True) elif opts.mode.lower() == 'upcont1': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=False) elif opts.mode.lower() == 'upcont3': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=True) elif opts.mode.lower() == 'upcont2': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=True) else: raise TypeError('Unknown model mode: {}'.format(opts.mode)) save_folder = osp.join(opts.save_dir, exp_name) util.mkdirs(save_folder) util.setup_logger(exp_name, save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger(exp_name) #### log info logger.info('Data: {}'.format(opts.test_dir)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() if n_gpus > 1: model = nn.DataParallel(model) model = model.to(device) avg_psnrs, avg_psnr_centers, avg_psnr_borders = [], [], [] avg_ssims, avg_ssim_centers, avg_ssim_borders = [], [], [] evaled_subs = [] # for each subfolder for sub in test_subs: evaled_subs.append(sub) test_sub_dir = osp.join(opts.test_dir, sub) gt_sub_dir = osp.join(opts.gt_dir, sub) img_names = sorted(os.listdir(test_sub_dir)) max_idx = len(img_names) if save_imgs: save_subfolder = osp.join(save_folder, sub) util.mkdirs(save_subfolder) #### get LQ and GT images if not cache_all_imgs: img_LQs, img_GTs = [], [] for i, name in enumerate(img_names): test_im_path = osp.join(test_sub_dir, name) gt_im_path = osp.join(gt_sub_dir, name) test_im = cv2.imread(test_im_path, cv2.IMREAD_UNCHANGED)[:, :, (2, 1, 0)] test_im = test_im.astype(np.float32).transpose( (2, 0, 1)) / 255. gt_im = cv2.imread(gt_im_path, cv2.IMREAD_UNCHANGED).astype(np.float32) img_LQs.append(test_im) img_GTs.append(gt_im) else: img_LQs = all_imgs[sub]['test'] img_GTs = all_imgs[sub]['gt'] avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 avg_ssim, avg_ssim_border, avg_ssim_center = 0, 0, 0 # process each image for i in range(0, max_idx, n_gpus): end = min(i + n_gpus, max_idx) select_idxs = [ data_util.index_generation(j, max_idx, N_in, padding=padding) for j in range(i, end) ] imgs = [] for select_idx in select_idxs: im = torch.from_numpy( np.stack([img_LQs[k] for k in select_idx])) imgs.append(im) if (i + n_gpus) > max_idx: for _ in range(max_idx, i + n_gpus): imgs.append(torch.zeros_like(im)) imgs = torch.stack(imgs, 0).to(device) if flip_test: output = util.flipx4_forward(model, imgs) else: if 'meta' in opts.mode.lower(): output = util.meta_single_forward( model, imgs, scale, n_gpus) if 'up' in opts.mode.lower(): output = util.up_single_forward(model, imgs, scale) else: output = util.single_forward(model, imgs) output = [ util.tensor2img(x).astype(np.float32) for x in output ] if save_imgs: for ii in range(i, end): cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_names[ii])), output[ii - i].astype(np.uint8)) # calculate PSNR GT = np.copy(img_GTs[i:end]) output = util.crop_border(output, crop_border) GT = util.crop_border(GT, crop_border) for m in range(i, end): crt_psnr = util.calculate_psnr(output[m - i], GT[m - i]) crt_ssim = util.calculate_ssim(output[m - i], GT[m - i]) logger.info( '{:3d} - {:25} \tPSNR: {:.6f} dB SSIM: {:.6}'. format(m + 1, img_names[m], crt_psnr, crt_ssim)) if m >= border_frame and m < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr avg_ssim_center += crt_ssim N_center += 1 else: # border frames avg_psnr_border += crt_psnr avg_ssim_border += crt_ssim N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border) avg_ssim_center = avg_ssim_center / N_center avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border avg_psnrs.append(avg_psnr) avg_psnr_centers.append(avg_psnr_center) avg_psnr_borders.append(avg_psnr_border) avg_ssims.append(avg_ssim) avg_ssim_centers.append(avg_ssim_center) avg_ssim_borders.append(avg_ssim_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( sub, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('Folder {} - Average SSIM: {:.6f} for {} frames; ' 'Center SSIM: {:.6f} for {} frames; ' 'Border SSIM: {:.6f} for {} frames.'.format( sub, avg_ssim, (N_center + N_border), avg_ssim_center, N_center, avg_ssim_border, N_border)) logger.info('################ Tidy Outputs ################') for sub_name, psnr, psnr_center, psnr_border, ssim, ssim_center, ssim_border in zip( evaled_subs, avg_psnrs, avg_psnr_centers, avg_psnr_borders, avg_ssims, avg_ssim_centers, avg_ssim_borders): logger.info( 'Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sub_name, psnr, psnr_center, psnr_border)) logger.info('Folder {} - Average SSIM: {:.6f} ' 'Center SSIM: {:.6f} Border SSIM: {:.6f} '.format( sub_name, ssim, ssim_center, ssim_border)) logger.info('################ Final Results ################') logger.info('Data: {}'.format(opts.test_dir)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnrs) / len(avg_psnrs), len(test_subs), sum(avg_psnr_centers) / len(avg_psnr_centers), sum(avg_psnr_borders) / len(avg_psnr_borders))) logger.info('Total Average SSIM: {:.6f} for {} clips. ' 'Center SSIM: {:.6f} Border SSIM: {:.6f} '.format( sum(avg_ssims) / len(avg_ssims), len(test_subs), sum(avg_ssim_centers) / len(avg_ssim_centers), sum(avg_ssim_borders) / len(avg_ssim_borders)))
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: model_path = '../experiments/pretrained_models/cinepak_small2.pth' else: raise ValueError('Vid4 does not support stage 2.') elif data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError if data_mode == 'Vid4': N_in = 7 # use N_in images to restore one HR image else: N_in = 5 predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### dataset if data_mode == 'Vid4': test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' else: if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) # for each subfolder for subfolder in subfolder_l: print('Processing video {:s}'.format(subfolder)) subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) # process each image for img_idx, img_path in enumerate(img_path_l): print('\tProcessing frame {:s}') img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output)
def main(): # Create object for parsing command-line options parser = argparse.ArgumentParser(description="Test with EDVR, requre path to test dataset folder.") # Add argument which takes path to a bag file as an input parser.add_argument("-i", "--input", type=str, help="Path to test folder") # Parse the command line arguments to an object args = parser.parse_args() # Safety if no parameter have been given if not args.input: print("No input paramater have been given.") print("For help type --help") exit() folder_name = args.input.split("/")[-1] if folder_name == '': index = len(args.input.split("/")) - 2 folder_name = args.input.split("/")[index] ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' else: raise ValueError('Vid4 does not support stage 2.') else: raise NotImplementedError if data_mode == 'Vid4': N_in = 7 # use N_in images to restore one HR image else: N_in = 5 predeblur, HR_in = False, False back_RBs = 40 model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### dataset if data_mode == 'Vid4': # debug test_dataset_folder = os.path.join(args.input, 'BIx4') GT_dataset_folder = os.path.join(args.input, 'GT') else: if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(folder_name) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(folder_name, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format(subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip(subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(folder_name, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main(): ################# # configurations ################# #torch.backends.cudnn.benchmark = True #torch.backends.cudnn.enabled = True device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '5' test_set = 'AI4K_test' # Vid4 | YouKu10 | REDS4 | AI4K_test data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic test_name = 'Contest2_Test18_A38_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd_165000' #'AI4K_TEST_Denoise_A02_265000' | AI4K_test_A01b_145000 N_in = 5 # load test set if test_set == 'AI4K_test': #test_dataset_folder = '/data1/yhliu/AI4K/Corrected_TestA_Contest2_001_ResNet_alpha_beta_gaussian_65000/' #'/data1/yhliu/AI4K/testA_LR_png/' test_dataset_folder = '/home/yhliu/AI4K/contest2/testA_LR_png/' flip_test = False #False #model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' #model_path = '../experiments/002_EDVR_EDVRwoTSAIni_lr4e-4_600k_REDS_LrCAR4S_fixTSA50k_new/models/latest_G.pth' #model_path = '../experiments/A02_predenoise/models/415000_G.pth' model_path = '../experiments/A38_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd/models/165000_G.pth' color_model_path = '/home/yhliu/BasicSR/experiments/35_ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW_re_100k/models/220000_G.pth' predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #model = my_EDVR_arch.MYEDVR_FusionDenoise(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in, deconv=False) color_model = SRResNet_arch.ResNet_alpha_beta_multi_in( structure='ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW') #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(test_name) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) model = nn.DataParallel(model) #### set up the models load_net = torch.load(color_model_path) load_net_clean = OrderedDict() # add prefix 'color_net.' for k, v in load_net.items(): k = 'color_net.' + k load_net_clean[k] = v color_model.load_state_dict(load_net_clean, strict=True) color_model.eval() color_model = color_model.to(device) color_model = nn.DataParallel(color_model) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) #print(subfolder_l) #print(subfolder_GT_l) #exit() # for each subfolder for subfolder in subfolder_l: subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) #print(img_path_l) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).cpu() print(imgs_in.size()) if flip_test: imgs_in = util.single_forward(color_model, imgs_in) output = util.flipx4_forward(model, imgs_in) else: start_time = time.time() imgs_in = util.single_forward(color_model, imgs_in) output = util.single_forward(model, imgs_in) end_time = time.time() print('Forward One image:', end_time - start_time) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) logger.info('{:3d} - {:25}'.format(img_idx + 1, img_name)) logger.info('################ Tidy Outputs ################') logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test))
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, help='Path to option YAML file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() opt = option.parse(args.opt, is_train=True) #### distributed training settings if args.launcher == 'none': # disabled distributed training opt['dist'] = False rank = -1 print('Disabled distributed training.') else: opt['dist'] = True init_dist() world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() #### loading resume state if exists if opt['path'].get('resume_state', None): # distributed resuming: all load into default GPU device_id = torch.cuda.current_device() resume_state = torch.load( opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id)) option.check_resume(opt, resume_state['iter']) # check resume options else: resume_state = None #### mkdir and loggers if rank <= 0: # normal training (rank -1) OR distributed training (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_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') 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') # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) #### random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) if rank <= 0: logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True ''' train_loader = loader.get_loader(opt, train=True) val_loader = loader.get_loader(opt, train=False) # train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size'])) total_iters = int(opt['train']['niter']) total_epochs = int(opt['train']['epochs']) ''' #### 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_set = loader.get_dataset(opt, train=True) 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, world_size, 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 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, scale=opt['scale']) # val_set = loader.get_dataset(opt, train=False) val_loader = create_dataloader(val_set, dataset_opt, opt, None) if 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 #### create model model = create_model(opt) center_idx = opt['datasets']['train']['N_frames'] // 2 #### 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): train_data['GT'] = train_data['GT'][:, center_idx] current_step += 1 if current_step > total_iters: break #### update learning rate model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter']) #### training if opt['network_G']['which_model_G'] == 'TOF': # Bicubic upsample to match the size LQs = train_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B * T, C, H, W) Bic_LQs = F.interpolate(LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) train_data['LQs'] = Bic_LQs.reshape(B, T, C, H * opt['scale'], W * opt['scale']) 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 = '[epoch:{:3d}, iter:{:8,d}, lr:('.format( epoch, current_step) for v in model.get_current_learning_rate(): message += '{:.3e},'.format(v) message += ')] ' for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: if rank <= 0: tb_logger.add_scalar(k, v, current_step) if rank <= 0: 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 # does not support multi-GPU validation pbar = util.ProgressBar(len(val_loader)) avg_psnr = 0. idx = 0 for val_data in val_loader: idx += 1 img_name = os.path.splitext( os.path.basename(val_data['LQ_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() sr_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # Save SR images for reference save_img_path = os.path.join( img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) util.save_img(sr_img, save_img_path) # calculate PSNR sr_img, gt_img = util.crop_border([sr_img, gt_img], opt['scale']) avg_psnr += util.calculate_psnr(sr_img, gt_img) pbar.update('Test {}'.format(img_name)) avg_psnr = avg_psnr / idx # log logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) else: # video restoration validation if opt['dist']: # multi-GPU testing psnr_rlt = {} # with border and center frames if rank == 0: pbar = util.ProgressBar(len(val_set)) for idx in range(rank, len(val_set), world_size): val_data = val_set[idx] val_data['LQs'].unsqueeze_(0) val_data['GT'].unsqueeze_(0) folder = val_data['folder'] idx_d, max_idx = val_data['idx'].split('/') idx_d, max_idx = int(idx_d), int(max_idx) name = '{}/{:08d}'.format(folder, idx_d) train_folder = os.path.join( '../results', opt['name'], name) if not os.path.isdir(train_folder): os.makedirs(train_folder, exist_ok=True) if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = torch.zeros( max_idx, dtype=torch.float32, device='cuda') ''' folder = val_data['folder'] idx_d, max_idx = val_data['idx'].split('/') idx_d, max_idx = int(idx_d), int(max_idx) if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = torch.zeros(max_idx, dtype=torch.float32, device='cuda') # tmp = torch.zeros(max_idx, dtype=torch.float32, device='cuda') ''' if max_idx < 80 or (idx_d < max_idx / 2 and max_idx >= 80): if opt['network_G']['which_model_G'] == 'TOF': # Bicubic upsample to match the size LQs = val_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B * T, C, H, W) Bic_LQs = F.interpolate( LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) val_data['LQs'] = Bic_LQs.reshape( B, T, C, H * opt['scale'], W * opt['scale']) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() rlt_img = util.tensor2img( visuals['rlt'], mode='rgb') # uint8, RGB gt_img = util.tensor2img( visuals['GT'], mode='rgb') # uint8, RGB imageio.imwrite( os.path.join(train_folder, 'hr.png'), gt_img) imageio.imwrite( os.path.join(train_folder, 'sr.png'), rlt_img) # calculate PSNR psnr_rlt[folder][idx_d] = util.calculate_psnr( rlt_img, gt_img) if rank == 0: for _ in range(world_size): pbar.update('Test {} - {}/{}'.format( folder, idx_d, max_idx)) # # collect data for _, v in psnr_rlt.items(): dist.reduce(v, 0) dist.barrier() if rank == 0: psnr_rlt_avg = {} psnr_total_avg = 0. for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = torch.mean(v).cpu().item() psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) if opt['use_tb_logger'] and 'debug' not in opt[ 'name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) else: pbar = util.ProgressBar(len(val_loader)) psnr_rlt = {} # with border and center frames psnr_rlt_avg = {} psnr_total_avg = 0. for val_data in val_loader: folder = val_data['folder'][0] idx_d = val_data['idx'][0] # border = val_data['border'].item() name = '{}/{:08d}'.format(folder, idx_d) train_folder = os.path.join( '../results', opt['name'], name) if not os.path.isdir(train_folder): os.makedirs(train_folder, exist_ok=True) if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = [] video_length = val_data['LQs'].size(1) print(video_length) print(val_data['LQs'].shape, val_data['GT'].shape) for i in range(video_length): val_seg = {} select_idx = data_util.index_generation( i, video_length, opt['datasets']['train']['N_frames']) val_seg['LQs'] = val_data['LQs'][:, select_idx] val_seg['GT'] = val_data['GT'][:, i] model.feed_data(val_seg) model.test() visuals = model.get_current_visuals() rlt_img = util.tensor2img( visuals['rlt'], mode='rgb') # uint8, RGB gt_img = util.tensor2img( visuals['GT'], mode='rgb') # uint8, RGB imageio.imwrite( os.path.join(train_folder, 'hr.png'), gt_img) imageio.imwrite( os.path.join(train_folder, 'sr.png'), rlt_img) # calculate PSNR psnr, _ = utility.calc_psnr(rlt_img, gt_img) pbar.update('Test {} - {}'.format(folder, idx_d)) for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) #### save models and training states if current_step % opt['logger']['save_checkpoint_freq'] == 0: if rank <= 0: logger.info('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) if rank <= 0: logger.info('Saving the final model.') model.save('latest') logger.info('End of training.') tb_logger.close()
def main(): #################### # arguments parser # #################### # [format] dataset(vid4, REDS4) N(number of frames) parser = argparse.ArgumentParser() parser.add_argument('dataset') parser.add_argument('n_frames') parser.add_argument('stage') args = parser.parse_args() data_mode = str(args.dataset) N_in = int(args.n_frames) stage = int(args.stage) #if args.command == 'start': # start(int(args.params[0])) #elif args.command == 'stop': # stop(args.params[0], int(args.params[1])) #elif args.command == 'stop_all': # stop_all(args.params[0]) ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' #data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). #stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' else: raise ValueError('Vid4 does not support stage 2.') elif data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError predeblur, HR_in = False, False back_RBs = 40 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 #### dataset if data_mode == 'Vid4': N_model_default = 7 test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' else: N_model_default = 5 if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True data_mode_t = copy.deepcopy(data_mode) if stage == 1 and data_mode_t != 'Vid4': data_mode = 'REDS-EDVR_REDS_SR_L_flipx4' save_folder = '../results/{}'.format(data_mode) data_mode = copy.deepcopy(data_mode_t) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models print([a for a in dir(model) if not callable(getattr(model, a))]) # not a.startswith('__') and #model.load_state_dict(torch.load(model_path), strict=True) raw_model.load_state_dict(torch.load(model_path), strict=True) # model.load_state_dict(torch.load(model_path), strict=True) #### change model so it can work with less input model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA #ResidualBlock_noBN_f = functools.partial(arch_util.ResidualBlock_noBN, nf=nf) #### extract features (for each frame) if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur #Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 #nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 #nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align #PCD_Align(nf=nf, groups=groups) ######## Resize TSA model.tsa_fusion.center = model.center # temporal attention (before fusion conv) model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 # fusion conv: using 1x1 to save parameters and computation #print(raw_model.tsa_fusion.fea_fusion.weight.shape) #print(raw_model.tsa_fusion.fea_fusion.weight.shape) #print(raw_model.tsa_fusion.fea_fusion.weight[127][639].shape) #print("MAIN SHAPE(FEA): ", raw_model.tsa_fusion.fea_fusion.weight.shape) model.tsa_fusion.fea_fusion = copy.deepcopy( raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy( torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) #[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) #model.tsa_fusion.fea_fusion.bias = raw_model.tsa_fusion.fea_fusion.bias # spatial attention (after fusion conv) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy( torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) #[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) #model.tsa_fusion.sAtt_1.bias = raw_model.tsa_fusion.sAtt_1.bias #print(N_in * 128) #print(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :].shape) print("MODEL TSA SHAPE: ", model.tsa_fusion.fea_fusion.weight.shape) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu #if model.w_TSA: # model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #TSA_Fusion(nf=nf, nframes=nframes, center=self.center) #else: # model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) # print(self.tsa_fusion) #### reconstruction model.recon_trunk = raw_model.recon_trunk # arch_util.make_layer(ResidualBlock_noBN_f, back_RBs) #### upsampling model.upconv1 = raw_model.upconv1 #nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True) model.upconv2 = raw_model.upconv2 #nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True) model.pixel_shuffle = raw_model.pixel_shuffle # nn.PixelShuffle(2) model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last #### activation function model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_ssim, avg_ssim_border, avg_ssim_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] if data_mode == "blur": select_idx = data_util.glarefree_index_generation( img_idx, max_idx, N_in, padding=padding) else: select_idx = data_util.index_generation( img_idx, max_idx, N_in, padding=padding) # HERE GOTCHA print("SELECT IDX: ", select_idx) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) # check this output = util.single_forward(model, imgs_in) # error here 1 output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate SSIM output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) output, GT = util.crop_border([output, GT], crop_border) crt_ssim = util.calculate_ssim(output * 255, GT * 255) logger.info('{:3d} - {:25} \tSSIM: {:.6f} dB'.format( img_idx + 1, img_name, crt_ssim)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_ssim_center += crt_ssim N_center += 1 else: # border frames avg_ssim_border += crt_ssim N_border += 1 avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border) avg_ssim_center = avg_ssim_center / N_center avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border avg_ssim_l.append(avg_ssim) avg_ssim_center_l.append(avg_ssim_center) avg_ssim_border_l.append(avg_ssim_border) logger.info('Folder {} - Average SSIM: {:.6f} dB for {} frames; ' 'Center SSIM: {:.6f} dB for {} frames; ' 'Border SSIM: {:.6f} dB for {} frames.'.format( subfolder_name, avg_ssim, (N_center + N_border), avg_ssim_center, N_center, avg_ssim_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, ssim, ssim_center, ssim_border in zip( subfolder_name_l, avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l): logger.info('Folder {} - Average SSIM: {:.6f} dB. ' 'Center SSIM: {:.6f} dB. ' 'Border SSIM: {:.6f} dB.'.format(subfolder_name, ssim, ssim_center, ssim_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average SSIM: {:.6f} dB for {} clips. ' 'Center SSIM: {:.6f} dB. Border SSIM: {:.6f} dB.'.format( sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l), sum(avg_ssim_center_l) / len(avg_ssim_center_l), sum(avg_ssim_border_l) / len(avg_ssim_border_l)))
def main(): ################# # configurations ################# device = torch.device('cuda') # os.environ['CUDA_VISIBLE_DEVICES'] = '0' prog = argparse.ArgumentParser() prog.add_argument('--dataset_mode', '-m', type=str, default='Vid4+REDS', help='data_mode') prog.add_argument('--degradation_mode', '-d', type=str, default='impulse', choices=('impulse', 'bicubic', 'preset'), help='path to image output directory.') prog.add_argument('--sigma_x', '-sx', type=float, default=1, help='sigma_x') prog.add_argument('--sigma_y', '-sy', type=float, default=0, help='sigma_y') prog.add_argument('--theta', '-t', type=float, default=0, help='theta') prog.add_argument('--scale', '-sc', type=int, default=2, choices=(2, 4), help='scale factor') args = prog.parse_args() data_modes = args.dataset_mode degradation_mode = args.degradation_mode # impulse | bicubic sig_x, sig_y, the = args.sigma_x, args.sigma_y, args.theta * math.pi / 180 if sig_y == 0: sig_y = sig_x scale = args.scale kernel_size = 21 N_frames = 7 padding = 'new_info' data_mode_l = data_modes.split('+') for i in range(len(data_mode_l)): data_mode = data_mode_l[i] #### dataset if data_mode == 'Vid4': kernel_folder = '../pretrained_models/Mixed/Vid4.npy' dataset_folder = '../dataset/Vid4' elif data_mode == 'REDS': kernel_folder = '../pretrained_models/Mixed/REDS.npy' dataset_folder = '../dataset/REDS/train' elif data_mode == 'Vimeo': if degradation_mode == 'preset': raise NotImplementedError( 'We do not support preset mode in Vimeo dataset') dataset_folder = '../dataset/vimeo_septuplet' else: raise NotImplementedError() save_folder_name = 'preset' if degradation_mode == 'preset' else degradation_mode + '_' + str( '{:.1f}'.format(sig_x)) + '_' + str('{:.1f}'.format( sig_y)) + '_' + str('{:.1f}'.format(args.theta)) save_folder = osp.join(dataset_folder, 'LR_' + save_folder_name, 'X2') if not osp.exists(save_folder): os.makedirs(save_folder) save_folder2 = osp.join(dataset_folder, 'LR_' + save_folder_name, 'X4') if not osp.exists(save_folder2): os.makedirs(save_folder2) if scale == 4: save_folder3 = osp.join(dataset_folder, 'LR_' + save_folder_name, 'X16') if not osp.exists(save_folder3): os.makedirs(save_folder3) if data_mode == 'Vimeo': GT_dataset_folder = osp.join(dataset_folder, 'sequences') meta = osp.join(dataset_folder, 'sep_testlist.txt') with open(meta, 'r') as f: seq_list = sorted(f.read().splitlines()) subfolder_GT_l = [ osp.join(GT_dataset_folder, seq_ind) for seq_ind in seq_list ] else: GT_dataset_folder = osp.join(dataset_folder, 'HR') subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) if data_mode == 'REDS': subfolder_GT_l = [ k for k in subfolder_GT_l if k.find('000') >= 0 or k.find('011') >= 0 or k.find('015') >= 0 or k.find('020') >= 0 ] sig_x, sig_y, the = float(sig_x), float(sig_y), float(the) for subfolder_GT in subfolder_GT_l: print(subfolder_GT) gen_kwargs = set_kernel_params(sigma_x=sig_x, sigma_y=sig_y, theta=the) if degradation_mode == 'impulse' or degradation_mode == 'preset': kernel_gen = rkg.Degradation(kernel_size, 2, **gen_kwargs) if degradation_mode == 'preset': kernel_preset = np.load(kernel_folder) else: kernel_gen = oldkg.Degradation(kernel_size, 2, type=0.7, **gen_kwargs) if data_mode == 'Vimeo': sub1 = osp.basename(osp.dirname(subfolder_GT)) sub2 = osp.basename(subfolder_GT) subfolder_name = osp.join(sub1, sub2) else: subfolder_name = osp.basename(subfolder_GT) save_subfolder = osp.join(save_folder, subfolder_name) if not osp.exists(save_subfolder): os.makedirs(save_subfolder) save_subfolder2 = osp.join(save_folder2, subfolder_name) if not osp.exists(save_subfolder2): os.makedirs(save_subfolder2) if scale == 4: save_subfolder3 = osp.join(save_folder3, subfolder_name) if not osp.exists(save_subfolder3): os.makedirs(save_subfolder3) img_GT_path_l = sorted(glob.glob(osp.join(subfolder_GT, '*'))) seq_length = len(img_GT_path_l) imgs_GT = data_util.read_img_seq(subfolder_GT) # T C H W if degradation_mode == 'preset': for index in range(seq_length): save_subsubfolder = osp.join( save_subfolder, osp.splitext(osp.basename(img_GT_path_l[index]))[0]) save_subsubfolder2 = osp.join( save_subfolder2, osp.splitext(osp.basename(img_GT_path_l[index]))[0]) if not osp.exists(save_subsubfolder): os.mkdir(save_subsubfolder) if not osp.exists(save_subsubfolder2): os.mkdir(save_subsubfolder2) if scale == 4 and not osp.exists(save_subsubfolder3): os.mkdir(save_subsubfolder3) kernel_gen.set_kernel_directly(kernel_preset[index]) imgs_HR = imgs_GT[data_util.index_generation( index, seq_length, N_frames, padding)] imgs_LR = kernel_gen.apply(imgs_HR) imgs_LR = imgs_LR.mul(255).clamp(0, 255).round().div(255) imgs_LR_np = imgs_LR.permute(0, 2, 3, 1).cpu().numpy() imgs_LR_np = (imgs_LR_np * 255).astype('uint8') for i, img_LR in enumerate(imgs_LR_np): imageio.imwrite( osp.join(save_subsubfolder, 'img{}.png'.format(i)), img_LR) imgs_SuperLR = kernel_gen.apply(imgs_LR) imgs_SuperLR = imgs_SuperLR.mul(255).clamp( 0, 255).round().div(255) imgs_SuperLR_np = imgs_SuperLR.permute(0, 2, 3, 1).cpu().numpy() imgs_SuperLR_np = (imgs_SuperLR_np * 255).astype('uint8') for i, img_SuperLR in enumerate(imgs_SuperLR_np): imageio.imwrite( osp.join(save_subsubfolder2, 'img{}.png'.format(i)), img_SuperLR) if scale == 4: imgs_SuperLR = kernel_gen.apply(imgs_LR) imgs_SuperLR = imgs_SuperLR.mul(255).clamp( 0, 255).round().div(255) imgs_SuperLR = kernel_gen.apply(imgs_LR) imgs_SuperLR = imgs_SuperLR.mul(255).clamp( 0, 255).round().div(255) imgs_SuperLR_np = imgs_SuperLR.permute( 0, 2, 3, 1).cpu().numpy() imgs_SuperLR_np = (imgs_SuperLR_np * 255).astype('uint8') for i, img_SuperLR in enumerate(imgs_SuperLR_np): imageio.imwrite( osp.join(save_subsubfolder3, 'img{}.png'.format(i)), img_SuperLR) else: count = 0 imgs_GT_l = imgs_GT.split(32) for img_batch in imgs_GT_l: img_lr_batch = kernel_gen.apply(img_batch) img_lr_batch = img_lr_batch.permute(0, 2, 3, 1).cpu().numpy() img_lr_batch = (img_lr_batch.clip(0, 1) * 255).round() img_lr_batch = img_lr_batch.astype('uint8') count_temp = count for img_lr in img_lr_batch: filename = osp.basename(img_GT_path_l[count]) imageio.imwrite(osp.join(save_subfolder, filename), img_lr) count += 1 img_lr_batch = img_lr_batch.astype('float32') / 255 img_lr_batch = torch.from_numpy(img_lr_batch).permute( 0, 3, 1, 2) img_superlr_batch = kernel_gen.apply(img_lr_batch) img_superlr_batch = img_superlr_batch.permute( 0, 2, 3, 1).cpu().numpy() img_superlr_batch = (img_superlr_batch.clip(0, 1) * 255).round() img_superlr_batch = img_superlr_batch.astype('uint8') count = count_temp for img_superlr in img_superlr_batch: filename = osp.basename(img_GT_path_l[count]) imageio.imwrite(osp.join(save_subfolder2, filename), img_superlr) count += 1 if scale == 4: img_superlr_batch = img_superlr_batch.astype( 'float32') / 255 img_superlr_batch = torch.from_numpy( img_superlr_batch).permute(0, 3, 1, 2) img_superlr_batch = kernel_gen.apply(img_superlr_batch) img_superlr_batch = img_superlr_batch.permute( 0, 2, 3, 1).cpu().numpy() img_superlr_batch = (img_superlr_batch.clip(0, 1) * 255).round() img_superlr_batch = img_superlr_batch.astype( 'float32') / 255 img_superlr_batch = torch.from_numpy( img_superlr_batch).permute(0, 3, 1, 2) img_superlr_batch = kernel_gen.apply(img_superlr_batch) img_superlr_batch = img_superlr_batch.permute( 0, 2, 3, 1).cpu().numpy() img_superlr_batch = (img_superlr_batch.clip(0, 1) * 255).round() img_superlr_batch = img_superlr_batch.astype('uint8') count = count_temp for img_superlr in img_superlr_batch: filename = osp.basename(img_GT_path_l[count]) imageio.imwrite( osp.join(save_subfolder2, filename), img_superlr) count += 1
def __getitem__(self, index): # path_LQ = self.data_info['path_LQ'][index] # path_GT = self.data_info['path_GT'][index] folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] if self.cache_data: select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'],padding=self.opt['padding']) #select_idx = util.index_generation_process_screen_change_withlog_fixbug(folder, self.frame_notation, idx, max_idx, self.opt['N_frames'], # padding=self.opt['padding']) # select_idx, log1, log2, nota = util.index_generation_process_screen_change_withlog_fixbug( # folder, self.frame_notation, idx, max_idx, self.opt['N_frames'], padding=self.opt['padding'], # enable=1) # if not log1 == None: # print('screen change') # print(nota) # print(log1) # print(log2) imgs_LQ = self.imgs_LQ[folder].index_select(0, torch.LongTensor(select_idx)) img_GT = self.imgs_GT[folder][idx] else: select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'],padding=self.opt['padding']) imgs_LQ_path = [] for v in select_idx: imgs_LQ_path.append(self.imgs_LQ[folder][v]) img_GT_path = self.imgs_GT[folder][idx] # read gt img_GT = util.read_img(None, img_GT_path) img_LQ_l = [] for v in imgs_LQ_path: img_LQ = util.read_img(None, v) img_LQ_l.append(img_LQ) #color if self.opt['color'] == 'YUV': img_LQ_l = [util.bgr2ycbcr(v,only_y=False) for v in img_LQ_l] img_GT = util.bgr2ycbcr(img_GT,only_y=False) # stack LQ images to NHWC, N is the frame number img_LQs = np.stack(img_LQ_l, axis=0) # BGR to RGB img_GT = img_GT[:, :, [2, 1, 0]] # HWC img_LQs = img_LQs[:, :, :, [2, 1, 0]] # HWC to CHW numpy to tensor # if YUV -> VUY img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float() imgs_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQs,(0, 3, 1, 2)))).float() return { 'LQs': imgs_LQ, 'GT': img_GT, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border }
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '1' data_mode = 'ai4khdr_test' flip_test = False ############################################################################ #### model ################# if data_mode == 'ai4khdr_test': model_path = '../experiments/002_EDVR_lr4e-4_600k_AI4KHDR/models/4000_G.pth' else: raise NotImplementedError N_in = 5 front_RBs = 5 back_RBs = 10 predeblur, HR_in = False, False model = EDVR_arch.EDVR(64, N_in, 8, front_RBs, back_RBs, predeblur=predeblur, HR_in=HR_in) ############################################################################ #### dataset ################# if data_mode == 'ai4khdr_test': test_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/test/540p_frames' else: raise NotImplementedError ############################################################################ #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'ai4khdr_test': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}_{}'.format(data_mode, util.get_timestamp()) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=False) model.eval() model = model.to(device) subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) # for each subfolder for subfolder in subfolder_l: subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) logger.info('Folder {}'.format(subfolder_name)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test))
def __call__(self, input_path, output_path): with torch.no_grad(): # read LQ images video_cap = cv2.VideoCapture(input_path) video_fps = video_cap.get(cv2.CAP_PROP_FPS) imgs_LQ = list() while video_cap.isOpened(): ret, frame = video_cap.read() if ret: imgs_LQ.append(torch.from_numpy( frame[:, :, [2, 1, 0]])) # BGR -> RGB else: break video_cap.release() imgs_LQ = torch.stack(imgs_LQ, dim=0) # imgs_LQ: Tensor[T,H,W,C] imgs_LQ = imgs_LQ.permute( 0, 3, 1, 2).contiguous().float() / 255.0 # imgs_LQ: Tensor[T,C,H,W] _, _, height, width = imgs_LQ.shape imgs_LQ = imgs_LQ[:, :, :height - (height % 16), :width - (width % 16)] max_idx = imgs_LQ.shape[0] input_batch = list() output_tensor = list() # process each image for img_idx in tqdm.tqdm(range(max_idx)): select_idx = data_util.index_generation( img_idx, max_idx, self.network_conf['nframes'], padding=self.padding) # imgs_in: Tensor[nframes,C,H,W] imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)) input_batch.append(imgs_in) if (img_idx + 1) % self.batch == 0 and len(input_batch) != 0: input_tensor = torch.stack( input_batch, dim=0) # Tensor[B,nframes,C,H,W] output = self.single_inference(input_tensor) output_tensor.append(output) input_batch = list() del input_tensor if len(input_batch) != 0: input_tensor = torch.stack(input_batch, dim=0) output = self.single_inference(input_tensor) output_tensor.append(output) del input_tensor del input_batch del imgs_LQ output_tensor = torch.cat(output_tensor, dim=0) # output_tensor: Tensor[T,C,H,W] # write video output = output_tensor.permute(0, 2, 3, 1) # output: Tensor[T,H,W,C] video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), video_fps, (output.shape[2], output.shape[1])) for i in range(output.shape[0]): video_writer.write(output[i][:, :, [2, 1, 0]].numpy()) video_writer.release()
def main(): ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '1' data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: #model_path = '../experiments/pretrained_models/EDVR_REDS_SR_M.pth' model_path = '../experiments/002_EDVR_lr4e-4_600k_AI4KHDR/models/4000_G.pth' else: raise ValueError('Vid4 does not support stage 2.') elif data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError if data_mode == 'Vid4': N_in = 5 # use N_in images to restore one HR image else: N_in = 5 predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 model = EDVR_arch.EDVR(64, 5, 8, 5, 10, predeblur=predeblur, HR_in=HR_in) #### dataset if data_mode == 'Vid4': test_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/valid/540p_frames' GT_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/valid/4k_frames' #test_dataset_folder = '../datasets/Vid4/BIx4' #GT_dataset_folder = '../datasets/Vid4/GT' else: if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=False) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) #logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip( subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main(): ################# # configurations ################# parser = argparse.ArgumentParser() parser.add_argument("--input_path", type=str, required=True) parser.add_argument("--gt_path", type=str, required=True) parser.add_argument("--output_path", type=str, required=True) parser.add_argument("--model_path", type=str, required=True) parser.add_argument("--gpu_id", type=str, required=True) parser.add_argument("--screen_notation", type=str, required=True) parser.add_argument('--opt', type=str, required=True, help='Path to option YAML file.') args = parser.parse_args() opt = option.parse(args.opt, is_train=False) PAD = 32 total_run_time = AverageMeter() print("GPU ", torch.cuda.device_count()) device = torch.device('cuda') # os.environ['CUDA_VISIBLE_DEVICES'] = '0' os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id) data_mode = 'sharp_bicubic' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False # Input_folder = "/DATA7_DB7/data/4khdr/data/Dataset/train_sharp_bicubic" # GT_folder = "/DATA7_DB7/data/4khdr/data/Dataset/train_4k" # Result_folder = "/DATA7_DB7/data/4khdr/data/Results" Input_folder = args.input_path GT_folder = args.gt_path Result_folder = args.output_path Model_path = args.model_path # create results folder if not os.path.exists(Result_folder): os.makedirs(Result_folder, exist_ok=True) ############################################################################ #### model # if data_mode == 'Vid4': # if stage == 1: # model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' # else: # raise ValueError('Vid4 does not support stage 2.') # elif data_mode == 'sharp_bicubic': # if stage == 1: # # model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' # else: # model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' # elif data_mode == 'blur_bicubic': # if stage == 1: # model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' # else: # model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' # elif data_mode == 'blur': # if stage == 1: # model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' # else: # model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' # elif data_mode == 'blur_comp': # if stage == 1: # model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' # else: # model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' # else: # raise NotImplementedError model_path = Model_path if data_mode == 'Vid4': N_in = 7 # use N_in images to restore one HR image else: N_in = 5 predeblur, HR_in = False, False back_RBs = 40 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 model = EDVR_arch.EDVR(nf=opt['network_G']['nf'], nframes=opt['network_G']['nframes'], groups=opt['network_G']['groups'], front_RBs=opt['network_G']['front_RBs'], back_RBs=opt['network_G']['back_RBs'], predeblur=opt['network_G']['predeblur'], HR_in=opt['network_G']['HR_in'], w_TSA=opt['network_G']['w_TSA']) # model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### dataset if data_mode == 'Vid4': test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' else: if stage == 1: # test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) # test_dataset_folder = '/DATA/wangshen_data/REDS/val_sharp_bicubic/X4' test_dataset_folder = Input_folder else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') # GT_dataset_folder = '../datasets/REDS4/GT' # GT_dataset_folder = '/DATA/wangshen_data/REDS/val_sharp' GT_dataset_folder = GT_folder #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True # save_folder = '../results/{}'.format(data_mode) # save_folder = '/DATA/wangshen_data/REDS/results/{}'.format(data_mode) save_folder = os.path.join(Result_folder, data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder # for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): end = time.time() for subfolder in subfolder_l: input_subfolder = os.path.split(subfolder)[1] subfolder_GT = os.path.join(GT_dataset_folder, input_subfolder) if not os.path.exists(subfolder_GT): continue print("Evaluate Folders: ", input_subfolder) subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) # Num x 3 x H x W img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to( device) # 960 x 540 # here we split the input images 960x540 into 9 320x180 patch gtWidth = 3840 gtHeight = 2160 intWidth_ori = imgs_in.shape[4] # 960 intHeight_ori = imgs_in.shape[3] # 540 split_lengthY = 180 split_lengthX = 320 scale = 4 intPaddingRight_ = int(float(intWidth_ori) / split_lengthX + 1) * split_lengthX - intWidth_ori intPaddingBottom_ = int(float(intHeight_ori) / split_lengthY + 1) * split_lengthY - intHeight_ori intPaddingRight_ = 0 if intPaddingRight_ == split_lengthX else intPaddingRight_ intPaddingBottom_ = 0 if intPaddingBottom_ == split_lengthY else intPaddingBottom_ pader0 = torch.nn.ReplicationPad2d( [0, intPaddingRight_, 0, intPaddingBottom_]) print("Init pad right/bottom " + str(intPaddingRight_) + " / " + str(intPaddingBottom_)) intPaddingRight = PAD # 32# 64# 128# 256 intPaddingLeft = PAD # 32#64 #128# 256 intPaddingTop = PAD # 32#64 #128#256 intPaddingBottom = PAD # 32#64 # 128# 256 pader = torch.nn.ReplicationPad2d([ intPaddingLeft, intPaddingRight, intPaddingTop, intPaddingBottom ]) imgs_in = torch.squeeze(imgs_in, 0) # N C H W imgs_in = pader0(imgs_in) # N C 540 960 imgs_in = pader(imgs_in) # N C 604 1024 assert (split_lengthY == int(split_lengthY) and split_lengthX == int(split_lengthX)) split_lengthY = int(split_lengthY) split_lengthX = int(split_lengthX) split_numY = int(float(intHeight_ori) / split_lengthY) split_numX = int(float(intWidth_ori) / split_lengthX) splitsY = range(0, split_numY) splitsX = range(0, split_numX) intWidth = split_lengthX intWidth_pad = intWidth + intPaddingLeft + intPaddingRight intHeight = split_lengthY intHeight_pad = intHeight + intPaddingTop + intPaddingBottom # print("split " + str(split_numY) + ' , ' + str(split_numX)) y_all = np.zeros((gtHeight, gtWidth, 3), dtype="float32") # HWC for split_j, split_i in itertools.product(splitsY, splitsX): # print(str(split_j) + ", \t " + str(split_i)) X0 = imgs_in[:, :, split_j * split_lengthY:(split_j + 1) * split_lengthY + intPaddingBottom + intPaddingTop, split_i * split_lengthX:(split_i + 1) * split_lengthX + intPaddingRight + intPaddingLeft] # y_ = torch.FloatTensor() X0 = torch.unsqueeze(X0, 0) # N C H W -> 1 N C H W if flip_test: output = util.flipx4_forward(model, X0) else: output = util.single_forward(model, X0) output_depadded = output[0, :, intPaddingTop * scale:(intPaddingTop + intHeight) * scale, intPaddingLeft * scale:(intPaddingLeft + intWidth) * scale] output_depadded = output_depadded.squeeze(0) output = util.tensor2img(output_depadded) y_all[split_j * split_lengthY * scale:(split_j + 1) * split_lengthY * scale, split_i * split_lengthX * scale:(split_i + 1) * split_lengthX * scale, :] = \ np.round(output).astype(np.uint8) # plt.figure(0) # plt.title("pic") # plt.imshow(y_all) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), y_all) print("*****************current image process time \t " + str(time.time() - end) + "s ******************") total_run_time.update(time.time() - end, 1) # calculate PSNR y_all = y_all / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) y_all = data_util.bgr2ycbcr(y_all, only_y=True) y_all, GT = util.crop_border([y_all, GT], crop_border) crt_psnr = util.calculate_psnr(y_all * 255, GT * 255) logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format( img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip( subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main(): #################### # arguments parser # #################### # [format] dataset(vid4, REDS4) N(number of frames) # data_mode = str(args.dataset) # N_in = int(args.n_frames) # metrics = str(args.metrics) # output_format = str(args.output_format) ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' #data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). # STAGE Vid4 # Collecting results for Vid4 model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False predeblur, HR_in = False, False back_RBs = 40 N_model_default = 7 data_mode = 'Vid4' # vid4_dir_map = {"calendar": 0, "city": 1, "foliage": 2, "walk": 3} vid4_results = {"calendar": {}, "city": {}, "foliage": {}, "walk": {}} #vid4_results = 4 * [[]] for N_in in range(1, N_model_default + 1): raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' aposterior_GT_dataset_folder = '../datasets/Vid4/GT_7' crop_border = 0 border_frame = N_in // 2 # border frames when evaluate padding = 'new_info' save_imgs = False raw_model.load_state_dict(torch.load(model_path), strict=True) model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups) model.tsa_fusion.center = model.center model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu model.recon_trunk = raw_model.recon_trunk model.upconv1 = raw_model.upconv1 model.upconv2 = raw_model.upconv2 model.pixel_shuffle = raw_model.pixel_shuffle model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) #avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*"))) # for each subfolder for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) img_GT_a = [] for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))): img_GT_a.append(data_util.read_img(None, img_GT_a_path)) #avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel #if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) GT_a = np.copy(img_GT_a[img_idx]) GT_a = data_util.bgr2ycbcr(GT_a, only_y=True) output_a = copy.deepcopy(output) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) crt_ssim = util.calculate_ssim(output * 255, GT * 255) output_a, GT_a = util.crop_border([output_a, GT_a], crop_border) crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE t = vid4_results[subfolder_name].get(str(img_name)) if t != None: vid4_results[subfolder_name][img_name].add_psnr(crt_psnr) vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) else: vid4_results[subfolder_name].update({img_name: metrics_file(img_name)}) vid4_results[subfolder_name][img_name].add_psnr(crt_psnr) vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) ############################################################################ #### model #### writing vid4 results util.mkdirs('../results/calendar') util.mkdirs('../results/city') util.mkdirs('../results/foliage') util.mkdirs('../results/walk') save_folder = '../results/' for i, dir_name in enumerate(["calendar", "city", "foliage", "walk"]): save_subfolder = osp.join(save_folder, dir_name) for j, value in vid4_results[dir_name].items(): # cur_result = json.dumps(_) with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile: json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4) #json.dump(cur_result, outfile) #cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) ################################################################################### # STAGE REDS reds4_results = {"000": {}, "011": {}, "015": {}, "020": {}} data_mode = 'sharp_bicubic' N_model_default = 5 for N_in in range(1, N_model_default + 1): for stage in range(1,3): flip_test = False if data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError predeblur, HR_in = False, False back_RBs = 40 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True data_mode_t = copy.deepcopy(data_mode) if stage == 1 and data_mode_t != 'Vid4': data_mode = 'REDS-EDVR_REDS_SR_L_flipx4' save_folder = '../results/{}'.format(data_mode) data_mode = copy.deepcopy(data_mode_t) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) aposterior_GT_dataset_folder = '../datasets/REDS4/GT_5' crop_border = 0 border_frame = N_in // 2 # border frames when evaluate raw_model.load_state_dict(torch.load(model_path), strict=True) model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups) model.tsa_fusion.center = model.center model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu model.recon_trunk = raw_model.recon_trunk model.upconv1 = raw_model.upconv1 model.upconv2 = raw_model.upconv2 model.pixel_shuffle = raw_model.pixel_shuffle model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) #avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*"))) # for each subfolder for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) print("SAVE FOLDER::::::", save_folder) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) img_GT_a = [] for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))): img_GT_a.append(data_util.read_img(None, img_GT_a_path)) #avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs and stage == 1: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR if stage == 2: output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel #if data_mode == 'Vid4': # bgr2y, [0, 1] GT_a = np.copy(img_GT_a[img_idx]) output_a = copy.deepcopy(output) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) crt_ssim = util.calculate_ssim(output * 255, GT * 255) output_a, GT_a = util.crop_border([output_a, GT_a], crop_border) crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE t = reds4_results[subfolder_name].get(str(img_name)) if t != None: reds4_results[subfolder_name][img_name].add_psnr(crt_psnr) reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) else: reds4_results[subfolder_name].update({img_name: metrics_file(img_name)}) reds4_results[subfolder_name][img_name].add_psnr(crt_psnr) reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) ############################################################################ #### model #### writing reds4 results util.mkdirs('../results/000') util.mkdirs('../results/011') util.mkdirs('../results/015') util.mkdirs('../results/020') save_folder = '../results/' for i, dir_name in enumerate(["000", "011", "015", "020"]): # + save_subfolder = osp.join(save_folder, dir_name) for j, value in reds4_results[dir_name].items(): # cur_result = json.dumps(value.__dict__) with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile: json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4)
def val(model_name, current_step, arch='EDVR'): ################# # configurations ################# device = torch.device('cuda') #os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3,4' test_set = 'REDS4' # Vid4 | YouKu10 | REDS4 | AI4K_val data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic N_in = 5 # load test set if test_set == 'Vid4': test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' elif test_set == 'YouKu10': test_dataset_folder = '../datasets/YouKu10/LR' GT_dataset_folder = '../datasets/YouKu10/HR' elif test_set == 'YouKu_val': test_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_lr_bmp' GT_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_hr_bmp' elif test_set == 'REDS4': test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) GT_dataset_folder = '../datasets/REDS4/GT' elif test_set == 'AI4K_val': test_dataset_folder = '/data0/yhliu/AI4K/contest1/val1_LR_png/' GT_dataset_folder = '/data0/yhliu/AI4K/contest1/val1_HR_png/' elif test_set == 'AI4K_val_small': test_dataset_folder = '/home/yhliu/AI4K/contest1/val1_LR_png_small/' GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png_small/' flip_test = False #model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' model_path = os.path.join('../experiments/', model_name, 'models/{}_G.pth'.format(current_step)) predeblur, HR_in = False, False back_RBs = 10 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if arch == 'EDVR': model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) elif arch == 'MY_EDVR': model = my_EDVR_arch.MYEDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = False save_folder = '../validation/{}'.format(test_set) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.to(device) model = nn.DataParallel(model, device_ids=[0, 1, 2, 3]) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) #print(subfolder_l) #print(subfolder_GT_l) #exit() # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) #print(img_path_l) max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): #print(img_GT_path) img_GT_l.append(data_util.read_img(None, img_GT_path)) #print(img_GT_l[0].shape) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) #print(imgs_in.size()) if flip_test: output = util.flipx4_forward(model, imgs_in) else: output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel ''' if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) ''' output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) #logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr N_center += 1 else: # border frames avg_psnr_border += crt_psnr N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip( subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l))) return sum(avg_psnr_l) / len(avg_psnr_l)