freddyaboulton HF staff commited on
Commit
e38752e
1 Parent(s): 8e0c0f6
Files changed (1) hide show
  1. inference.py +1 -1
inference.py CHANGED
@@ -68,7 +68,7 @@ class YOLOv10:
68
  return [], [], []
69
 
70
  # Get the class with the highest confidence
71
- class_ids = np.argmax(predictions[:, 4:], axis=1)
72
 
73
  # Get bounding boxes for each object
74
  boxes = self.extract_boxes(predictions)
 
68
  return [], [], []
69
 
70
  # Get the class with the highest confidence
71
+ class_ids = predictions[:, 5].astype(int)
72
 
73
  # Get bounding boxes for each object
74
  boxes = self.extract_boxes(predictions)