demo: only keep detections above threshold

aeitel · Aug 18, 2015 · 3d079c2 · 3d079c2
1 parent 5989b41
commit 3d079c2
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tools/demo.py b/tools/demo.py
@@ -98,6 +98,9 @@ def demo(net, image_name, classes):
         cls_ind = CLASSES.index(cls)
         cls_boxes = boxes[:, 4*cls_ind:4*(cls_ind + 1)]
         cls_scores = scores[:, cls_ind]
+        keep = np.where(cls_scores >= CONF_THRESH)[0]
+        cls_boxes = cls_boxes[keep, :]
+        cls_scores = cls_scores[keep]
         dets = np.hstack((cls_boxes,
                           cls_scores[:, np.newaxis])).astype(np.float32)
         keep = nms(dets, NMS_THRESH)