set gt boxes in both cases

IDEA-Research · Feb 26, 2023 · a526c23 · a526c23
1 parent 21e98f9
commit a526c23
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/maskdino/data/dataset_mappers/mask_former_semantic_dataset_mapper.py b/maskdino/data/dataset_mappers/mask_former_semantic_dataset_mapper.py
@@ -11,7 +11,7 @@
 from detectron2.data import detection_utils as utils
 from detectron2.data import transforms as T
 from detectron2.projects.point_rend import ColorAugSSDTransform
-from detectron2.structures import BitMasks, Instances
+from detectron2.structures import BitMasks, Boxes, Instances
 
 __all__ = ["MaskFormerSemanticDatasetMapper"]
 
@@ -173,12 +173,13 @@ def __call__(self, dataset_dict):
             if len(masks) == 0:
                 # Some image does not have annotation (all ignored)
                 instances.gt_masks = torch.zeros((0, sem_seg_gt.shape[-2], sem_seg_gt.shape[-1]))
+                instances.gt_boxes = Boxes(torch.zeros((0,4)))
             else:
                 masks = BitMasks(
                     torch.stack([torch.from_numpy(np.ascontiguousarray(x.copy())) for x in masks])
                 )
                 instances.gt_masks = masks.tensor
-            instances.gt_boxes = masks.get_bounding_boxes()
+                instances.gt_boxes = masks.get_bounding_boxes()
 
             dataset_dict["instances"] = instances