def convert_coco_poly_to_mask(segmentations: Sequence, height: int,
width: int) -> NamedTensorTuple:
"""
:param segmentations:
:type segmentations:
:param height:
:type height:
:param width:
:type width:
:return:
:rtype:"""
masks = []
for polygons in segmentations:
rles = coco_mask.frPyObjects(polygons, height, width)
mask = coco_mask.decode(rles)
if len(mask.shape) < 3:
mask = mask[..., None]
mask = torch.as_tensor(mask, dtype=torch.uint8)
mask = mask.any(dim=2)
masks.append(mask)
if masks:
masks = torch.stack(masks, dim=0)
else:
masks = torch.zeros((0, height, width), dtype=torch.uint8)
return NamedTensorTuple(masks=masks)
def __getitem__(self, idx):
img, target = super().__getitem__(idx)
image_id = self.ids[idx]
target = NamedTensorTuple(image_id=image_id, annotations=target)
if self._transforms is not None:
img, target = self._transforms(img, target)
return img, target
def __getitem__(self, index):
image_id = self._ids[index]
boxes, labels = self._get_annotation(image_id)
image = self._read_image(image_id)
if self._img_transforms:
image, boxes, labels = self._img_transforms(image, boxes, labels)
if self._annotation_transforms:
boxes, labels = self._annotation_transforms(boxes, labels)
return image, NamedTensorTuple(boxes=boxes, labels=labels), index
def __call__(self, batch: Iterable) -> Tuple:
transposed_batch = list(zip(*batch))
images = default_collate(transposed_batch[0])
img_ids = default_collate(transposed_batch[2])
if self.wrap:
list_targets = transposed_batch[1]
targets = {
key: default_collate([d[key] for d in list_targets])
for key in list_targets[0]
}
targets = NamedTensorTuple(**targets)
else:
targets = None
return images, targets, img_ids