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

fix issue 13587 : 'dict' object has no attribute 'shape'

Browse Source
This commit is contained in:
qianyue76 2024-06-17 18:08:54 +08:00
parent 2c36ab0f10
commit e7c22252a3
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ultralytics/utils/ops.py
View File

@ -210,6 +210,10 @@ def non_max_suppression(
# Checks
assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
if isinstance(prediction, dict) and 'one2one' in prediction: # model.pt, output = {"one2many": torch.Tensor, "one2one": torch.Tensor }
prediction = prediction['one2one']
if isinstance(prediction, (list, tuple)): # YOLOv8 model in validation model, output = (inference_out, loss_out)
prediction = prediction[0] # select only inference output