xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix predict with class filter

Browse Source
This commit is contained in:
wa22 2024-05-27 13:43:44 +08:00
parent 4dabf44202
commit 9f73bc7768
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ultralytics/models/yolov10/predict.py
View File

@ -21,6 +21,8 @@ class YOLOv10DetectionPredictor(DetectionPredictor):
preds = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
mask = preds[..., 4] > self.args.conf
if self.args.classes is not None:
mask = mask & (preds[..., 5:6] == torch.tensor(self.args.classes, device=preds.device).unsqueeze(0)).any(2)
preds = [p[mask[idx]] for idx, p in enumerate(preds)]