diff --git a/fcos_core/data/datasets/coco.py b/fcos_core/data/datasets/coco.py index 4d08a325703b2a06870dbba9030ffcacbefabfee..ac697e74da091e8b3fc893b4deb55b101d83dcb3 100644 --- a/fcos_core/data/datasets/coco.py +++ b/fcos_core/data/datasets/coco.py @@ -61,7 +61,7 @@ class COCODataset(torchvision.datasets.coco.CocoDetection): v: k for k, v in self.json_category_id_to_contiguous_id.items() } self.id_to_img_map = {k: v for k, v in enumerate(self.ids)} - self.transforms = transforms + self._transforms = transforms def __getitem__(self, idx): img, anno = super(COCODataset, self).__getitem__(idx) @@ -90,8 +90,8 @@ class COCODataset(torchvision.datasets.coco.CocoDetection): target = target.clip_to_image(remove_empty=True) - if self.transforms is not None: - img, target = self.transforms(img, target) + if self._transforms is not None: + img, target = self._transforms(img, target) return img, target, idx diff --git a/fcos_core/data/transforms/transforms.py b/fcos_core/data/transforms/transforms.py index 179723eeef6ca12f134a8e2922393bc47cb132ab..102c47e46b78f12f435ad73ea86b3e7e39926a1b 100644 --- a/fcos_core/data/transforms/transforms.py +++ b/fcos_core/data/transforms/transforms.py @@ -57,9 +57,12 @@ class Resize(object): def __call__(self, image, target=None): size = self.get_size(image.size) image = F.resize(image, size) - if target is None: + if isinstance(target, list): + target = [t.resize(image.size) for t in target] + elif target is None: return image - target = target.resize(image.size) + else: + target = target.resize(image.size) return image, target diff --git a/fcos_core/structures/segmentation_mask.py b/fcos_core/structures/segmentation_mask.py index aca97a9693201c0908786d7c1c9c9f006bb7f4e8..31486b63bc035a1daaaf4c3654dacf033d210546 100644 --- a/fcos_core/structures/segmentation_mask.py +++ b/fcos_core/structures/segmentation_mask.py @@ -414,7 +414,8 @@ class PolygonList(object): else: # advanced indexing on a single dimension selected_polygons = [] - if isinstance(item, torch.Tensor) and item.dtype == torch.uint8: + if isinstance(item, torch.Tensor) and \ + item.dtype == torch.uint8 or item.dtype == torch.bool: item = item.nonzero() item = item.squeeze(1) if item.numel() > 0 else item item = item.tolist()