modify senario where dataset has no mask provided (facebookresearch#814)

* modify senario where dataset has no mask provided * change to check segmentation mask in the way of keypoint
miguelvr · May 24, 2019 · 38d08ca · 38d08ca
1 parent 16b8b09
commit 38d08ca
Showing 1 changed file with 4 additions and 3 deletions.
diff --git 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]