introduce thresholds on init to fix onnx

Author: datumbox

test/test_onnx.py

@@ -184,7 +184,7 @@ def _init_test_generalized_rcnn_transform(self):
         transform = GeneralizedRCNNTransform(min_size, max_size, image_mean, image_std)
         return transform
 
-    def _init_test_rpn(self):
+    def _init_test_rpn(self, score_threshold=0.0):
         anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
         aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
         rpn_anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios)
@@ -197,7 +197,7 @@ def _init_test_rpn(self):
         rpn_pre_nms_top_n = dict(training=2000, testing=1000)
         rpn_post_nms_top_n = dict(training=2000, testing=1000)
         rpn_nms_thresh = 0.7
-        rpn_score_thresh = 0.0
+        rpn_score_thresh = score_threshold
 
         rpn = RegionProposalNetwork(
             rpn_anchor_generator, rpn_head,
@@ -260,7 +260,7 @@ def test_rpn(self):
         class RPNModule(torch.nn.Module):
             def __init__(self_module):
                 super(RPNModule, self_module).__init__()
-                self_module.rpn = self._init_test_rpn()
+                self_module.rpn = self._init_test_rpn(0.5)
 
             def forward(self_module, images, features):
                 images = ImageList(images, [i.shape[-2:] for i in images])

torchvision/models/detection/rpn.py

@@ -264,7 +264,8 @@ def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_
             boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]
 
             # remove low scoring boxes
-            keep = torch.where(scores > self.score_thresh)[0]
+            # use >= for Backwards compatibility
+            keep = torch.where(scores >= self.score_thresh)[0]
             boxes, scores, lvl = boxes[keep], scores[keep], lvl[keep]
 
             # non-maximum suppression, independently done per level