new interface of single_proir

mmdet/core/anchor/anchor_generator.py

@@ -225,13 +225,45 @@ def grid_priors(self, featmap_sizes, device='cuda'):
         multi_level_anchors = []
         for i in range(self.num_levels):
             anchors = self.single_level_grid_priors(
-                self.base_anchors[i].to(device),
-                featmap_sizes[i],
-                self.strides[i],
-                device=device)
+                featmap_sizes[i], level_idx=i, device=device)
             multi_level_anchors.append(anchors)
         return multi_level_anchors
 
+    def single_level_grid_priors(self, featmap_size, level_idx, device='cuda'):
+        """Generate grid anchors of a single level.
+
+        Note:
+            This function is usually called by method ``self.grid_priors``.
+
+        Args:
+            featmap_size (tuple[int]): Size of the feature maps.
+            level_idx (int): The index of corresponding feature map level.
+            device (str, optional): The device the tensor will be put on.
+                Defaults to 'cuda'.
+
+        Returns:
+            torch.Tensor: Anchors in the overall feature maps.
+        """
+
+        base_anchors = self.base_anchors[level_idx]
+        feat_h, feat_w = featmap_size
+        stride_w, stride_h = self.strides[level_idx]
+        shift_x = torch.arange(0, feat_w, device=device) * stride_w
+        shift_y = torch.arange(0, feat_h, device=device) * stride_h
+
+        shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
+        shifts = torch.stack([shift_xx, shift_yy, shift_xx, shift_yy], dim=-1)
+        shifts = shifts.type_as(base_anchors)
+        # first feat_w elements correspond to the first row of shifts
+        # add A anchors (1, A, 4) to K shifts (K, 1, 4) to get
+        # shifted anchors (K, A, 4), reshape to (K*A, 4)
+
+        all_anchors = base_anchors[None, :, :] + shifts[:, None, :]
+        all_anchors = all_anchors.view(-1, 4)
+        # first A rows correspond to A anchors of (0, 0) in feature map,
+        # then (0, 1), (0, 2), ...
+        return all_anchors
+
     def grid_anchors(self, featmap_sizes, device='cuda'):
 
         warnings.warn('``grid_anchors`` would be deprecated soon. Please use '
@@ -243,21 +275,6 @@ def single_level_grid_anchors(self,
                                   featmap_size,
                                   stride=(16, 16),
                                   device='cuda'):
-        warnings.warn(
-            '``single_level_grid_anchors`` would be deprecated soon. '
-            'Please use ``single_level_grid_priors`` ')
-
-        return self.single_level_grid_priors(
-            base_anchors=base_anchors,
-            featmap_size=featmap_size,
-            stride=stride,
-            device=device)
-
-    def single_level_grid_priors(self,
-                                 base_anchors,
-                                 featmap_size,
-                                 stride=(16, 16),
-                                 device='cuda'):
         """Generate grid anchors of a single level.
 
         Note:
@@ -275,6 +292,10 @@ def single_level_grid_priors(self,
             torch.Tensor: Anchors in the overall feature maps.
         """
         # keep as Tensor, so that we can covert to ONNX correctly
+        warnings.warn(
+            '``single_level_grid_anchors`` would be deprecated soon. '
+            'Please use ``single_level_grid_priors`` ')
+
         feat_h, feat_w = featmap_size
         stride_w, stride_h = stride
         shift_x = torch.arange(0, feat_w, device=device) * stride_w
@@ -353,15 +374,15 @@ def single_level_valid_flags(self,
         return valid
 
     def sparse_priors(self,
-                      prior_indexs,
+                      prior_idx,
                       featmap_size,
                       level_idx,
                       dtype,
                       device='cuda'):
-        """Generate sparse anchors according to the ``prior_indexs``.
+        """Generate sparse anchors according to the ``prior_idx``.
 
         Args:
-            prior_indexs (Tensor): The index of corresponding anchors
+            prior_idx (Tensor): The index of corresponding anchors
                 in the feature map.
             featmap_size (tuple[int]): feature map size arrange as (h, w).
             level_idx (int): The level index of corresponding feature
@@ -371,15 +392,15 @@ def sparse_priors(self,
                 located.
         Returns:
             Tensor: Anchor with shape (N, 4), N should be equal to
-                the length of ``prior_indexs``.
+                the length of ``prior_idx``.
         """
 
         height, width = featmap_size
         num_base_anchors = self.num_base_anchors[level_idx]
-        base_anchor_id = prior_indexs % num_base_anchors
-        x = (prior_indexs //
+        base_anchor_id = prior_idx % num_base_anchors
+        x = (prior_idx //
              num_base_anchors) % width * self.strides[level_idx][0]
-        y = (prior_indexs // width //
+        y = (prior_idx // width //
              num_base_anchors) % height * self.strides[level_idx][1]
         priors = torch.stack([x, y, x, y], 1).to(
             dtype) + self.base_anchors[level_idx][base_anchor_id, :].to(device)

mmdet/core/anchor/point_generator.py

@@ -93,15 +93,15 @@ def grid_priors(self, featmap_sizes, device='cuda', with_stride=False):
         for i in range(self.num_levels):
             priors = self.single_level_grid_priors(
                 featmap_sizes[i],
-                stride=self.strides[i],
+                level_idx=i,
                 device=device,
                 with_stride=with_stride)
             multi_level_priors.append(priors)
         return multi_level_priors
 
     def single_level_grid_priors(self,
                                  featmap_size,
-                                 stride=None,
+                                 level_idx,
                                  device='cuda',
                                  with_stride=False):
         """Generate grid Points of a single level.
@@ -112,8 +112,7 @@ def single_level_grid_priors(self,
         Args:
             featmap_size (tuple[int]): Size of the feature maps, arrange as
                 (h, w).
-            stride (int, tuple[int], optional): Stride of the feature map
-                in order (w, h). Defaults to None.
+            level_idx (int): The index of corresponding feature map level.
             device (str, optional): The device the tensor will be put on.
                 Defaults to 'cuda'.
             with_stride (bool): Concatenate the stride to the last dimension
@@ -130,7 +129,7 @@ def single_level_grid_priors(self,
             (coord_x, coord_y, stride_w, stride_h).
         """
         feat_h, feat_w = featmap_size
-        stride_w, stride_h = stride
+        stride_w, stride_h = self.strides[level_idx]
         shift_x = (torch.arange(0., feat_w, device=device) +
                    self.offset) * stride_w
         shift_y = (torch.arange(0., feat_h, device=device) +
@@ -204,15 +203,15 @@ def single_level_valid_flags(self,
         return valid
 
     def sparse_priors(self,
-                      prior_indexs,
+                      prior_idx,
                       featmap_size,
                       level_idx,
                       dtype,
                       device='cuda'):
-        """Generate sparse points according to the ``prior_indexs``.
+        """Generate sparse points according to the ``prior_idx``.
 
         Args:
-            prior_indexs (Tensor): The index of corresponding anchors
+            prior_idx (Tensor): The index of corresponding anchors
                 in the feature map.
             featmap_size (tuple[int]): feature map size arrange as (w, h).
             level_idx (int): The level index of corresponding feature
@@ -222,12 +221,12 @@ def sparse_priors(self,
                 located.
         Returns:
             Tensor: Anchor with shape (N, 2), N should be equal to
-            the length of ``prior_indexs``. And last dimension
+            the length of ``prior_idx``. And last dimension
             2 represent (coord_x, coord_y).
         """
         height, width = featmap_size
-        x = (prior_indexs % width + self.offset) * self.strides[level_idx][0]
-        y = ((prior_indexs // width) % height +
+        x = (prior_idx % width + self.offset) * self.strides[level_idx][0]
+        y = ((prior_idx // width) % height +
              self.offset) * self.strides[level_idx][1]
         prioris = torch.stack([x, y], 1).to(dtype)
         prioris = prioris.to(device)

tests/test_utils/test_anchor.py

@@ -59,6 +59,19 @@ def test_standard_points_generator():
     assert (priors_half_offset[1][0] - priors[1][0]).sum() == 10 * 0.5 * 2
 
 
+def test_sparse_prior():
+    from mmdet.core.anchor import MlvlPointGenerator
+    mlvl_points = MlvlPointGenerator(strides=[4, 10], offset=0)
+    prior_indexs = torch.Tensor([[0, 2, 4, 5, 6, 9]], dtype=torch.Long)
+
+    featmap_sizes = [(3, 5), (6, 4)]
+    grid_anchors = mlvl_points.grid_priors(
+        featmap_sizes=featmap_sizes, with_stride=False)
+    sparse_prior = mlvl_points.sparse_priors(
+        prior_indexs=prior_indexs, featmap_size=featmap_sizes[0], level_idx=0)
+    assert (sparse_prior == grid_anchors[0][prior_indexs]).all()
+
+
 def test_standard_anchor_generator():
     from mmdet.core.anchor import build_anchor_generator
     anchor_generator_cfg = dict(
@@ -71,10 +84,6 @@ def test_standard_anchor_generator():
     assert anchor_generator is not None
 
 
-def test_sparse_prior():
-    pass
-
-
 def test_strides():
     from mmdet.core import AnchorGenerator
     # Square strides