Reputation: 295
I try add on the bounding boxes the score for the Object tracker on the link: https://github.com/pinczakko/yolov4-deepsort
How to get the score from the track object? Here the method to process the detections:
def process_detections(tracker, detections, nms_max_overlap, frame):
"""
Process objects detected by YOLO with deepsort
Returns processed frame
"""
#initialize color map
cmap = plt.get_cmap('tab20b')
colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]
# run non-maxima supression
boxs = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
classes = np.array([d.class_name for d in detections])
indices = preprocessing.non_max_suppression(boxs, classes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
# Call the tracker
tracker.predict()
tracker.update(detections)
# update tracks
for track in tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
class_name = track.get_class()
# draw bbox on screen
color = colors[int(track.track_id) % len(colors)]
color = [i * 255 for i in color]
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 1)
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1]-30)),
(int(bbox[0])+(len(class_name)+len(str(track.track_id)))*17, int(bbox[1])), color, -1)
cv2.putText(frame, class_name + "-" + str(track.track_id),(int(bbox[0]),
int(bbox[1]-10)),0, 0.5, (255,255,255), 1)
# if enable info flag then print details about each track
if FLAGS.info:
print("Tracker ID: {}, Class: {}, BBox Coords (xmin, ymin, xmax, ymax): {}".format(str(track.track_id),
class_name, (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))))
return frame
Upvotes: 2
Views: 1057
Reputation: 126
If you want to print the confidence score, change this line :
cv2.putText(frame, class_name + "-" + str(track.track_id),(int(bbox[0]), int(bbox[1]-10)),0, 0.5, (255,255,255), 1)
Into :
cv2.putText(frame, class_name + "-" + str(track.track_id) + str(scores[track.track_id]),(int(bbox[0]), int(bbox[1]-10)),0, 0.5, (255,255,255), 1)
=========
EDIT :
After some local testing, this function is working :
def process_detections(tracker, detections, nms_max_overlap, frame):
"""
Process objects detected by YOLO with deepsort
Returns processed frame
"""
# initialize color map
cmap = plt.get_cmap('tab20b')
colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]
# run non-maxima supression
boxs = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
classes = np.array([d.class_name for d in detections])
indices = preprocessing.non_max_suppression(boxs, classes, nms_max_overlap, scores)
detections = [detections[i] for i in indices]
# Call the tracker
tracker.predict()
tracker.update(detections)
# DEBUG: I do not know why there are tracks than detections objects its why i verify the length if the scores
# print(len(detections), len(scores), len(tracker.tracks))
# update tracks
for i, track in enumerate(tracker.tracks):
if not track.is_confirmed() or track.time_since_update > 1:
continue
bbox = track.to_tlbr()
class_name = track.get_class()
# draw bbox on screen
color = colors[int(track.track_id) % len(colors)]
color = [i * 255 for i in color]
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 1)
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1] - 30)),
(int(bbox[0]) + (len(class_name) + len(str(track.track_id))) * 17, int(bbox[1])), color, -1)
if i < len(scores):
cv2.putText(frame,
class_name + "-" + str(track.track_id) + " - {:.2f}".format(float(scores[i])),
(int(bbox[0]),
int(bbox[1] - 10)), 0, 0.5, (255, 255, 255), 1)
else:
cv2.putText(frame,
class_name + "-" + str(track.track_id),
(int(bbox[0]),
int(bbox[1] - 10)), 0, 0.5, (255, 255, 255), 1)
# if enable info flag then print details about each track
if FLAGS.info:
print("Tracker ID: {}, Class: {}, BBox Coords (xmin, ymin, xmax, ymax): {}".format(str(track.track_id),
class_name, (
int(bbox[0]),
int(bbox[1]),
int(bbox[2]),
int(bbox[3]))))
return frame
Upvotes: 1