diff --git a/references/detection/coco_utils.py b/references/detection/coco_utils.py index 44b917a6ec6..e6eaaa25297 100644 --- a/references/detection/coco_utils.py +++ b/references/detection/coco_utils.py @@ -26,6 +26,9 @@ def convert_coco_poly_to_mask(segmentations, height, width): class ConvertCocoPolysToMask: + def __init__(self, with_masks=False): + self.with_masks = with_masks + def __call__(self, image, target): w, h = image.size @@ -45,8 +48,9 @@ def __call__(self, image, target): classes = [obj["category_id"] for obj in anno] classes = torch.tensor(classes, dtype=torch.int64) - segmentations = [obj["segmentation"] for obj in anno] - masks = convert_coco_poly_to_mask(segmentations, h, w) + if self.with_masks: + segmentations = [obj["segmentation"] for obj in anno] + masks = convert_coco_poly_to_mask(segmentations, h, w) keypoints = None if anno and "keypoints" in anno[0]: @@ -59,14 +63,16 @@ def __call__(self, image, target): keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) boxes = boxes[keep] classes = classes[keep] - masks = masks[keep] + if self.with_masks: + masks = masks[keep] if keypoints is not None: keypoints = keypoints[keep] target = {} target["boxes"] = boxes target["labels"] = classes - target["masks"] = masks + if self.with_masks: + target["masks"] = masks target["image_id"] = image_id if keypoints is not None: target["keypoints"] = keypoints @@ -218,8 +224,7 @@ def get_coco(root, image_set, transforms, mode="instances", use_v2=False, with_m target_keys += ["masks"] dataset = wrap_dataset_for_transforms_v2(dataset, target_keys=target_keys) else: - # TODO: handle with_masks for V1? - t = [ConvertCocoPolysToMask()] + t = [ConvertCocoPolysToMask(with_masks=with_masks)] if transforms is not None: t.append(transforms) transforms = T.Compose(t) diff --git a/test/test_references_detection.py b/test/test_references_detection.py new file mode 100644 index 00000000000..457b485b062 --- /dev/null +++ b/test/test_references_detection.py @@ -0,0 +1,35 @@ +import sys +from pathlib import Path + +from PIL import Image + +sys.path.insert(0, str(Path(__file__).parents[1] / "references" / "detection")) + +import coco_utils + +from coco_utils import ConvertCocoPolysToMask + + +def test_convert_coco_polys_to_mask_skips_masks_when_not_requested(monkeypatch): + def fail_if_called(*args, **kwargs): + raise AssertionError("mask decoding should be skipped") + + monkeypatch.setattr(coco_utils, "convert_coco_poly_to_mask", fail_if_called) + + _, target = ConvertCocoPolysToMask()( + Image.new("RGB", (4, 4)), + { + "image_id": 1, + "annotations": [ + { + "bbox": [0, 0, 2, 2], + "category_id": 1, + "iscrowd": 0, + "area": 4, + "segmentation": [], + } + ], + }, + ) + + assert "masks" not in target