def define_C(opt):
opt_net = opt['network_C']
which_model = opt_net['which_model_C']
if which_model == 'dfn':
netC = DynamicF.DFN_Color_correction()
elif 'ResNet' in which_model:
netC = SRResNet_arch.ResNet_alpha_beta_multi_in(which_model)
else:
raise NotImplementedError(
'Discriminator model [{:s}] not recognized'.format(which_model))
return netC
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():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '5'
#os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3,4'
test_set = 'AI4K_val' # Vid4 | YouKu10 | REDS4 | AI4K_val
test_name = 'A38_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd_165000'
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_bic':
test_dataset_folder = '/home/yhliu/AI4K/contest2/val2_LR_png_bic/'
GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png_bic/'
elif test_set == 'AI4K_testA':
test_dataset_folder = '/data0/yhliu/AI4K/val2_LR_png' #'/home/yhliu/AI4K/val2_LR_png/'
GT_dataset_folder = '/data0/yhliu/AI4K/val1_HR_png/'
flip_test = False
#model_path = '../experiments/pretrained_models/A01xxx/900000_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)
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, '*')))
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:
imgs_in = util.single_forward(color_model, imgs_in)
output = util.flipx4_forward(model, imgs_in)
else:
imgs_in = util.single_forward(color_model, imgs_in)
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)))
