|
- import argparse
- import cv2
- import glob
- import numpy as np
- import os
- import torch
- from tqdm import tqdm
-
- from basicsr.archs.ridnet_arch import RIDNet
- from basicsr.utils.img_util import img2tensor, tensor2img
-
- if __name__ == '__main__':
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- parser = argparse.ArgumentParser()
- parser.add_argument('--test_path', type=str, default='datasets/denoise/RNI15')
- parser.add_argument('--noise_g', type=int, default=25)
- parser.add_argument(
- '--model_path',
- type=str,
- default= # noqa: E251
- 'experiments/pretrained_models/RIDNet/RIDNet.pth')
- args = parser.parse_args()
- if args.test_path.endswith('/'): # solve when path ends with /
- args.test_path = args.test_path[:-1]
- test_root = os.path.join(args.test_path, f'X{args.noise_g}')
- result_root = f'results/RIDNet/{os.path.basename(args.test_path)}'
- os.makedirs(result_root, exist_ok=True)
-
- # set up the RIDNet
- net = RIDNet(3, 64, 3).to(device)
- checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage)
- net.load_state_dict(checkpoint)
- net.eval()
-
- # scan all the jpg and png images
- img_list = sorted(glob.glob(os.path.join(test_root, '*.[jp][pn]g')))
- pbar = tqdm(total=len(img_list), desc='')
- for idx, img_path in enumerate(img_list):
- img_name = os.path.basename(img_path).split('.')[0]
- pbar.update(1)
- pbar.set_description(f'{idx}: {img_name}')
- # read image
- img = cv2.imread(img_path, cv2.IMREAD_COLOR)
- img = img2tensor(img, bgr2rgb=True, float32=True).unsqueeze(0).to(device)
- # inference
- with torch.no_grad():
- output = net(img)
- # save image
- output = tensor2img(output, rgb2bgr=True, out_type=np.uint8, min_max=(0, 255))
- save_img_path = os.path.join(result_root, f'{img_name}_x{args.noise_g}_RIDNet.png')
- cv2.imwrite(save_img_path, output)
|