Merge pull request #12 from sony/fix_empty

fix all scores are invalid

sony_custom_layers/pytorch/object_detection/nms.py

@@ -159,10 +159,12 @@ def _image_multiclass_nms(boxes: Tensor, scores: Tensor, score_threshold: float,
 
     """
     x = _convert_inputs(boxes, scores, score_threshold)
+    out = torch.zeros(max_detections, 6, device=boxes.device)
+    if x.size(0) == 0:
+        return out, 0
     idxs = _nms_with_class_offsets(x, iou_threshold=iou_threshold)
     idxs = idxs[:max_detections]
     valid_dets = idxs.numel()
-    out = torch.zeros(max_detections, 6, device=boxes.device)
     out[:valid_dets] = x[idxs]
     return out, valid_dets
 

sony_custom_layers/pytorch/tests/object_detection/test_multiclass_nms.py

@@ -127,6 +127,14 @@ def test_image_multiclass_nms(self, mocker, max_detections, mock_tv_op):
         assert torch.all(ret[:, 5][exp_valid_dets:] == 0)
         assert ret_valid_dets == exp_valid_dets
 
+    def test_empty_tensors(self):
+        # empty inputs
+        ret = nms.multiclass_nms(torch.rand(1, 0, 4), torch.rand(1, 0, 10), 0.55, 0.6, 50)
+        assert ret.n_valid[0] == 0 and ret.boxes.size(1) == 50
+        # no valid scores
+        ret = nms.multiclass_nms(torch.rand(1, 100, 4), torch.rand(1, 100, 20) / 2, 0.55, 0.6, 50)
+        assert ret.n_valid[0] == 0 and ret.boxes.size(1) == 50
+
     def test_batch_multiclass_nms(self, mocker):
         input_boxes, input_scores = self._generate_random_inputs(batch=3, n_boxes=20, n_classes=10)
         max_dets = 5