fix predict with class filter

This commit is contained in:
wa22 2024-05-27 13:43:44 +08:00
parent 4dabf44202
commit 9f73bc7768

View File

@ -21,7 +21,9 @@ 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)]
if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list