Closed kenifxyz closed 1 year ago
Hello @kenifxyz, a YOLOv8 demo would indeed be beneficial. For now, I'm providing you with code to make Norfair works for your model.
As you pointed out, it's necessary to modify the function that transforms YOLO detections into Norfair detections. The code below adjusts this function to accommodate the new model.
import logging
from typing import List
import numpy as np
import torch
from ultralytics import YOLO
from norfair import Detection, Tracker, Video, draw_boxes
# Disable logging from the ultralytics package
logger_name = "ultralytics"
logging.getLogger(logger_name).setLevel(logging.CRITICAL)
def yolov8_detections_to_norfair_detections(
yolo_detections: torch.tensor, track_points: str = "bbox" # bbox or centroid
) -> List[Detection]:
"""convert yolo_detections to Norfair's detections"""
norfair_detections: List[Detection] = []
if track_points == "centroid":
for object in yolo_detections:
detections_as_xywh = object.boxes.xywh
scores = object.boxes.conf
classes = object.boxes.cls
for detection_as_xywh, score, class_id in zip(detections_as_xywh, scores, classes):
centroid = np.array([detection_as_xywh[0].item(), detection_as_xywh[1].item()])
scores = np.array([score.item()])
norfair_detections.append(
Detection(points=centroid, scores=scores, label=int(class_id))
)
elif track_points == "bbox":
for object in yolo_detections:
detections_as_xyxy = object.boxes.xyxy
scores = object.boxes.conf
classes = object.boxes.cls
for detection_as_xyxy, score, class_id in zip(detections_as_xyxy, scores, classes):
bbox = np.array(
[
[detection_as_xyxy[0].item(), detection_as_xyxy[1].item()],
[detection_as_xyxy[2].item(), detection_as_xyxy[3].item()],
]
)
scores = np.array([score, score])
norfair_detections.append(
Detection(points=bbox, scores=scores, label=int(class_id))
)
return norfair_detections
video = Video(input_path=<video-path>)
tracker = Tracker(distance_function="iou", distance_threshold=0.7)
model = YOLO(<model-path>)
for frame in video:
yolo_detections = model(frame)
detections = yolov8_detections_to_norfair_detections(yolo_detections)
tracked_objects = tracker.update(detections=detections)
frame = draw_boxes(frame, tracked_objects)
video.write(frame)
If you have any further questions, please don't hesitate to ask.
Hey @DiegoFernandezC, that works perfectly! Thanks so much!
This works great. Able to do with yolov8. But if in case you come up with this error, "Can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first", you can solve this by adding this before line 44,
score = score.cpu().numpy()
If I may add something useful to this discussion, I found this snippet of code super handy but kind of sluggish. Therefore, I optimized a bit the function _yolov8_detections_to_norfairdetections with a few (tiny) changes. Just for reference, it went from around 2fps rate to more than 5fps. Again thanks for your support @DiegoFernandezC and hope someone else will find this code helpful.
def yolov8_detections_to_norfair_detections(
yolo_detections: torch.tensor, track_points: str = "bbox" # bbox or centroid
) -> List[Detection]:
"""convert yolo_detections to Norfair's detections"""
norfair_detections: List[Detection] = []
if track_points == "centroid":
for object in yolo_detections:
detections_as_xywh = object.boxes.xywh
scores = object.boxes.conf
classes = object.boxes.cls
norfair_detections += [
Detection(
points=np.array([detection_as_xywh[0].item(), detection_as_xywh[1].item()]),
scores=np.array([score.item()]),
label=int(class_id)
)
for detection_as_xywh, score, class_id in zip(detections_as_xywh, scores, classes)
]
elif track_points == "bbox":
for object in yolo_detections:
detections_as_xyxy = object.boxes.xyxy
scores = object.boxes.conf
classes = object.boxes.cls
norfair_detections += [
Detection(
points=np.array([
[detection_as_xyxy[0].item(), detection_as_xyxy[1].item()],
[detection_as_xyxy[2].item(), detection_as_xyxy[3].item()],
]),
scores=np.array([score, score]),
label=int(class_id)
)
for detection_as_xyxy, score, class_id in zip(detections_as_xyxy, scores, classes)
]
return norfair_detections
Would be good if this repo had a demo for how to implement YOLOv8.
I've started with modifying the code from the YOLOv7 demo, but am hitting some issues.
from ultralytics import YOLO
...# self.model = torch.hub.load("WongKinYiu/yolov7", "custom", model_path)
self.model = YOLO(model_path)
Here's how the model can be loaded with the new ultralytics repo compared to the YOLOv7 one. However, at the step of the code that converts the bounding boxes to norfair_detections, particularly here:
detections_as_xyxy = yolo_detections.xyxy[0]
I'm having some issues extracting the xyxy from the YOLOv8 results object to a valid format.
Any help would be greatly appreciated!