import time
from scipy.io import loadmat
use_gpu = True
C = 3
dep_U = 4
# load the pretrained model
print('Loading the Model')
checkpoint = torch.load('./model_state/model_state_SIDD')
net = VDN(C, dep_U=dep_U, wf=64)
if use_gpu:
net = torch.nn.DataParallel(net).cuda()
net.load_state_dict(checkpoint)
else:
load_state_dict_cpu(net, checkpoint)
net.eval()
im_noisy = loadmat('./test_data/DND/1.mat')['InoisySRGB']
H, W, _ = im_noisy.shape
if H % 2**dep_U != 0:
H -= H % 2**dep_U
if W % 2**dep_U != 0:
W -= W % 2**dep_U
im_noisy = im_noisy[:H, :W, ]
im_noisy = torch.from_numpy(im_noisy.transpose((2,0,1))[np.newaxis,])
if use_gpu:
im_noisy = im_noisy.cuda()
print('Begin Testing on GPU')
else:
help="Wave choice for wavelet transform")
args = parser.parse_args()
# load the pre-trained model
print('Loading the model')
checkpoint = torch.load(args.ckpt)
model = FAN(depth_S=args.depth_S,
depth_U=args.depth_U,
feature_dims=64,
wave_pattern=args.wave_mode)
if args.gpu:
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(checkpoint)
print('Begin Testing on GPU')
else:
load_state_dict_cpu(model, checkpoint)
print('Begin Testing on CPU')
model.eval()
img_suffixes = (".png", ".bmp", ".jpg", ".jpeg")
img_paths = [
str(i) for i in Path(args.img_dir).glob('*noisy*')
if i.suffix.lower() in img_suffixes
]
for img_path in img_paths:
img_noisy = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) / 255
height, width, _ = img_noisy.shape
if height % 2**(args.depth_U - 1) != 0:
height -= height % 2**(args.depth_U - 1)
if width % 2**(args.depth_U - 1) != 0:
width -= width % 2**(args.depth_U - 1)