def __getitem__(self, index):
img_path = self.imgs[index]
pt = img_path.split('_', 1)[0]
p = re.findall('\d+', pt)[0]
# print(p)
clear_img_path = self.clear_imgs[int(p) - 1]
# print('hazypath:'+img_path+'clearpath:'+clear_img_path)
img = Image.open(img_path).convert('RGB')
clear_img = Image.open(clear_img_path).convert('RGB')
im_w, im_h = img.size
if im_w % 4 != 0 or im_h % 4 != 0:
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
img = np.array(img).astype('float')
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
edge_data = edge_compute(img_data)
in_data = torch.cat((img_data, edge_data), dim=0)
# print(in_data.shape)
cim_w, cim_h = clear_img.size
if cim_w % 4 != 0 or cim_h % 4 != 0:
clear_img = clear_img.resize(
(int(cim_w // 4 * 4), int(cim_h // 4 * 4)))
clear_img = np.array(clear_img).astype('float')
clear_img_data = torch.from_numpy(clear_img.transpose(
(2, 0, 1))).float()
# if self.transforms:
# in_data = self.transforms(in_data)
return in_data, clear_img_data - img_data
def get_images(self, index):
haze_name = self.haze_names[index][:-1]
#gt_name = haze_name.split('_')[0] + '.png'
haze_img = Image.open(self.haze_dir + haze_name).convert('RGB')
width, height = haze_img.size
im_w, im_h = haze_img.size
if im_w % 4 != 0 or im_h % 4 != 0:
haze_img = haze_img.resize(
(int(im_w // 4 * 4), int(im_h // 4 * 4)))
img = np.array(haze_img).astype('float')
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
edge_data = edge_compute(img_data)
in_data = torch.cat((img_data, edge_data), dim=0) - 128
return in_data, haze_name
def get_images(self, index):
haze_name = self.haze_names[index]
gt_name = haze_name.split('_')[0] + '.png'
haze_img = Image.open(self.haze_dir + haze_name)
gt_img = Image.open(self.gt_dir + gt_name)
haze_reshaped = haze_img
haze_reshaped = haze_reshaped.resize((256, 256), Image.ANTIALIAS)
# --- Transform to tensor --- #
transform_haze = Compose(
[ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
transform_gt = Compose([ToTensor()])
haze = transform_haze(haze_img)
haze_reshaped = transform_haze(haze_reshaped)
gt = transform_gt(gt_img)
haze_edge = edge_compute(haze)
haze = torch.cat((haze, haze_edge), 0)
return haze, haze_reshaped, gt, haze_name
def __getitem__(self, index):
crop_width, crop_height = self.crop_size
haze_name = self.haze_names[index]
haze_gt_img = Image.open(self.haze_dir + haze_name).convert('RGB')
haze_img = haze_gt_img.crop((0, 0, 400, 400))
gt_img = haze_gt_img.crop((400, 0, 800, 400))
width, height = haze_img.size
if width < crop_width or height < crop_height:
if width < height:
haze_img = haze_img.resize((260, (int)(height * 260 / width)),
Image.ANTIALIAS)
gt_img = gt_img.resize((260, (int)(height * 260 / width)),
Image.ANTIALIAS)
elif width >= height:
haze_img = haze_img.resize(((int)(width * 260 / height), 260),
Image.ANTIALIAS)
gt_img = gt_img.resize(((int)(width * 260 / height), 260),
Image.ANTIALIAS)
# --- random crop --- #
x, y = randrange(0, width - crop_width + 1), randrange(
0, height - crop_height + 1)
haze_crop_img = haze_img.crop((x, y, x + crop_width, y + crop_height))
gt_crop_img = gt_img.crop((x, y, x + crop_width, y + crop_height))
transform_haze = Compose(
[ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
transform_gt = Compose([ToTensor()])
haze = transform_haze(haze_crop_img)
gt = transform_gt(gt_crop_img)
edge = edge_compute(haze)
haze = torch.cat((haze, edge), 0)
return haze, gt
else:
net.float()
net.load_state_dict(torch.load(opt.model, map_location='cpu'))
net.eval()
for img_path in test_img_paths:
print(img_path)
time_start = time.time()
img = Image.open(img_path).convert('RGB')
im_w, im_h = img.size
if im_w % 4 != 0 or im_h % 4 != 0:
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
img = np.array(img).astype('float')
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
edge_data = edge_compute(img_data)
in_data = torch.cat((img_data, edge_data), dim=0).unsqueeze(0) - 128
in_data = in_data.cuda() if opt.use_cuda else in_data.float()
with torch.no_grad():
pred = net(Variable(in_data))
if opt.only_residual:
out_img_data = (pred.data[0].cpu().float() + img_data).round().clamp(
0, 255)
else:
out_img_data = pred.data[0].cpu().float().round().clamp(0, 255)
time_end = time.time()
print("fps = " + str(1 / (time_end - time_start)))
out_img = Image.fromarray(out_img_data.numpy().astype(np.uint8).transpose(
1, 2, 0))
out_img.save(
