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.pil_open_rgb(self.img_list[index])
label = ImageHelper.pil_open_p(self.label_list[index])
if self.aug_transform is not None:
img, label = self.aug_transform(img, label=label)
if self.img_transform is not None:
img = self.img_transform(img)
if self.label_transform is not None:
label = self.label_transform(label)
return img, label
