From 38d08ca9982c1cd812e2c7a37c6e0070f1769392 Mon Sep 17 00:00:00 2001 From: CoinCheung <867153576@qq.com> Date: Fri, 24 May 2019 17:39:00 +0800 Subject: [PATCH] modify senario where dataset has no mask provided (#814) * modify senario where dataset has no mask provided * change to check segmentation mask in the way of keypoint --- maskrcnn_benchmark/data/datasets/coco.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/maskrcnn_benchmark/data/datasets/coco.py b/maskrcnn_benchmark/data/datasets/coco.py index 313924413..cd9fc835e 100644 --- a/maskrcnn_benchmark/data/datasets/coco.py +++ b/maskrcnn_benchmark/data/datasets/coco.py @@ -79,9 +79,10 @@ def __getitem__(self, idx): classes = torch.tensor(classes) target.add_field("labels", classes) - masks = [obj["segmentation"] for obj in anno] - masks = SegmentationMask(masks, img.size, mode='poly') - target.add_field("masks", masks) + if anno and "segmentation" in anno[0]: + masks = [obj["segmentation"] for obj in anno] + masks = SegmentationMask(masks, img.size, mode='poly') + target.add_field("masks", masks) if anno and "keypoints" in anno[0]: keypoints = [obj["keypoints"] for obj in anno]