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