def main(cfg):
video_name = cfg.video_name
upscale_factor = cfg.upscale_factor
use_gpu = cfg.gpu_mode
test_set = TestsetLoader('data/'+ video_name, upscale_factor)
test_loader = DataLoader(test_set, num_workers=1, batch_size=1, shuffle=False)
net = SOFVSR(upscale_factor=upscale_factor)
ckpt = torch.load('./log/SOFVSR_x' + str(upscale_factor) + '.pth')
net.load_state_dict(ckpt)
if use_gpu:
net.cuda()
for idx_iter, (LR_y_cube, SR_cb, SR_cr) in enumerate(test_loader):
LR_y_cube = Variable(LR_y_cube)
if use_gpu:
LR_y_cube = LR_y_cube.cuda()
SR_y = net(LR_y_cube)
SR_y = np.array(SR_y.data)
SR_y = SR_y[np.newaxis, :, :]
SR_ycbcr = np.concatenate((SR_y, SR_cb, SR_cr), axis=0).transpose(1,2,0)
SR_rgb = ycbcr2rgb(SR_ycbcr) * 255.0
SR_rgb = np.clip(SR_rgb, 0, 255)
SR_rgb = ToPILImage()(SR_rgb.astype(np.uint8))
if not os.path.exists('results/' + video_name):
os.mkdir('results/' + video_name)
SR_rgb.save('results/'+video_name+'/sr_'+ str(idx_iter+2).rjust(2,'0') + '.png')
def main(cfg):
video_name = cfg.video_name
upscale_factor = cfg.upscale_factor
use_gpu = cfg.gpu_mode
test_set = TestsetLoader('data/test/' + video_name, upscale_factor)
test_loader = DataLoader(test_set,
num_workers=1,
batch_size=1,
shuffle=False)
net = SOFVSR(upscale_factor=upscale_factor)
ckpt = torch.load('./log/SOFVSR_x' + str(upscale_factor) + '.pth')
net.load_state_dict(ckpt)
if use_gpu:
net.cuda()
for idx_iter, (LR_y_cube, SR_cb, SR_cr) in enumerate(test_loader):
LR_y_cube = Variable(LR_y_cube)
if use_gpu:
LR_y_cube = LR_y_cube.cuda()
if cfg.chop_forward:
# crop borders to ensure each patch can be divisible by 2
_, _, h, w = LR_y_cube.size()
h = int(h // 16) * 16
w = int(w // 16) * 16
LR_y_cube = LR_y_cube[:, :, :h, :w]
SR_cb = SR_cb[:, :h * upscale_factor, :w * upscale_factor]
SR_cr = SR_cr[:, :h * upscale_factor, :w * upscale_factor]
SR_y = chop_forward(LR_y_cube, net, cfg.upscale_factor)
else:
SR_y = net(LR_y_cube)
SR_y = SR_y.cpu()
else:
SR_y = net(LR_y_cube)
SR_y = np.array(SR_y.data)
SR_y = SR_y[np.newaxis, :, :]
SR_ycbcr = np.concatenate((SR_y, SR_cb, SR_cr),
axis=0).transpose(1, 2, 0)
SR_rgb = ycbcr2rgb(SR_ycbcr) * 255.0
SR_rgb = np.clip(SR_rgb, 0, 255)
SR_rgb = ToPILImage()(SR_rgb.astype(np.uint8))
if not os.path.exists('results/' + video_name):
os.mkdir('results/' + video_name)
SR_rgb.save('results/' + video_name + '/sr_' +
str(idx_iter + 2).rjust(2, '0') + '.png')
def main():
'''
训练时时并行的，测试时也应当并行，不然会报告如下的错误：
Missing key(s) in state_dict: ...(如：conv1.weight)
'''
print('testing processing....')
#加载模型
test_model = VRCNN(opt.upscale_factor)
test_model = torch.nn.DataParallel(test_model,device_ids=gpus_list,output_device=gpus_list[1])
test_model = test_model.cuda(gpus_list[0])
print('---------- Networks architecture -------------')
print_network(test_model)
print('----------------------------------------------')
#加载预训练模型
model_name = os.path.join(opt.model_save_folder,opt.exp_name,opt.test_model)
print('model_name=',model_name)
if os.path.exists(model_name):
pretrained_dict=torch.load(model_name,map_location=lambda storage, loc: storage)
model_dict=test_model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
test_model.load_state_dict(model_dict)
print('Pre-trained SR model is loaded.')
if not os.path.exists(opt.pre_result):
os.mkdir(opt.pre_result)
with open(opt.train_log + '/psnr_ssim-xr-200.txt', 'a') as psnr_ssim:
with torch.no_grad():
ave_psnr = 0
ave_ssim = 0
single_ave_psnr = 0
single_ave_ssim = 0
numb = 2
valSet = ValidationsetLoader(opt.val_dataset_hr,opt.val_dataset_lr)
valLoader = DataLoader(dataset=valSet,batch_size=opt.test_val_batchSize,shuffle=False)
val_bar = tqdm(valLoader)
for data in val_bar:
test_model.eval()
# dual_net.eval()
batch_lr_y, label, SR_cb,SR_cr,idx,bicubic_restore = data
batch_lr_y,label = Variable(batch_lr_y).cuda(gpus_list[0]), Variable(label).cuda(gpus_list[0])
output = test_model(batch_lr_y)
SR_ycbcr = np.concatenate((np.array(output.squeeze(0).data.cpu()), SR_cb, SR_cr), axis=0).transpose(1,2,0)
SR_rgb = ycbcr2rgb(SR_ycbcr) * 255.0
SR_rgb = np.clip(SR_rgb, 0, 255)
SR_rgb = ToPILImage()(SR_rgb.astype(np.uint8))
#ToTensor() ---image(0-255)==>image(0-1), (H,W,C)==>(C,H,W)
SR_rgb = ToTensor()(SR_rgb)
#将给定的Tensor保存成image文件。如果给定的是mini-batch tensor，那就用make-grid做成雪碧图，再保存。与utils.make_grid()配套使用
if not os.path.exists(opt.pre_result+'/'+opt.exp_name):
os.mkdir(opt.pre_result+'/'+opt.exp_name)
utils.save_image(SR_rgb, opt.pre_result+'/' +opt.exp_name +'/' + 'my'+str(numb).rjust(3,'0')+'.png')
numb = numb + 1
psnr_value = psnr(np.array(torch.squeeze(label).data.cpu())*255,np.array(torch.squeeze(output).data.cpu())*255)
ave_psnr = ave_psnr + psnr_value
single_ave_psnr = single_ave_psnr + psnr_value
ssim_value = calculate_ssim(np.array(torch.squeeze(label).data.cpu())*255,np.array(torch.squeeze(output).data.cpu())*255)
ave_ssim = ave_ssim + ssim_value
single_ave_ssim = single_ave_ssim + ssim_value
val_bar.set_description('===>{}th video {}th frame, wsPSNR:{:.4f} dB,wsSSIM:{:.6f}'.format(idx // 98 + 1,idx % 98 + 1,psnr_value,ssim_value))
if idx == 293 or idx == 97 or idx == 195 or idx == 391:
print("===> {}th video Avg. wsPSNR: {:.4f} dB".format(idx // 98+1,single_ave_psnr / 98))
print("===> {}th video Avg. wsSSIM: {:.6f}".format(idx // 98+1,single_ave_ssim / 98))
psnr_ssim.write('===>{}th video avg wsPSNR:{:.4f} dB,wsSSIM:{:.6f}\n'.format(idx // 98+1,single_ave_psnr / 98,single_ave_ssim / 98))
single_ave_psnr = 0
single_ave_ssim = 0
print("===> All Avg. wsPSNR: {:.4f} dB".format(ave_psnr / len(valLoader)))
print("===> ALL Avg. wsSSIM: {:.6f}".format(ave_ssim / len(valLoader)))
psnr_ssim.write('===>all videos avg wsPSNR:{:.4f} dB,wsSSIM:{:.6f}\n'.format(ave_psnr / len(valLoader),ave_ssim / len(valLoader)))
print('testing finished!')
def main(cfg):
# model
net = SOFVSR(cfg, is_training=False)
ckpt = torch.load('./log/' + cfg.degradation + '_x' + str(cfg.scale) +
'.pth')
net.load_state_dict(ckpt)
if cfg.gpu_mode:
net.cuda()
with torch.no_grad():
video_list = os.listdir(cfg.testset_dir)
for idx_video in range(len(video_list)):
video_name = video_list[idx_video]
# dataloader
test_set = TestsetLoader(cfg, video_name)
test_loader = DataLoader(test_set,
num_workers=1,
batch_size=1,
shuffle=False)
for idx_iter, (LR_y_cube, SR_cb, SR_cr) in enumerate(test_loader):
# data
b, n_frames, h_lr, w_lr = LR_y_cube.size()
LR_y_cube = Variable(LR_y_cube)
LR_y_cube = LR_y_cube.view(b, -1, 1, h_lr, w_lr)
if cfg.gpu_mode:
LR_y_cube = LR_y_cube.cuda()
if cfg.chop_forward:
# crop borders to ensure each patch can be divisible by 2
_, _, _, h, w = LR_y_cube.size()
h = int(h // 16) * 16
w = int(w // 16) * 16
LR_y_cube = LR_y_cube[:, :, :, :h, :w]
SR_cb = SR_cb[:, :h * cfg.scale, :w * cfg.scale]
SR_cr = SR_cr[:, :h * cfg.scale, :w * cfg.scale]
SR_y = chop_forward(LR_y_cube, net,
cfg.scale).squeeze(0)
else:
SR_y = net(LR_y_cube).squeeze(0)
else:
SR_y = net(LR_y_cube).squeeze(0)
SR_y = np.array(SR_y.data.cpu())
SR_ycbcr = np.concatenate((SR_y, SR_cb, SR_cr),
axis=0).transpose(1, 2, 0)
SR_rgb = ycbcr2rgb(SR_ycbcr) * 255.0
SR_rgb = np.clip(SR_rgb, 0, 255)
SR_rgb = ToPILImage()(np.round(SR_rgb).astype(np.uint8))
if not os.path.exists('results/Vid4'):
os.mkdir('results/Vid4')
if not os.path.exists('results/Vid4/' + cfg.degradation +
'_x' + str(cfg.scale)):
os.mkdir('results/Vid4/' + cfg.degradation + '_x' +
str(cfg.scale))
if not os.path.exists('results/Vid4/' + cfg.degradation +
'_x' + str(cfg.scale) + '/' +
video_name):
os.mkdir('results/Vid4/' + cfg.degradation + '_x' +
str(cfg.scale) + '/' + video_name)
SR_rgb.save('results/Vid4/' + cfg.degradation + '_x' +
str(cfg.scale) + '/' + video_name + '/sr_' +
str(idx_iter + 2).rjust(2, '0') + '.png')