提交 0ea6ea06 编写于 作者: H Hao Chen

compatibility for torchvision 0.4.0

上级 bb6e97bc
...@@ -61,7 +61,7 @@ class COCODataset(torchvision.datasets.coco.CocoDetection): ...@@ -61,7 +61,7 @@ class COCODataset(torchvision.datasets.coco.CocoDetection):
v: k for k, v in self.json_category_id_to_contiguous_id.items() 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.id_to_img_map = {k: v for k, v in enumerate(self.ids)}
self.transforms = transforms self._transforms = transforms
def __getitem__(self, idx): def __getitem__(self, idx):
img, anno = super(COCODataset, self).__getitem__(idx) img, anno = super(COCODataset, self).__getitem__(idx)
...@@ -90,8 +90,8 @@ class COCODataset(torchvision.datasets.coco.CocoDetection): ...@@ -90,8 +90,8 @@ class COCODataset(torchvision.datasets.coco.CocoDetection):
target = target.clip_to_image(remove_empty=True) target = target.clip_to_image(remove_empty=True)
if self.transforms is not None: if self._transforms is not None:
img, target = self.transforms(img, target) img, target = self._transforms(img, target)
return img, target, idx return img, target, idx
......
...@@ -57,9 +57,12 @@ class Resize(object): ...@@ -57,9 +57,12 @@ class Resize(object):
def __call__(self, image, target=None): def __call__(self, image, target=None):
size = self.get_size(image.size) size = self.get_size(image.size)
image = F.resize(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 return image
target = target.resize(image.size) else:
target = target.resize(image.size)
return image, target return image, target
......
...@@ -414,7 +414,8 @@ class PolygonList(object): ...@@ -414,7 +414,8 @@ class PolygonList(object):
else: else:
# advanced indexing on a single dimension # advanced indexing on a single dimension
selected_polygons = [] 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.nonzero()
item = item.squeeze(1) if item.numel() > 0 else item item = item.squeeze(1) if item.numel() > 0 else item
item = item.tolist() item = item.tolist()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册