tryolabs / norfair

Lightweight Python library for adding real-time multi-object tracking to any detector.

Home Page:https://tryolabs.github.io/norfair/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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.