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 get_outputs(multiplier, img, model, preprocess):
"""Computes the averaged heatmap and paf for the given image
:param multiplier:
:param origImg: numpy array, the image being processed
:param model: pytorch model
:returns: numpy arrays, the averaged paf and heatmap
"""
heatmap_avg = np.zeros((img.shape[0], img.shape[1], 19))
paf_avg = np.zeros((img.shape[0], img.shape[1], 38))
max_scale = multiplier[-1]
max_size = max_scale * img.shape[0]
# padding
max_cropped, _, _ = im_transform.crop_with_factor(img,
max_size,
factor=8,
is_ceil=True)
batch_images = np.zeros(
(len(multiplier), 3, max_cropped.shape[0], max_cropped.shape[1]))
for m in range(len(multiplier)):
scale = multiplier[m]
inp_size = scale * img.shape[0]
# padding
im_croped, im_scale, real_shape = im_transform.crop_with_factor(
img, inp_size, factor=8, is_ceil=True)
if preprocess == 'rtpose':
im_data = rtpose_preprocess(im_croped)
elif preprocess == 'vgg':
im_data = vgg_preprocess(im_croped)
elif preprocess == 'inception':
im_data = inception_preprocess(im_croped)
elif preprocess == 'ssd':
im_data = ssd_preprocess(im_croped)
batch_images[m, :, :im_data.shape[1], :im_data.shape[2]] = im_data
# several scales as a batch
batch_var = torch.from_numpy(batch_images).cuda().float()
predicted_outputs, _ = model(batch_var)
output1, output2 = predicted_outputs[-2], predicted_outputs[-1]
heatmaps = output2.cpu().data.numpy().transpose(0, 2, 3, 1)
pafs = output1.cpu().data.numpy().transpose(0, 2, 3, 1)
for m in range(len(multiplier)):
scale = multiplier[m]
inp_size = scale * img.shape[0]
# padding
im_cropped, im_scale, real_shape = im_transform.crop_with_factor(
img, inp_size, factor=8, is_ceil=True)
heatmap = heatmaps[m, :int(im_cropped.shape[0] /
8), :int(im_cropped.shape[1] / 8), :]
heatmap = cv2.resize(heatmap,
None,
fx=8,
fy=8,
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[0:real_shape[0], 0:real_shape[1], :]
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = pafs[m, :int(im_cropped.shape[0] / 8), :int(im_cropped.shape[1] /
8), :]
paf = cv2.resize(paf, None, fx=8, fy=8, interpolation=cv2.INTER_CUBIC)
paf = paf[0:real_shape[0], 0:real_shape[1], :]
paf = cv2.resize(paf, (img.shape[1], img.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return paf_avg, heatmap_avg
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
