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
Exemple #2
0
    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