Skip to content

Commit 5b628ea

Browse files
authored
Merge pull request #319 from lyuwenyu/fix_postprocess
[fix] bbox_pred
2 parents 2b88d5d + 186d207 commit 5b628ea

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

rtdetr_pytorch/src/zoo/rtdetr/rtdetr_postprocessor.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,12 @@ def forward(self, outputs, orig_target_sizes):
4747
else:
4848
scores = F.softmax(logits)[:, :, :-1]
4949
scores, labels = scores.max(dim=-1)
50+
boxes = bbox_pred
5051
if scores.shape[1] > self.num_top_queries:
5152
scores, index = torch.topk(scores, self.num_top_queries, dim=-1)
5253
labels = torch.gather(labels, dim=1, index=index)
5354
boxes = torch.gather(boxes, dim=1, index=index.unsqueeze(-1).tile(1, 1, boxes.shape[-1]))
54-
55+
5556
# TODO for onnx export
5657
if self.deploy_mode:
5758
return labels, boxes, scores

0 commit comments

Comments
 (0)