Merge pull request #317 from sarlinpe/fix-tf-thresh

Propagate the detection threshold to box_nms
rpautrat · Feb 29, 2024 · 361799f · 361799f
2 parents d8ebb90 + 9d4f8bf
commit 361799f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/convert_to_pytorch.ipynb b/convert_to_pytorch.ipynb
diff --git a/superpoint/models/super_point.py b/superpoint/models/super_point.py
@@ -48,7 +48,8 @@ def net(image):
         prob = results['prob']
         if config['nms']:
             prob = tf.map_fn(lambda p: utils.box_nms(
-                p, config['nms'], keep_top_k=config['top_k']), prob)
+                p, config['nms'], keep_top_k=config['top_k'],
+                min_prob=config['detection_threshold']), prob)
             results['prob_nms'] = prob
         results['pred'] = tf.to_int32(tf.greater_equal(
             prob, config['detection_threshold']))