def set_input(self):
"""Set the input for the network"""
# get the test mask from painter
self.PaintPanel.saveDraw()
buffer = QtCore.QBuffer()
buffer.open(QtCore.QBuffer.ReadWrite)
self.PaintPanel.map.save(buffer, 'PNG')
pil_im = Image.open(io.BytesIO(buffer.data()))
# transform the image to the tensor
img = self.transform(self.img)
value = self.comboBox.currentIndex()
if value > 4:
mask = torch.autograd.Variable(self.transform(pil_im)).unsqueeze(0)
# mask from the random mask
# mask = Image.open(self.mname)
# mask = torch.autograd.Variable(self.transform(mask)).unsqueeze(0)
mask = (mask < 1).float()
else:
mask = task.center_mask(img).unsqueeze(0)
if len(self.opt.gpu_ids) > 0:
img = img.unsqueeze(0).cuda(self.opt.gpu_ids[0], async=True)
mask = mask.cuda(self.opt.gpu_ids[0], async=True)
# get I_m and I_c for image with mask and complement regions for training
mask = mask
self.img_truth = img * 2 - 1
self.img_m = mask * self.img_truth
self.img_c = (1 - mask) * self.img_truth
return self.img_m, self.img_c, self.img_truth, mask
def load_mask(self, img, index, annotations, img_path):
"""Load different mask types for training and testing"""
mask_type_index = random.randint(0, len(self.opt.mask_type) - 1)
mask_type = self.opt.mask_type[mask_type_index]
# center mask
if mask_type == 0:
return task.center_mask(img)
# random regular mask
if mask_type == 1:
return task.random_regular_mask(img)
# random irregular mask
if mask_type == 2:
return task.random_irregular_mask(img)
# external mask from "Image Inpainting for Irregular Holes Using Partial Convolutions (ECCV18)"
if mask_type == 3:
if self.opt.isTrain:
mask_index = random.randint(0, self.mask_size - 1)
else:
mask_index = index
mask_pil = Image.open(self.mask_paths[mask_index]).convert('RGB')
size = mask_pil.size[0]
if size > mask_pil.size[1]:
size = mask_pil.size[1]
mask_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(10),
transforms.CenterCrop([size, size]),
transforms.Resize(self.opt.fineSize),
transforms.ToTensor()
])
mask = (mask_transform(mask_pil) == 0).float()
mask_pil.close()
return mask
## add new option 4
if mask_type == 4:
return task.nose_mask(img, annotations['nose_x'],
annotations['nose_y'], img_path)
def set_input(self):
"""Set the input for the network"""
# get the test mask from painter
text = self.textEdit.toPlainText()
text_idx, text_len = util._caption_to_idx(self.model.wordtoix, text,
len(text))
self.text_idx = torch.Tensor([text_idx]).long()
self.text_len = torch.Tensor([text_len]).long()
self.PaintPanel.saveDraw()
buffer = QtCore.QBuffer()
buffer.open(QtCore.QBuffer.ReadWrite)
self.PaintPanel.map.save(buffer, 'PNG')
pil_im = Image.open(io.BytesIO(buffer.data()))
# transform the image to the tensor
img = self.transform(self.img)
value = self.comboBox.currentIndex()
if value < 4:
mask = torch.autograd.Variable(self.transform(pil_im)).unsqueeze(0)
# mask from the random mask
# mask = Image.open(self.mname)
# mask = torch.autograd.Variable(self.transform(mask)).unsqueeze(0)
mask = (mask < 1).float()
else:
mask = task.center_mask(img).unsqueeze(0)
if len(self.opt.gpu_ids) > 0:
img = img.unsqueeze(0).cuda(self.opt.gpu_ids[0])
mask = mask.cuda(self.opt.gpu_ids[0])
# get I_m and I_c for image with mask and complement regions for training
mask = mask
self.img_truth = img * 2 - 1
self.img_m = mask * self.img_truth
self.img_c = (1 - mask) * self.img_truth
return self.img_m, self.img_c, self.img_truth, mask, self.text_idx, self.text_len