浏览代码

fix a bug in rtdetr post-process

yjh0410 1 年之前
父节点
当前提交
031a711c57
共有 1 个文件被更改,包括 1 次插入2 次删除
  1. 1 2
      models/detectors/rtdetr/rtdetr.py

+ 1 - 2
models/detectors/rtdetr/rtdetr.py

@@ -81,12 +81,11 @@ class RT_DETR(nn.Module):
 
             # Filter out the proposals with low confidence score
             keep_idxs = topk_scores > self.conf_thresh
-            scores = topk_scores[keep_idxs]
+            topk_scores = topk_scores[keep_idxs]
             topk_idxs = topk_idxs[keep_idxs]
             topk_box_idxs = torch.div(topk_idxs, self.num_classes, rounding_mode='floor')
 
             ## Top-k results
-            topk_scores = predicted_prob[:self.num_topk]
             topk_labels = topk_idxs % self.num_classes
             topk_bboxes = box_pred[topk_box_idxs]