-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
74 lines (57 loc) · 2.73 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import numpy as np
import os
import argparse
from tqdm import tqdm
import torch.nn as nn
import torch
from torch.utils.data import DataLoader
import torch.nn.functional as F
import utils
from data_RGB import get_test_data
from Net import Net
from skimage import img_as_ubyte
from pdb import set_trace as stx
parser = argparse.ArgumentParser(description='Image Deblurring using Net')
parser.add_argument('--input_dir', default='./Datasets/', type=str, help='Directory of validation images')
parser.add_argument('--result_dir', default='./results/', type=str, help='Directory for results')
parser.add_argument('--weights', default='./pretrained_models/model_deblurring.pth', type=str, help='Path to weights')
parser.add_argument('--dataset', default='GoPro', type=str, help='Test Dataset') # ['GoPro', 'HIDE', 'RealBlur_J', 'RealBlur_R']
parser.add_argument('--gpus', default='0', type=str, help='CUDA_VISIBLE_DEVICES')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
model_restoration = Net()
utils.load_checkpoint(model_restoration,args.weights)
print("===>Testing using weights: ",args.weights)
model_restoration.cuda()
model_restoration = nn.DataParallel(model_restoration)
model_restoration.eval()
dataset = args.dataset
rgb_dir_test = os.path.join(args.input_dir, dataset, 'test', 'input')
test_dataset = get_test_data(rgb_dir_test, img_options={})
test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=4, drop_last=False, pin_memory=True)
result_dir = os.path.join(args.result_dir, dataset)
utils.mkdir(result_dir)
with torch.no_grad():
for ii, data_test in enumerate(tqdm(test_loader), 0):
torch.cuda.ipc_collect()
torch.cuda.empty_cache()
input_ = data_test[0].cuda()
filenames = data_test[1]
# Padding in case images are not multiples of 8
if dataset == 'RealBlur_J' or dataset == 'RealBlur_R':
factor = 8
h,w = input_.shape[2], input_.shape[3]
H,W = ((h+factor)//factor)*factor, ((w+factor)//factor)*factor
padh = H-h if h%factor!=0 else 0
padw = W-w if w%factor!=0 else 0
input_ = F.pad(input_, (0,padw,0,padh), 'reflect')
restored = model_restoration(input_)
restored = torch.clamp(restored[0],0,1)
# Unpad images to original dimensions
if dataset == 'RealBlur_J' or dataset == 'RealBlur_R':
restored = restored[:,:,:h,:w]
restored = restored.permute(0, 2, 3, 1).cpu().detach().numpy()
for batch in range(len(restored)):
restored_img = img_as_ubyte(restored[batch])
utils.save_img((os.path.join(result_dir, filenames[batch]+'.png')), restored_img)