Diego Fernandez commited on
Commit
57a060c
1 Parent(s): 040e4ee

fix: paths with centroids

Browse files
Files changed (1) hide show
  1. custom_models/yolo.py +7 -2
custom_models/yolo.py CHANGED
@@ -51,8 +51,13 @@ def yolo_detections_to_norfair_detections(
51
  if track_points == "centroid":
52
  detections_as_xywh = yolo_detections.xywh[0]
53
  for detection_as_xywh in detections_as_xywh:
54
- centroid = np.array([detection_as_xywh[0].item(), detection_as_xywh[1].item()])
55
- scores = np.array([detection_as_xywh[4].item()])
 
 
 
 
 
56
  norfair_detections.append(Detection(points=centroid, scores=scores))
57
  elif track_points == "bbox":
58
  detections_as_xyxy = yolo_detections.xyxy[0]
51
  if track_points == "centroid":
52
  detections_as_xywh = yolo_detections.xywh[0]
53
  for detection_as_xywh in detections_as_xywh:
54
+ centroid = np.array(
55
+ [
56
+ [detection_as_xywh[0].item(), detection_as_xywh[1].item()],
57
+ [detection_as_xywh[0].item(), detection_as_xywh[1].item()],
58
+ ]
59
+ )
60
+ scores = np.array([detection_as_xywh[4].item(), detection_as_xywh[4].item()])
61
  norfair_detections.append(Detection(points=centroid, scores=scores))
62
  elif track_points == "bbox":
63
  detections_as_xyxy = yolo_detections.xyxy[0]