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.')
if args.ref_path and not os.path.isfile(args.ref_path):
raise ValueError('It seems that you did not input a valid '
'"ref_path". Please double check your input, or '
'you may want to use "ref_path=None" '
'for single restoration.')
model = init_model(args.config,
args.checkpoint,
device=torch.device('cuda', args.device))
if args.ref_path: # Ref-SR
output = restoration_inference(model, args.img_path, args.ref_path)
else: # SISR
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')
def inference(self, data, *args, **kwargs):
# generate temp image path for restoration_inference
temp_name = ''.join(
random.sample(string.ascii_letters + string.digits, 18))
temp_path = f'./{temp_name}.png'
data.save(temp_path)
results = restoration_inference(self.model, temp_path)
# delete the temp image path
os.remove(temp_path)
return results
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)
mmcv.imwrite(output, args.save_path)
if args.imshow:
mmcv.imshow(output, 'predicted restoration result')
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