def __getitem__(self, index):
img = ImageHelper.pil_open_rgb(self.img_list[index])
if os.path.exists(self.mask_list[index]):
maskmap = ImageHelper.pil_open_p(self.mask_list[index])
else:
maskmap = ImageHelper.np2img(np.ones((img.size[1], img.size[0]), dtype=np.uint8))
kpts, bboxes = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None and len(bboxes) > 0:
img, maskmap, kpts, bboxes = self.aug_transform(img, mask=maskmap, kpts=kpts, bboxes=bboxes)
elif self.aug_transform is not None:
img, maskmap, kpts = self.aug_transform(img, mask=maskmap, kpts=kpts)
width, height = maskmap.size
maskmap = ImageHelper.resize(maskmap, (width // self.configer.get('network', 'stride'),
height // self.configer.get('network', 'stride')), Image.NEAREST)
maskmap = np.expand_dims(np.array(maskmap, dtype=np.float32), axis=2)
heatmap = self.pose_data_utilizer.generate_heatmap(kpts=kpts, mask=maskmap)
vecmap = self.pose_data_utilizer.generate_paf(kpts=kpts, mask=maskmap)
if self.img_transform is not None:
img = self.img_transform(img)
if self.label_transform is not None:
heatmap = self.label_transform(heatmap)
vecmap = self.label_transform(vecmap)
maskmap = self.label_transform(maskmap)
return img, heatmap, maskmap, vecmap
def __getitem__(self, index):
img = ImageHelper.read_image(self.img_list[index],
tool=self.configer.get('data', 'image_tool'),
mode=self.configer.get('data', 'input_mode'))
if os.path.exists(self.mask_list[index]):
maskmap = ImageHelper.read_image(self.mask_list[index],
tool=self.configer.get('data', 'image_tool'), mode='P')
else:
maskmap = np.ones((img.size[1], img.size[0]), dtype=np.uint8)
if self.configer.get('data', 'image_tool') == 'pil':
maskmap = ImageHelper.np2img(maskmap)
kpts, bboxes = self.__read_json_file(self.json_list[index])
if self.aug_transform is not None and len(bboxes) > 0:
img, maskmap, kpts, bboxes = self.aug_transform(img, maskmap=maskmap, kpts=kpts, bboxes=bboxes)
elif self.aug_transform is not None:
img, maskmap, kpts = self.aug_transform(img, maskmap=maskmap, kpts=kpts)
width, height = maskmap.size
maskmap = ImageHelper.resize(maskmap,
(width // self.configer.get('network', 'stride'),
height // self.configer.get('network', 'stride')),
interpolation='nearest')
maskmap = torch.from_numpy(np.array(maskmap, dtype=np.float32))
kpts = torch.from_numpy(kpts).float()
if self.img_transform is not None:
img = self.img_transform(img)
return img, maskmap, kpts
def _reduce_zero_label(self, labelmap):
if not self.configer.get('data', 'reduce_zero_label'):
return labelmap
labelmap = np.array(labelmap)
encoded_labelmap = labelmap - 1
if self.configer.get('data', 'image_tool') == 'pil':
encoded_labelmap = ImageHelper.np2img(
encoded_labelmap.astype(np.uint8))
return encoded_labelmap
def _reduce_zero_label(self, labelmap):
if not self.configer.get('data', 'reduce_zero_label'):
return labelmap
labelmap = np.array(labelmap)
labelmap[labelmap == 0] = 255
labelmap = labelmap - 1
labelmap[labelmap == 254] = 255
if self.configer.get('data', 'image_tool') == 'pil':
labelmap = ImageHelper.np2img(labelmap.astype(np.uint8))
return labelmap
def _encode_label(self, labelmap):
labelmap = np.array(labelmap)
shape = labelmap.shape
encoded_labelmap = np.ones(shape=(shape[0], shape[1]),
dtype=np.float32) * 255
for i in range(self.configer.get('data', 'num_classes')):
class_id = self.configer.get('data', 'label_list')[i]
encoded_labelmap[labelmap == class_id] = i
if self.configer.get('data', 'image_tool') == 'pil':
encoded_labelmap = ImageHelper.np2img(
encoded_labelmap.astype(np.uint8))
return encoded_labelmap
