File size: 4,479 Bytes
c245188 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import cv2
import numpy as np
from sort import Sort
# SORT tracking algorithm initialization
sort_max_age = 50
sort_min_hits = 2
sort_iou_thresh = 0.2
sort_tracker = Sort(max_age=sort_max_age,
min_hits=sort_min_hits,
iou_threshold=sort_iou_thresh)
track_color_id = 0
# Sorting algorithm
'''Computer Color for every box and track'''
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
def compute_color_for_labels(label):
color = [int(int(p * (label ** 2 - label + 1)) % 255) for p in palette]
return tuple(color)
"""" Calculates the relative bounding box from absolute pixel values. """
def bbox_rel(*xyxy):
bbox_left = min([xyxy[0].item(), xyxy[2].item()])
bbox_top = min([xyxy[1].item(), xyxy[3].item()])
bbox_w = abs(xyxy[0].item() - xyxy[2].item())
bbox_h = abs(xyxy[1].item() - xyxy[3].item())
x_c = (bbox_left + bbox_w / 2)
y_c = (bbox_top + bbox_h / 2)
w = bbox_w
h = bbox_h
return x_c, y_c, w, h
"""Function to Draw Bounding boxes"""
def draw_boxes(img, bbox, identities=None, categories=None,
names=None, color_box=None, offset=(0, 0)):
for i, box in enumerate(bbox):
x1, y1, x2, y2 = [int(i) for i in box]
x1 += offset[0]
x2 += offset[0]
y1 += offset[1]
y2 += offset[1]
cat = int(categories[i]) if categories is not None else 0
id = int(identities[i]) if identities is not None else 0
data = (int((box[0]+box[2])/2), (int((box[1]+box[3])/2)))
label = str(id)
if color_box:
color = compute_color_for_labels(id)
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), (255, 191, 0), -1)
cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
[255, 255, 255], 1)
cv2.circle(img, data, 3, color, -1)
else:
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 191, 0), 2)
# cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), (255, 191, 0), -1)
# cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
# [255, 255, 255], 1)
cv2.rectangle(img, (x1, y1 - 3*h), (x1 + w, y1), (255, 191, 0), -1)
cv2.putText(img, label, (x1, y1 - 2*h), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
[255, 255, 255], 1)
cv2.circle(img, data, 3, (255, 191, 0), -1)
return img
def _display_detected_tracks(dets, img, color_box=None):
# pass an empty array to sort
dets_to_sort = np.empty((0, 6))
# NOTE: We send in detected object class too
for x1, y1, x2, y2, conf, detclass in dets.cpu().detach().numpy():
# for x1, y1, x2, y2, conf, detclass in dets.cpu().detach().numpy():
dets_to_sort = np.vstack((dets_to_sort,
np.array([x1, y1, x2, y2,
conf, detclass])))
# Run SORT
tracked_dets = sort_tracker.update(dets_to_sort)
# tracks = sort_tracker.getTrackers()
# loop over tracks
# for track in tracks:
# if color_box:
# color = compute_color_for_labels(track_color_id)
# [cv2.line(img, (int(track.centroidarr[i][0]), int(track.centroidarr[i][1])),
# (int(track.centroidarr[i+1][0]),
# int(track.centroidarr[i+1][1])),
# color, thickness=3) for i, _ in enumerate(track.centroidarr)
# if i < len(track.centroidarr)-1]
# track_color_id = track_color_id+1
# else:
# [cv2.line(img, (int(track.centroidarr[i][0]), int(track.centroidarr[i][1])),
# (int(track.centroidarr[i+1][0]),
# int(track.centroidarr[i+1][1])),
# (124, 252, 0), thickness=3) for i, _ in enumerate(track.centroidarr)
# if i < len(track.centroidarr)-1]
# draw boxes for visualization
if len(tracked_dets) > 0:
bbox_xyxy = tracked_dets[:, :4]
identities = tracked_dets[:, 8]
categories = tracked_dets[:, 4]
names = None
draw_boxes(img, bbox_xyxy, identities,
categories, names, color_box)
|