fix post processing bug

bowenc0221 · Aug 18, 2020 · 5b3dd8c · 5b3dd8c
1 parent 9c25646
commit 5b3dd8c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/segmentation/model/post_processing/instance_post_processing.py b/segmentation/model/post_processing/instance_post_processing.py
@@ -46,7 +46,7 @@ def find_instance_center(ctr_hmp, threshold=0.1, nms_kernel=3, top_k=None):
         return ctr_all
     else:
         # find top k centers.
-        top_k_scores, _ = torch.topk(torch.flatten(ctr_all), top_k)
+        top_k_scores, _ = torch.topk(torch.flatten(ctr_hmp), top_k)
         return torch.nonzero(ctr_hmp > top_k_scores[-1])