fix BitMasks indexing

Summary: fix #3591

Differential Revision: D31771140

fbshipit-source-id: cefca0b49e9b061fe45be6bc6d34ed8e1853c7fa

Author: ppwwyyxx

detectron2/structures/masks.py

@@ -130,7 +130,7 @@ def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "BitMasks":
         subject to Pytorch's indexing semantics.
         """
         if isinstance(item, int):
-            return BitMasks(self.tensor[item].view(1, -1))
+            return BitMasks(self.tensor[item].unsqueeze(0))
         m = self.tensor[item]
         assert m.dim() == 3, "Indexing on BitMasks with {} returns a tensor with shape {}!".format(
             item, m.shape

docs/requirements.txt

@@ -15,7 +15,7 @@ cloudpickle
 Pillow
 future
 git+git://github.com/facebookresearch/fvcore.git
-https://download.pytorch.org/whl/cpu/torch-1.8.0%2Bcpu-cp37-cp37m-linux_x86_64.whl
+https://download.pytorch.org/whl/cpu/torch-1.8.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
 https://download.pytorch.org/whl/cpu/torchvision-0.9.1%2Bcpu-cp37-cp37m-linux_x86_64.whl
 omegaconf>=2.1.0.dev24
 hydra-core>=1.1.0.dev5

tests/structures/test_masks.py

@@ -42,6 +42,12 @@ def test_from_empty_polygons(self):
         masks = BitMasks.from_polygon_masks([], 100, 100)
         self.assertEqual(masks.tensor.shape, (0, 100, 100))
 
+    def test_getitem(self):
+        masks = BitMasks(torch.ones(3, 10, 10))
+        self.assertEqual(masks[1].tensor.shape, (1, 10, 10))
+        self.assertEqual(masks[1:3].tensor.shape, (2, 10, 10))
+        self.assertEqual(masks[torch.tensor([True, False, False])].tensor.shape, (1, 10, 10))
+
 
 if __name__ == "__main__":
     unittest.main()