소스 검색

fix a bug in rtdetr post-process

yjh0410 1 년 전
부모
커밋
31bc9e7dea
1개의 변경된 파일1개의 추가작업 그리고 2개의 파일을 삭제
  1. 1 2
      models/detectors/rtpdetr/rtpdetr.py

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

@@ -154,12 +154,11 @@ class RT_PDETR(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]