assert statement breaks script unnecessarily
Arno989 opened this issue · comments
In the yolov7 demo, the yolo_detections_to_norfair_detections
function makes a (2, 1) np array for the scores, although i adapted this to yolov8 this function itself gives no problems
In tracker.py on line 641 in the hit
funtion there's assert len(detection.scores.shape) == 1
, which fails on this array.
If I comment out this line in the library files my script starts working again.
Now I assume one of the two, either the assert statement or the demo function, is wrong. But I cannot for the life of me figure out how to make my script work without commenting out the assert line in the library itself, which is less than ideal.
Complete script:
#%%
import cv2, torch
import matplotlib.pyplot as plt
import numpy as np
from IPython.display import display, clear_output
from typing import List
from ultralytics import YOLO
from norfair import Detection, Tracker, Video, draw_boxes, draw_points, print_objects_as_table
#%%
in_file = "traffic-short.mp4"
save_frames = False
n_frames = 10
show_frames = False
# filter_classes = False
# class_filter = 2 # 2 = car, 0 = person
# Model
model = YOLO('yolov8s.pt')
#%%
def yolo_detections_to_norfair_detections(yolo_detections: torch.tensor) -> List[Detection]:
"""convert detections_as_xyxy to norfair detections"""
norfair_detections: List[Detection] = []
# detections_as_xyxy = yolo_detections.xyxy
detections_as_xyxy = yolo_detections.xyxy
for i, detection_as_xyxy in enumerate(detections_as_xyxy):
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(
[yolo_detections[i].conf, yolo_detections[i].conf]
)
norfair_detections.append(
Detection(
points=bbox, scores=scores, label=int(yolo_detections[i].cls)
)
)
return norfair_detections
#%%
# Init tracker
tracker = Tracker(
distance_function="euclidean",
distance_threshold=100,
hit_counter_max=1000,
initialization_delay=3,
detection_threshold=0.3,
past_detections_length=10)
# Open video
video = Video(input_path=in_file)
# Init plots
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
# Iterate over frames
f = 0
for frame in video:
if f < n_frames:
# Move model to gpu
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
# Predict and move back to cpu
results = model(frame)[0].cpu().numpy().boxes
# Convert to norfair detections
detections = yolo_detections_to_norfair_detections(results)
# Update tracker
tracked_objects = tracker.update(detections=detections)
# Draw tracked objects's bounding boxes
frame = draw_boxes(frame, tracked_objects)
if show_frames:
ax.imshow(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
ax.axis("off") # turns off axes
ax.axis("tight") # gets rid of white border
ax.axis("image") # square up the image instead of filling the "figure" space
clear_output(wait = True)
print("Frame:", str(f+1) + ", Objects:", len(detections))
display(fig)
if save_frames:
video.write(frame)
f += 1
else:
break
# plt.show()
# clear_output()
print_objects_as_table(tracked_objects)
Only commented line in the library is ~/miniconda3/envs/torch/lib/python3.1/site-packages/norfair/tracker.py
line #641
Hi @Arno989, nice to see you've integrated YOLOv8!
Simply changing the way you're defining the scores fixes your problem. This should work:
scores = np.array(
[yolo_detections[i].conf.item(), yolo_detections[i].conf.item()]
)
This makes the scores
array have a (2,)
shape instead of having a (2,1)
shape as before.
We will review this internally since it doesn't seem to be intuitive.
Works perfectly, thank you!
I did notice the .item()
everywhere but I guess I didn't think too much (or enough) about why it was there.