def main():
args = parse_args()
model = init_model(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
result = inpainting_inference(model, args.masked_img_path, args.mask_path)
result = tensor2img(result, min_max=(-1, 1))[..., ::-1]
mmcv.imwrite(result, args.save_path)
if args.imshow:
mmcv.imshow(result, 'predicted inpainting result')
def main():
args = parse_args()
model = init_model(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
output = restoration_inference(model, args.img_path)
output = tensor2img(output)
# print(np.shape(output))
mmcv.imwrite(output, args.save_path)
if args.imshow:
mmcv.imshow(output, 'predicted restoration result')
def main():
args = parse_args()
model = init_model(
args.config, args.checkpoint, device=torch.device('cuda', args.device))
output = restoration_video_inference(model, args.input_dir,
args.window_size, args.filename_tmpl)
for i in range(0, output.size(1)):
output_i = output[:, i, :, :, :]
output_i = tensor2img(output_i)
save_path_i = f'{args.output_dir}/{i:08d}.png'
mmcv.imwrite(output_i, save_path_i)
def forward_test(self,
lq,
gt=None,
meta=None,
save_image=False,
save_path=None,
iteration=None):
"""Testing forward function.
Args:
lq (Tensor): LQ Tensor with shape (n, c, h, w).
gt (Tensor): GT Tensor with shape (n, c, h, w). Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results.
"""
output = self.generator(lq)
# normalize from [-1, 1] to [0, 1]
output = (output + 1) / 2.0
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
gt = (gt + 1) / 2.0 # normalize from [-1, 1] to [0, 1]
results = dict(eval_result=self.evaluate(output, gt))
else:
results = dict(lq=lq.cpu(), output=output.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
lq_path = meta[0]['lq_path']
folder_name = osp.splitext(osp.basename(lq_path))[0]
if isinstance(iteration, numbers.Number):
save_path = osp.join(save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{folder_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(output), save_path)
return results
def _save_image(meta, iteration, save_path, output):
"""Save the image.
Args:
meta (list[dict]): Meta data, such as path of target file.
Default: None. These dictionaries should contain
'target_path' (str of a path) or 'inputs_path' (list of str)
iteration (int): Iteration for the saving image name.
Default: None.
save_path (str): Path to save image. Default: None.
output (Tensor): Output image.
"""
if output.ndim == 4: # an image
img_name = meta[0]['key'].replace('/', '_')
if isinstance(iteration, numbers.Number):
save_path = osp.join(save_path,
f'{img_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{img_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(output), save_path)
elif output.ndim == 5: # a sequence
folder_name = meta[0]['key'].split('/')[0]
for i in range(0, output.size(1)):
if isinstance(iteration, numbers.Number):
save_path_i = osp.join(save_path, folder_name,
f'{i:08d}-{iteration + 1:06d}.png')
elif iteration is None:
save_path_i = osp.join(save_path, folder_name,
f'{i:08d}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(output[:, i, :, :, :]), save_path_i)
def evaluate(self, output, gt):
"""Evaluation function.
If the output contains multiple frames, we compute the metric
one by one and take an average.
Args:
output (Tensor): Model output with shape (n, t, c, h, w).
gt (Tensor): GT Tensor with shape (n, t, c, h, w).
Returns:
dict: Evaluation results.
"""
crop_border = self.test_cfg.crop_border
convert_to = self.test_cfg.get('convert_to', None)
eval_result = dict()
for metric in self.test_cfg.metrics:
if output.ndim == 5: # a sequence: (n, t, c, h, w)
avg = []
for i in range(0, output.size(1)):
output_i = tensor2img(output[:, i, :, :, :])
gt_i = tensor2img(gt[:, i, :, :, :])
avg.append(self.allowed_metrics[metric](
output_i, gt_i, crop_border, convert_to=convert_to))
eval_result[metric] = np.mean(avg)
elif output.ndim == 4: # an image: (n, c, t, w), for Vimeo-90K-T
output_img = tensor2img(output)
gt_img = tensor2img(gt)
value = self.allowed_metrics[metric](output_img,
gt_img,
crop_border,
convert_to=convert_to)
eval_result[metric] = value
return eval_result
def main():
""" Demo for video restoration models.
Note that we accept video as input/output, when 'input_dir'/'output_dir'
is set to the path to the video. But using videos introduces video
compression, which lowers the visual quality. If you want actual quality,
please save them as separate images (.png).
"""
args = parse_args()
model = init_model(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
output = restoration_video_inference(model, args.input_dir,
args.window_size, args.start_idx,
args.filename_tmpl, args.max_seq_len)
file_extension = os.path.splitext(args.output_dir)[1]
if file_extension in VIDEO_EXTENSIONS: # save as video
h, w = output.shape[-2:]
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video_writer = cv2.VideoWriter(args.output_dir, fourcc, 25, (w, h))
for i in range(0, output.size(1)):
img = tensor2img(output[:, i, :, :, :])
video_writer.write(img.astype(np.uint8))
cv2.destroyAllWindows()
video_writer.release()
else:
for i in range(args.start_idx, args.start_idx + output.size(1)):
output_i = output[:, i - args.start_idx, :, :, :]
output_i = tensor2img(output_i)
save_path_i = f'{args.output_dir}/{args.filename_tmpl.format(i)}'
mmcv.imwrite(output_i, save_path_i)
def merge_frames(input_tensors, output_tensors):
"""merge input frames and output frames.
This is a basic function, interpolate a frame between the given two
frames.
Args:
input_tensors (Tensor): The input frames with shape [n, 2, c, h, w]
output_tensors (Tensor): The output frames with shape
[n, 1, c, h, w].
Returns:
list[np.array]: The final frames.
in_frame, out_frame, in_frame, out_frame, in_frame ...
"""
num_frames = input_tensors.shape[0]
result = []
for i in range(num_frames):
result.append(tensor2img(input_tensors[i, 0]))
result.append(tensor2img(output_tensors[i, 0]))
result.append(tensor2img(input_tensors[-1, 1]))
return result
def main():
args = parse_args()
if not os.path.isfile(args.img_path):
raise ValueError('It seems that you did not input a valid '
'"image_path". Please double check your input, or '
'you may want to use "restoration_video_demo.py" '
'for video restoration.')
model = init_model(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
output = restoration_inference(model, args.img_path)
output = tensor2img(output)
mmcv.imwrite(output, args.save_path)
if args.imshow:
mmcv.imshow(output, 'predicted restoration result')
else:
torch.cuda.empty_cache()
# Inpainting
t0 = time.time()
## create mask for inpainting
mask = b_boxes_manager.mask(image)
## Save temp mask file temporarly
mask_path = "temp.png"
Image.fromarray(mask).save("temp.png")
## init the inpainter
model = init_model(config["Inpainter"]["configuration"],
config["Inpainter"]["pretrained_model"],
device=device)
## inpaint
result = inpainting_inference(model, image_path, mask_path)[0]
result = tensor2img(result, min_max=(-1, 1))
os.remove("temp.png")
print("Inpainting time : {:.3f}".format(time.time() - t0))
# Translation
t0 = time.time()
result = b_boxes_manager.fill_translation(
result,
same_block_th=config["same_block_th"],
translation_service=config["translation_service"],
src_language=src_language,
dest_language=dst_language,
result_path=result_path)
print("Translation time : {:.3f}".format(time.time() - t0))
# open results
def forward_test(self,
lq,
gt=None,
meta=None,
save_image=False,
save_path=None,
iteration=None,
multi_scale=False):
"""Testing forward function.
Args:
lq (Tensor): LQ Tensor with shape (n, c, h, w). (n, t, c, h, w).
gt (Tensor): GT Tensor with shape (n, c, h, w). Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results.
"""
# todo add multi_scale test code
if multi_scale:
output = self.multi_scale_test(lq)
else:
output = self.generator(lq)
# output = self.generator(torch.flip(lq, (-1, ))) # w
# output = self.generator(torch.flip(lq, (-2, ))) # h
# output = self.generator(torch.flip(lq, (-2, -1))) # wh
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
results = dict(eval_result=self.evaluate(output, gt))
else:
results = dict(lq=lq.cpu(), output=output.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
gt_path = meta[0]['gt_path'][0]
folder_name = meta[0]['key'].split('/')[0]
frame_name = osp.splitext(osp.basename(gt_path))[0]
center_frame_idx = int(frame_name)
# if center_frame_idx % 10 == 0:
if True:
if isinstance(iteration, numbers.Number):
save_path = osp.join(
save_path, folder_name,
f'{frame_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, folder_name,
f'{frame_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
output = tensor2img(output)
# print(np.shape(output))
# output_crop = output[:536, :, :]
output_crop = output[:-8, :, :]
# print(np.shape(output_crop))
mmcv.imwrite(output_crop, save_path)
return results
def forward_test(self,
img_a,
img_b,
meta,
save_image=False,
save_path=None,
iteration=None):
"""Forward function for testing.
Args:
img_a (Tensor): Input image from domain A.
img_b (Tensor): Input image from domain B.
meta (list[dict]): Input meta data.
save_image (bool, optional): If True, results will be saved as
images. Default: False.
save_path (str, optional): If given a valid str path, the results
will be saved in this path. Default: None.
iteration (int, optional): Iteration number. Default: None.
Returns:
dict: Dict of forward and evaluation results for testing.
"""
# No need for metrics during training for pix2pix. And
# this is a special trick in pix2pix original paper & implementation,
# collecting the statistics of the test batch at test time.
self.train()
# necessary setup
real_a, real_b, image_path = self.setup(img_a, img_b, meta)
fake_b = self.generator(real_a)
results = dict(real_a=real_a.cpu(),
fake_b=fake_b.cpu(),
real_b=real_b.cpu())
# save image
if save_image:
assert save_path is not None
folder_name = osp.splitext(osp.basename(image_path[0]))[0]
if self.show_input:
if iteration:
save_path = osp.join(
save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}-ra-fb-rb.png')
else:
save_path = osp.join(save_path,
f'{folder_name}-ra-fb-rb.png')
output = np.concatenate([
tensor2img(results['real_a'], min_max=(-1, 1)),
tensor2img(results['fake_b'], min_max=(-1, 1)),
tensor2img(results['real_b'], min_max=(-1, 1))
],
axis=1)
else:
if iteration:
save_path = osp.join(
save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}-fb.png')
else:
save_path = osp.join(save_path, f'{folder_name}-fb.png')
output = tensor2img(results['fake_b'], min_max=(-1, 1))
flag = mmcv.imwrite(output, save_path)
results['saved_flag'] = flag
return results
def test_tensor2img():
tensor_4d_1 = torch.FloatTensor(2, 3, 4, 4).uniform_(0, 1)
tensor_4d_2 = torch.FloatTensor(1, 3, 4, 4).uniform_(0, 1)
tensor_4d_3 = torch.FloatTensor(3, 1, 4, 4).uniform_(0, 1)
tensor_4d_4 = torch.FloatTensor(1, 1, 4, 4).uniform_(0, 1)
tensor_3d_1 = torch.FloatTensor(3, 4, 4).uniform_(0, 1)
tensor_3d_2 = torch.FloatTensor(3, 6, 6).uniform_(0, 1)
tensor_3d_3 = torch.FloatTensor(1, 6, 6).uniform_(0, 1)
tensor_2d = torch.FloatTensor(4, 4).uniform_(0, 1)
with pytest.raises(TypeError):
# input is not a tensor
tensor2img(4)
with pytest.raises(TypeError):
# input is not a list of tensors
tensor2img([tensor_3d_1, 4])
with pytest.raises(ValueError):
# unsupported 5D tensor
tensor2img(torch.FloatTensor(2, 2, 3, 4, 4).uniform_(0, 1))
# 4d
rlt = tensor2img(tensor_4d_1, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_4d_1_np = make_grid(tensor_4d_1, nrow=1, normalize=False).numpy()
tensor_4d_1_np = np.transpose(tensor_4d_1_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_1_np * 255).round())
rlt = tensor2img(tensor_4d_2, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_4d_2_np = tensor_4d_2.squeeze().numpy()
tensor_4d_2_np = np.transpose(tensor_4d_2_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_2_np * 255).round())
rlt = tensor2img(tensor_4d_3, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_4d_3_np = make_grid(tensor_4d_3, nrow=1, normalize=False).numpy()
tensor_4d_3_np = np.transpose(tensor_4d_3_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_3_np * 255).round())
rlt = tensor2img(tensor_4d_4, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_4d_4_np = tensor_4d_4.squeeze().numpy()
np.testing.assert_almost_equal(rlt, (tensor_4d_4_np * 255).round())
# 3d
rlt = tensor2img([tensor_3d_1, tensor_3d_2],
out_type=np.uint8,
min_max=(0, 1))
assert rlt[0].dtype == np.uint8
tensor_3d_1_np = tensor_3d_1.numpy()
tensor_3d_1_np = np.transpose(tensor_3d_1_np[[2, 1, 0], :, :], (1, 2, 0))
tensor_3d_2_np = tensor_3d_2.numpy()
tensor_3d_2_np = np.transpose(tensor_3d_2_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt[0], (tensor_3d_1_np * 255).round())
np.testing.assert_almost_equal(rlt[1], (tensor_3d_2_np * 255).round())
rlt = tensor2img(tensor_3d_3, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_3d_3_np = tensor_3d_3.squeeze().numpy()
np.testing.assert_almost_equal(rlt, (tensor_3d_3_np * 255).round())
# 2d
rlt = tensor2img(tensor_2d, out_type=np.uint8, min_max=(0, 1))
assert rlt.dtype == np.uint8
tensor_2d_np = tensor_2d.numpy()
np.testing.assert_almost_equal(rlt, (tensor_2d_np * 255).round())
rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0, 1))
assert rlt.dtype == np.float32
np.testing.assert_almost_equal(rlt, tensor_2d_np)
rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0.1, 0.5))
assert rlt.dtype == np.float32
tensor_2d_np = (np.clip(tensor_2d_np, 0.1, 0.5) - 0.1) / 0.4
np.testing.assert_almost_equal(rlt, tensor_2d_np)
def forward_test(self,
masked_img,
mask,
save_image=False,
save_path=None,
iteration=None,
**kwargs):
"""Forward function for testing.
Args:
masked_img (torch.Tensor): Tensor with shape of (n, 3, h, w).
mask (torch.Tensor): Tensor with shape of (n, 1, h, w).
save_image (bool, optional): If True, results will be saved as
image. Defaults to False.
save_path (str, optional): If given a valid str, the reuslts will
be saved in this path. Defaults to None.
iteration (int, optional): Iteration number. Defaults to None.
Returns:
dict: Contain output results and eval metrics (if have).
"""
mask_input = mask.expand_as(masked_img)
mask_input = 1. - mask_input
fake_res, final_mask = self.generator(masked_img, mask_input)
fake_img = fake_res * mask + masked_img * (1. - mask)
output = dict()
eval_results = {}
if self.eval_with_metrics:
gt_img = kwargs['gt_img']
data_dict = dict(gt_img=gt_img, fake_res=fake_res, mask=mask)
for metric_name in self.test_cfg['metrics']:
if metric_name in ['ssim', 'psnr']:
eval_results[metric_name] = self._eval_metrics[
metric_name](tensor2img(fake_img, min_max=(-1, 1)),
tensor2img(gt_img, min_max=(-1, 1)))
else:
eval_results[metric_name] = self._eval_metrics[
metric_name]()(data_dict).item()
output['eval_results'] = eval_results
else:
output['fake_res'] = fake_res
output['fake_img'] = fake_img
output['final_mask'] = final_mask
output['meta'] = None if 'meta' not in kwargs else kwargs['meta'][0]
if save_image:
assert save_image and save_path is not None, (
'Save path should been given')
assert output['meta'] is not None, (
'Meta information should be given to save image.')
tmp_filename = output['meta']['gt_img_path']
filestem = Path(tmp_filename).stem
if iteration is not None:
filename = f'{filestem}_{iteration}.png'
else:
filename = f'{filestem}.png'
mmcv.mkdir_or_exist(save_path)
if kwargs.get('gt_img', None) is not None:
img_list = [kwargs['gt_img']]
else:
img_list = []
img_list.extend(
[masked_img,
mask.expand_as(masked_img), fake_res, fake_img])
img = torch.cat(img_list, dim=3).cpu()
self.save_visualization(img, osp.join(save_path, filename))
output['save_img_path'] = osp.abspath(
osp.join(save_path, filename))
return output
def forward_test(self,
lq,
gt=None,
meta=None,
save_image=False,
save_path=None,
iteration=None):
"""Testing forward function.
Args:
lq (Tensor): LQ image.
gt (Tensor): GT image.
meta (list[dict]): Meta data, such as path of GT file.
Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results, which contain either key(s)
1. 'eval_result'.
2. 'lq', 'pred'.
3. 'lq', 'pred', 'gt'.
"""
# generator
with torch.no_grad():
sr_list, _ = self.generator.forward(lq)
pred = sr_list[-1]
pred = self.img_denormalize(pred)
if gt is not None:
gt = self.img_denormalize(gt)
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
results = dict(eval_result=self.evaluate(pred, gt))
else:
results = dict(lq=lq.cpu(), output=pred.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
if 'gt_path' in meta[0]:
pred_path = meta[0]['gt_path']
else:
pred_path = meta[0]['lq_path']
folder_name = osp.splitext(osp.basename(pred_path))[0]
if isinstance(iteration, numbers.Number):
save_path = osp.join(save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{folder_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(pred), save_path)
return results
def run_inpainting(self,
img_or_path,
mask_or_path,
dilate_kernel_size=19,
dilate_iter_num=3):
"""
Args:
img_or_path (str or np.ndarray): (h, w, 3) is in the range of [0, 255] with BGR channel;
mask_or_path (str or np.ndarray): (h, w) is in the range of [0, 255], np.uint8;
dilate_kernel_size (int): the kernel size of dilation;
dilate_iter_num (int): the iterations of dilation;
Returns:
inpainting_result (np.ndarray): (h, w, 3), is in the range of [0, 255] with BGR channel.
"""
# TODO, do not write the middle outputs to disk, and make them in memory.
# scaled_src_path, scaled_mask_path, scaled_inpainting_result_path
img_name = str(time.time())
img_path = os.path.join(self.temp_dir, img_name)
if isinstance(img_or_path, str):
src_img = cv2.imread(img_or_path)
else:
src_img = img_or_path.copy()
"""
scaled image
"""
scaled_src_path = f"{img_path}_scaled.png"
scaled_mask_path = f"{img_path}_mask.png"
scaled_inpainting_result_path = f"{img_path}_inpainting.png"
origin_h, origin_w = src_img.shape[:2]
scaled_size = compute_scaled_size(
(origin_w, origin_h), control_size=self.inpainting_control_size)
scaled_src_img = cv2.resize(src_img, scaled_size)
cv2.imwrite(scaled_src_path, scaled_src_img)
"""
dilate mask
"""
if isinstance(mask_or_path, str):
mask = cv2.imread(mask_or_path, cv2.IMREAD_GRAYSCALE)
else:
# mask = (mask * 255).astype(np.uint8)
mask = mask_or_path
scaled_mask = cv2.resize(mask,
scaled_size,
interpolation=cv2.INTER_NEAREST)
kernel = cv2.getStructuringElement(
cv2.MORPH_ELLIPSE, (dilate_kernel_size, dilate_kernel_size))
dilated_scaled_mask = cv2.dilate(scaled_mask,
kernel,
iterations=dilate_iter_num)
cv2.imwrite(scaled_mask_path, dilated_scaled_mask)
"""
inpainting result
"""
scaled_result = inpainting_inference(self.inpainting_model,
scaled_src_path, scaled_mask_path)
# (h, w, 3) 0-255 bgr
scaled_result = tensor2img(scaled_result, min_max=(-1, 1))[..., ::-1]
cv2.imwrite(scaled_inpainting_result_path, scaled_result)
"""
super-resolution
"""
if self.cfg["use_sr"]:
result = restoration_inference(self.sr_model,
scaled_inpainting_result_path)
result = tensor2img(result)
result = cv2.resize(result, (origin_w, origin_h))
result = result.astype(np.uint8)
else:
result = cv2.resize(scaled_result, (origin_w, origin_h))
result = result.astype(np.uint8)
os.remove(scaled_src_path)
os.remove(scaled_mask_path)
os.remove(scaled_inpainting_result_path)
return result, dilated_scaled_mask
def forward_test(self,
lq,
gt=None,
meta=None,
save_image=False,
save_path=None,
iteration=None):
"""Testing forward function.
Args:
lq (Tensor): LQ Tensor with shape (n, t, c, h, w).
gt (Tensor): GT Tensor with shape (n, t, c, h, w). Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results.
"""
with torch.no_grad():
output = self.generator(lq)
# If the GT is an image (i.e. the center frame), the output sequence is
# turned to an image.
if gt is not None and gt.ndim == 4:
t = output.size(1)
if self.check_if_mirror_extended(lq): # with mirror extension
output = 0.5 * (output[:, t // 4] + output[:, -1 - t // 4])
else: # without mirror extension
output = output[:, t // 2]
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
results = dict(eval_result=self.evaluate(output, gt))
else:
results = dict(lq=lq.cpu(), output=output.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
if output.ndim == 4: # an image, key = 000001/0000 (Vimeo-90K)
img_name = meta[0]['key'].replace('/', '_')
if isinstance(iteration, numbers.Number):
save_path = osp.join(
save_path, f'{img_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{img_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(output), save_path)
elif output.ndim == 5: # a sequence, key = 000
folder_name = meta[0]['key'].split('/')[0]
for i in range(0, output.size(1)):
if isinstance(iteration, numbers.Number):
save_path_i = osp.join(
save_path, folder_name,
f'{i:08d}-{iteration + 1:06d}.png')
elif iteration is None:
save_path_i = osp.join(save_path, folder_name,
f'{i:08d}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(output[:, i, :, :, :]),
save_path_i)
return results
def forward_test(self,
lq,
lq_up,
ref,
ref_downup,
gt=None,
meta=None,
save_image=False,
save_path=None,
iteration=None):
"""Testing forward function.
Args:
lq (Tensor): LQ image
gt (Tensor): GT image
lq_up (Tensor): Upsampled LQ image
ref (Tensor): Reference image
ref_downup (Tensor): Image generated by sequentially applying
bicubic down-sampling and up-sampling on reference image
meta (list[dict]): Meta data, such as path of GT file.
Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results, which contain either key(s)
1. 'eval_result'.
2. 'lq', 'pred'.
3. 'lq', 'pred', 'gt'.
"""
# generator
with torch.no_grad():
pred = self.forward_dummy(lq=lq,
lq_up=lq_up,
ref=ref,
ref_downup=ref_downup)
pred = (pred + 1.) / 2.
if gt is not None:
gt = (gt + 1.) / 2.
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
results = dict(eval_result=self.evaluate(pred, gt))
else:
results = dict(lq=lq.cpu(), output=pred.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
if 'gt_path' in meta[0]:
the_path = meta[0]['gt_path']
else:
the_path = meta[0]['lq_path']
folder_name = osp.splitext(osp.basename(the_path))[0]
if isinstance(iteration, numbers.Number):
save_path = osp.join(save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{folder_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(pred), save_path)
return results
def forward_test(self,
lq,
gt,
coord,
cell,
meta=None,
save_image=False,
save_path=None,
iteration=None):
"""Testing forward function.
Args:
lq (Tensor): LQ image.
gt (Tensor): GT image.
coord (Tensor): Coord tensor.
cell (Tensor): Cell tensor.
meta (list[dict]): Meta data, such as path of GT file.
Default: None.
save_image (bool): Whether to save image. Default: False.
save_path (str): Path to save image. Default: None.
iteration (int): Iteration for the saving image name.
Default: None.
Returns:
dict: Output results, which contain either key(s)
1. 'eval_result'.
2. 'lq', 'pred'.
3. 'lq', 'pred', 'gt'.
"""
# norm
self.lq_mean = self.lq_mean.to(lq)
self.lq_std = self.lq_std.to(lq)
lq = (lq - self.lq_mean) / self.lq_std
# generator
with torch.no_grad():
pred = self.generator(lq, coord, cell, test_mode=True)
self.gt_mean = self.gt_mean.to(pred)
self.gt_std = self.gt_std.to(pred)
pred = pred * self.gt_std + self.gt_mean
pred.clamp_(0, 1)
# reshape for eval
ih, iw = lq.shape[-2:]
s = math.sqrt(coord.shape[1] / (ih * iw))
shape = [lq.shape[0], round(ih * s), round(iw * s), 3]
pred = pred.view(*shape).permute(0, 3, 1, 2).contiguous()
if gt is not None:
gt = gt.view(*shape).permute(0, 3, 1, 2).contiguous()
if self.test_cfg is not None and self.test_cfg.get('metrics', None):
assert gt is not None, (
'evaluation with metrics must have gt images.')
results = dict(eval_result=self.evaluate(pred, gt))
else:
results = dict(lq=lq.cpu(), output=pred.cpu())
if gt is not None:
results['gt'] = gt.cpu()
# save image
if save_image:
gt_path = meta[0]['gt_path']
folder_name = osp.splitext(osp.basename(gt_path))[0]
if isinstance(iteration, numbers.Number):
save_path = osp.join(save_path, folder_name,
f'{folder_name}-{iteration + 1:06d}.png')
elif iteration is None:
save_path = osp.join(save_path, f'{folder_name}.png')
else:
raise ValueError('iteration should be number or None, '
f'but got {type(iteration)}')
mmcv.imwrite(tensor2img(pred), save_path)
return results