HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.85 YOLOv10 max_det arg fix (#15917)

Browse source 
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Francesco Mattioli 2024-09-01 02:04:22 +02:00 • committed by GitHub
parent 0546c08102
commit ea13dc6208
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/utils/ops.py
View file

@ -218,7 +218,7 @@ def non_max_suppression(
classes = torch.tensor(classes, device=prediction.device)
if prediction.shape[-1] == 6: # end-to-end model (BNC, i.e. 1,300,6)
output = [pred[pred[:, 4] > conf_thres] for pred in prediction]
output = [pred[pred[:, 4] > conf_thres][:max_det] for pred in prediction]
if classes is not None:
output = [pred[(pred[:, 5:6] == classes).any(1)] for pred in output]
return output