def __getitem__(self, idx):
ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
img_path = Path(self.img_folder) / ann_info['file_name'].replace('.png', '.jpg')
ann_path = Path(self.ann_folder) / ann_info['file_name']
img = Image.open(img_path).convert('RGB')
w, h = img.size
if "segments_info" in ann_info:
masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
masks = rgb2id(masks)
ids = np.array([ann['id'] for ann in ann_info['segments_info']])
masks = masks == ids[:, None, None]
masks = torch.as_tensor(masks, dtype=torch.uint8)
labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
target = {}
target['image_id'] = torch.tensor([ann_info['image_id'] if "image_id" in ann_info else ann_info["id"]])
if self.return_masks:
target['masks'] = masks
target['labels'] = labels
target["boxes"] = masks_to_boxes(masks)
target['size'] = torch.as_tensor([int(h), int(w)])
target['orig_size'] = torch.as_tensor([int(h), int(w)])
if "segments_info" in ann_info:
for name in ['iscrowd', 'area']:
target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __getitem__(self, idx):
# ann_info = self.coco['annotations'][idx] if "annotations" in self.coco else self.coco['images'][idx]
filename = self.file_names[idx]
img_path = osp.join(self.img_folder, filename + self.postfix_img)
ann_path = osp.join(self.ann_folder, filename + self.postfix_ann)
ann_info = list(read_json(ann_path).values())[0]
# import ipdb
# ipdb.set_trace()
img = Image.open(img_path).convert('RGB')
w, h = img.size
masks = []
labels = []
# import ipdb; ipdb.set_trace()
mask = np.zeros([h,w]).astype('uint8')
for ann in ann_info['annotation'][:100]:
mask *= 0
cnt = np.array(ann['segmentation']).astype(int)
cv2.drawContours(mask, [cnt], -1, 1, -1)
masks.append(mask.copy())
labels.append(self.cat2id(ann['tags'][0]))
# if "segments_info" in ann_info:
# # masks = np.asarray(Image.open(ann_path), dtype=np.uint32)
# masks = mmcv.imread(ann_path)
# masks = rgb2id(masks)
# ids = np.array([ann['id'] for ann in ann_info['segments_info']])
# # import ipdb; ipdb.set_trace()
# assert len(ids) == len(np.unique(masks)), f"{len(ids)},{len(np.unique(masks))}"
# masks = masks == ids[:, None, None]
# masks = torch.as_tensor(masks, dtype=torch.uint8)
# labels = torch.tensor([ann['category_id'] for ann in ann_info['segments_info']], dtype=torch.int64)
# assert len(labels) == len(ids)
target = {}
target['image_id'] = torch.tensor([idx])
masks = np.array(masks)
masks = torch.as_tensor(masks, dtype=torch.uint8)
target['labels'] = torch.tensor(labels, dtype=torch.int64)
if self.return_masks:
target['masks'] = masks
target["boxes"] = masks_to_boxes(masks)
target['size'] = torch.as_tensor([int(h), int(w)])
target['orig_size'] = torch.as_tensor([int(h), int(w)])
if "segments_info" in ann_info:
for name in ['iscrowd', 'area']:
target[name] = torch.tensor([ann[name] for ann in ann_info['segments_info']])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target