def main(opt):
""" Main process of test.py """
# Load Model
model = utils.loadModel(opt.checkpoint,
ImproveNet(opt.rb),
dataparallel=True)
device = utils.selectDevice(
) if opt.cuda and torch.cuda.is_available() else 'cpu'
if opt.normalize:
model = nn.Sequential(
model,
InverseMeanShift(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])).to(device)
# Inference the images on the test set
predict(opt, model, opt.hazy, opt.dehazy, device)
# Measure the performance on the validation set
# predict(opt, net, opt.hazy, opt.dehazy, device)
# Measure the performance on the test set
gts = sorted(utils.load_all_image(opt.gt))
dehazes = sorted(utils.load_all_image(opt.dehazy))
if opt.record is None:
validate(dehazes, gts)
if opt.record is not None:
validate(dehazes, gts, os.path.join(opt.dehazy, opt.record))
return
def main(opt):
if os.path.isdir(opt.dehaze):
dehazes = utils.load_all_image(opt.dehaze).sort()
elif os.path.isfile(opt.dehaze):
dehazes = [opt.dehaze]
if os.path.isdir(opt.gt):
gts = utils.load_all_image(opt.gt).sort()
elif os.path.isfile(opt.gt):
gts = [opt.gt]
if len(dehazes) != len(gts):
raise ValueError(
"The image of dehaze should be equal to GTs, found {} and {}".
format(len(dehazes), len(gts)))
val(dehazes, gts, opt.output)
def predict(opt, net, folder, output_folder, device, *args):
"""
Generate the result.
Parameters
----------
folder, output : str
Input directory and output directory
net : nn.Module
(...)
normalized : bool
(...)
"""
# Test photos: Default Reside
images = utils.load_all_image(opt.hazy)
mean = utils.mean.to(device)
std = utils.std.to(device)
if opt.normalize:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
else:
transform = transforms.ToTensor()
# Inference
for im_path in images:
filename = im_path.split('/')[-1]
im = Image.open(im_path)
h, w = im.size
print("==========> Input filename: {}".format(filename))
with torch.no_grad():
im = transform(im).view(1, -1, w, h).to(device)
im_dehaze = net(im)
im = im.cpu()
im_dehaze = torch.clamp(im_dehaze, 0., 1.).cpu().data[0]
im_dehaze = transforms.ToPILImage()(im_dehaze)
im_dehaze.save(os.path.join(opt.dehazy, filename))
print("==========> File saved: {}".format(
os.path.join(opt.dehazy, filename)))
return
parser = argparse.ArgumentParser(description="PyTorch DeepDehazing")
parser.add_argument("--data",
type=str,
default="output",
help="path to load data images")
parser.add_argument("--gt", type=str, help="path to load gt images")
parser.add_argument("--view", type=str, help="path to save data gt images")
opt = parser.parse_args()
print(opt)
if not os.path.exists(opt.view):
os.makedirs(opt.view)
datas = utils.load_all_image(opt.data)
datas.sort()
def output_psnr_mse(img_orig, img_out):
squared_error = np.square(img_orig - img_out)
mse = np.mean(squared_error)
psnr = 10 * np.log10(1.0 / mse)
return psnr
psnrs = []
for data_p in tqdm(datas):
data = Image.open(data_p)
gt = Image.open(join(opt.gt, data_p.split('/')[-1][:-3] + 'jpg'))
w, h = data.size
parser = argparse.ArgumentParser(description="PyTorch DeepDehazing")
parser.add_argument("--rb", type=int, default=18, help="number of residual blocks")
parser.add_argument("--checkpoint", type=str, help="path to load model checkpoint")
parser.add_argument("--test", type=str, help="path to load test images")
opt = parser.parse_args()
print(opt)
net = Net(opt.rb)
net.load_state_dict(torch.load(opt.checkpoint)['state_dict'])
net.eval()
net = nn.DataParallel(net, device_ids=[0, 1, 2, 3]).cuda()
print(net)
images = utils.load_all_image(opt.test)
for im_path in tqdm(images):
filename = im_path.split('/')[-1]
print(filename)
im = Image.open(im_path)
h, w = im.size
print(h, w)
im = ToTensor()(im)
im = Variable(im).view(1, -1, w, h)
im = im.cuda()
with torch.no_grad():
im = net(im)
im = torch.clamp(im, 0., 1.)
im = im.cpu()
im = im.data[0]
def __init__(self, data_path):
self.images_all_list = glob.glob(data_path + 'images\*.jpg')
self.img_all = load_all_image(data_path)
self.img_size = np.array(self.img_all[0].shape)
self.data_type = 'Pred'
def __init__(self, data_path):
self.img_all = load_all_image(data_path)
self.label_all = load_all_labels(data_path, 'mat')
self.img_size = np.array(self.img_all[0].shape)
self.data_type = 'Train'
parser = argparse.ArgumentParser(description="Test Script")
parser.add_argument("--output",
type=str,
required=True,
help="path to save output images")
parser.add_argument("--datas",
type=str,
required=True,
help="path to load datas images")
opt = parser.parse_args()
print(opt)
if not os.path.exists(opt.output):
os.makedirs(opt.output)
datas = opt.datas.split(',')
data_images = [utils.load_all_image(data) for data in datas]
[data_image.sort() for data_image in data_images]
for i, p in enumerate(data_images[0]):
filename = p.split('/')[-1]
image_paths = [ps[i] for ps in data_images]
images = [Image.open(ip) for ip in image_paths]
images_np = [np.asarray(image) for image in images]
output = np.mean(images_np, axis=0)
output = output.round()
output[output >= 255] = 255
output[output <= 0] = 0
output = Image.fromarray(output.astype(np.uint8), mode='RGB')
output.save(os.path.join(opt.output, filename))