def get_current_visuals(self, size=512):
out = []
visual_imgs = []
out.append(utils.tensor_to_numpy(self.img_LR))
out.append(utils.tensor_to_numpy(self.img_SR))
out.append(utils.tensor_to_numpy(self.img_HR))
out_imgs = [utils.batch_numpy_to_image(x, size) for x in out]
visual_imgs += out_imgs
visual_imgs.append(utils.color_parse_map(self.ref_mask_onehot, size))
visual_imgs.append(utils.color_parse_map(self.hr_mask, size))
return visual_imgs
def get_current_visuals(self, size=512):
out = []
visual_imgs = []
out.append(utils.tensor_to_numpy(self.img_LR))
out.append(utils.tensor_to_numpy(self.img_SR))
out.append(utils.tensor_to_numpy(self.img_HR))
out_imgs = [utils.batch_numpy_to_image(x, size) for x in out]
visual_imgs.append(out_imgs[0])
visual_imgs.append(out_imgs[1])
visual_imgs.append(utils.color_parse_map(self.pred_Parse))
visual_imgs.append(utils.color_parse_map(self.gt_Parse.unsqueeze(1)))
visual_imgs.append(out_imgs[2])
return visual_imgs
def enhance_faces(LQ_faces, model):
hq_faces = []
lq_parse_maps = []
for lq_face in tqdm(LQ_faces):
with torch.no_grad():
lq_tensor = torch.tensor(lq_face.transpose(2, 0, 1)) / 255. * 2 - 1
lq_tensor = lq_tensor.unsqueeze(0).float().to(model.device)
parse_map, _ = model.netP(lq_tensor)
parse_map_onehot = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float()
_, output_SR = model.netG(lq_tensor, parse_map_onehot)
hq_faces.append(utils.tensor_to_img(output_SR))
lq_parse_maps.append(utils.color_parse_map(parse_map_onehot)[0])
return hq_faces, lq_parse_maps
model.eval()
max_size = 9999
os.makedirs(os.path.join(save_dir, 'sr'), exist_ok=True)
for i, data in tqdm(enumerate(dataset),
total=len(dataset) // opt.batch_size):
inp = data['LR']
with torch.no_grad():
parse_map, _ = netP(inp)
parse_map_sm = (parse_map == parse_map.max(
dim=1, keepdim=True)[0]).float()
output_parse, output_SR = netG(inp, parse_map_sm)
img_path = data['LR_paths'] # get image paths
for i in tqdm(range(len(img_path))):
inp_img = utils.batch_tensor_to_img(inp)
output_sr_img = utils.batch_tensor_to_img(output_SR)
ref_parse_img = utils.color_parse_map(parse_map_sm)
save_path = os.path.join(save_dir, 'lq',
img_path[i].split('/')[-1])
os.makedirs(os.path.join(save_dir, 'lq'), exist_ok=True)
save_img = Image.fromarray(inp_img[i])
save_img.save(save_path)
save_path = os.path.join(save_dir, 'hq',
img_path[i].split('/')[-1])
os.makedirs(os.path.join(save_dir, 'hq'), exist_ok=True)
save_img = Image.fromarray(output_sr_img[i])
save_img.save(save_path)
save_path = os.path.join(save_dir, 'parse',
img_path[i].split('/')[-1])
