ultralytics 8.0.106
(#2736)
Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: vyskocj <whiskey1939@seznam.cz> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: triple Mu <gpu@163.com> Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
@ -17,9 +17,11 @@ class RTDETRPredictor(BasePredictor):
|
||||
results = []
|
||||
for i, bbox in enumerate(bboxes): # (300, 4)
|
||||
bbox = ops.xywh2xyxy(bbox)
|
||||
score, cls = scores[i].max(-1) # (300, )
|
||||
idx = score > self.args.conf
|
||||
pred = torch.cat([bbox, score[..., None], cls[..., None]], dim=-1)[idx] # filter
|
||||
score, cls = scores[i].max(-1, keepdim=True) # (300, 1)
|
||||
idx = score.squeeze(-1) > self.args.conf # (300, )
|
||||
if self.args.classes is not None:
|
||||
idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
|
||||
pred = torch.cat([bbox, score, cls], dim=-1)[idx] # filter
|
||||
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
|
||||
oh, ow = orig_img.shape[:2]
|
||||
if not isinstance(orig_imgs, torch.Tensor):
|
||||
|
Reference in New Issue
Block a user