def __getitem__(self, index):
idx = self.index_list[index]
img = cv2.imread(os.path.join(self.root, self.data[idx]['img_paths']))
img_idx = self.data[idx]['img_paths'][-16:-3]
# print( img.shape)
# print("idx: {} dataset:{} img_idx: {}{}".format(idx, self.data[idx]['dataset'], self.root, self.data[idx]['img_paths']))
if "COCO_val" in self.data[idx]['dataset']:
mask_miss = cv2.imread(
self.mask_dir + 'mask2014/val2014/mask_COCO_val2014_' +
img_idx + 'jpg', 0)
elif "COCO" in self.data[idx]['dataset']:
mask_miss = cv2.imread(
self.mask_dir + 'mask2014/train2014/mask_COCO_train2014_' +
img_idx + 'jpg', 0)
# print (self.root + 'mask2014/val2014_mask_miss_' + img_idx + 'png')
meta_data = self.get_anno(self.data[idx])
meta_data = self.add_neck(meta_data)
meta_data, img, mask_miss = aug_scale(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_rotate(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_croppad(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_flip(meta_data, img, mask_miss,
self.params_transform)
meta_data = self.remove_illegal_joint(meta_data)
heat_mask, heatmaps, paf_mask, pafs = self.get_ground_truth(
meta_data, mask_miss)
# image preprocessing, which comply the model
# trianed on Imagenet dataset
if self.preprocess == 'rtpose':
img = rtpose_preprocess(img)
elif self.preprocess == 'vgg':
img = vgg_preprocess(img)
elif self.preprocess == 'inception':
img = inception_preprocess(img)
elif self.preprocess == 'ssd':
img = ssd_preprocess(img)
img = torch.from_numpy(img)
heatmaps = torch.from_numpy(
heatmaps.transpose((2, 0, 1)).astype(np.float32))
heat_mask = torch.from_numpy(
heat_mask.transpose((2, 0, 1)).astype(np.float32))
pafs = torch.from_numpy(pafs.transpose((2, 0, 1)).astype(np.float32))
paf_mask = torch.from_numpy(
paf_mask.transpose((2, 0, 1)).astype(np.float32))
return img, heatmaps, heat_mask, pafs, paf_mask
def __getitem__(self, index):
idx = self.index_list[index]
img = cv2.imread(os.path.join(self.root, self.data[idx]['img_paths']))
img_idx = self.data[idx]['img_paths'][:-3]
mask_miss = cv2.imread(
self.mask_dir + '/masks/mask_' + img_idx + 'jpg', 0)
# fix mask problems
mask_miss[mask_miss > 0] = 200.
mask_miss[mask_miss == 0] = 255.
mask_miss[mask_miss == 200] = 0.
meta_data = self.get_anno(self.data[idx])
meta_data = self.remove_zero(meta_data)
meta_data = self.get_anno_inmyorder(meta_data)
meta_data, img, mask_miss = aug_scale(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_rotate(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_croppad(meta_data, img, mask_miss,
self.params_transform)
meta_data, img, mask_miss = aug_flip(meta_data,
img,
mask_miss,
self.params_transform,
coco=False,
neworder=True)
meta_data = self.remove_illegal_joint(meta_data)
heat_mask, heatmaps, paf_mask, pafs = self.get_ground_truth(
meta_data, mask_miss)
##############################################
########## check labels ############
##############################################
# mask = cv2.resize(mask_miss, (384, 384))
# image = img
# heatmaps = cv2.resize(heatmaps, (384, 384))
# pafs = cv2.resize(pafs, (384, 384))
# mask = mask / 255
# mask = mask.astype(np.uint8)
# mask = np.expand_dims(mask, axis=2)
# mask = np.repeat(mask, 3, axis=2)
# image = cv2.multiply(mask, image)
# for j in range(0, 16, 4):
# heatmap = heatmaps[:, :, j]
# heatmap = heatmap.reshape((384, 384, 1))
# heatmap *= 255
# heatmap = heatmap.astype(np.uint8)
# heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
# plt.imshow(image)
# plt.imshow(heatmap, alpha=0.5)
# plt.show()
# plt.close()
# for j in range(0, 32, 2):
# paf = np.abs(pafs[:, :, j])
# paf += np.abs(pafs[:, :, j + 1])
# paf[paf > 1] = 1
# paf = paf.reshape((384, 384, 1))
# paf *= 255
# paf = paf.astype(np.uint8)
# paf = cv2.applyColorMap(paf, cv2.COLORMAP_JET)
# plt.imshow(image)
# plt.imshow(paf, alpha=0.5)
# plt.show()
# plt.close()
##############################################
########## check labels ############
##############################################
# trianed on Imagenet dataset
if self.preprocess == 'rtpose':
img = rtpose_preprocess(img)
elif self.preprocess == 'vgg':
img = vgg_preprocess(img)
elif self.preprocess == 'inception':
img = inception_preprocess(img)
elif self.preprocess == 'ssd':
img = ssd_preprocess(img)
elif self.preprocess == 'atrous_pose':
img = dense_pose_preprocess(img)
img = torch.from_numpy(img)
heatmaps = torch.from_numpy(
heatmaps.transpose((2, 0, 1)).astype(np.float32))
heat_mask = torch.from_numpy(
heat_mask.transpose((2, 0, 1)).astype(np.float32))
pafs = torch.from_numpy(pafs.transpose((2, 0, 1)).astype(np.float32))
paf_mask = torch.from_numpy(
paf_mask.transpose((2, 0, 1)).astype(np.float32))
return img, heatmaps, heat_mask, pafs, paf_mask, self.data[idx][
'img_paths'], idx
def __getitem__(self, index):
idx = self.index_list[index]
img_path = os.path.join(self.root, self.data[idx]['img_paths'])
img = cv2.imread(img_path)
# print('img path: ', self.root, os.path.join(self.root, self.data[idx]['img_paths']), self.data[idx]['img_paths'])
img_idx = self.data[idx]['img_paths'][-16:-3]
if "COCO_val" in self.data[idx]['dataset']:
# mask_miss = cv2.imread(self.mask_dir + 'mask2014/val2014_mask_miss_' + img_idx + 'png', 0)
# val_img_path = self.mask_dir + 'mask2014/val2014_mask_miss_' + img_idx + 'png'
val_img_path = self.mask_dir + '/val2014/mask_COCO_val2014_' + img_idx + 'jpg'
mask_miss = cv2.imread(val_img_path, 0)
# print("COCO_val", val_img_path)
# mask_COCO_val2014_000000000074.jpg
elif "COCO" in self.data[idx]['dataset']:
# mask_miss = cv2.imread(self.mask_dir + 'mask2014/train2014_mask_miss_' + img_idx + 'png', 0)
# train_img_path = self.mask_dir + 'mask2014/train2014_mask_miss_' + img_idx + 'png'
train_img_path = self.mask_dir + '/train2014/mask_COCO_train2014_' + img_idx + 'jpg'
mask_miss = cv2.imread(train_img_path, 0)
# print("COCO_train", train_img_path)
# mask_COCO_train2014_000000000036.jpg
# print(self.root + 'mask2014/val2014_mask_miss_' + img_idx + 'png')
# print('get-1', img.shape, mask_miss.shape) # (375, 500, 3) (375, 500). (425, 640, 3) (425, 640)
meta_data = self.get_anno(self.data[idx])
meta_data = self.add_neck(meta_data)
meta_data, img, mask_miss = aug_scale(meta_data, img, mask_miss,
self.params_transform) # TODO
#
meta_data, img, mask_miss = aug_rotate(meta_data, img, mask_miss,
self.params_transform)
#
meta_data, img, mask_miss = aug_croppad(meta_data, img, mask_miss,
self.params_transform) # TODO
#
meta_data, img, mask_miss = aug_flip(meta_data, img, mask_miss,
self.params_transform)
meta_data = self.remove_illegal_joint(meta_data)
heat_mask, heatmaps, paf_mask, pafs = self.get_ground_truth(
meta_data, mask_miss)
# print('get-2', img.shape, mask_miss.shape) # (368, 368, 3) (369, 369)
# image preprocessing, which comply the model
# trianed on Imagenet dataset
if self.preprocess == 'rtpose':
img = rtpose_preprocess(img)
elif self.preprocess == 'vgg':
img = vgg_preprocess(img)
elif self.preprocess == 'inception':
img = inception_preprocess(img)
elif self.preprocess == 'ssd':
img = ssd_preprocess(img)
img = torch.from_numpy(img)
heatmaps = torch.from_numpy(
heatmaps.transpose((2, 0, 1)).astype(np.float32))
heat_mask = torch.from_numpy(
heat_mask.transpose((2, 0, 1)).astype(np.float32))
pafs = torch.from_numpy(pafs.transpose((2, 0, 1)).astype(np.float32))
paf_mask = torch.from_numpy(
paf_mask.transpose((2, 0, 1)).astype(np.float32))
return img, heatmaps, heat_mask, pafs, paf_mask, img_path