def forward(self, boxes, scores):
return soft_nms(boxes, scores, 0.5, 0.05)
def test_indexing_works(self):
# Making sure that we can use keep for indexing
_, keep = soft_nms(self.boxes, self.scores, 0.5, 0.1)
_ = self.boxes[keep, :]
def test_gpu(self):
_ = soft_nms(self.boxes.cuda(), self.scores.cuda(), 0.5, 0.1)
def test_correct_keep(self):
_, keep = soft_nms(self.boxes, self.scores, 0.5, 0.1)
self.assertTrue(torch.allclose(keep, torch.tensor([1, 0])))