user3410517
user3410517

Reputation: 295

Print score value in bounding box

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

Answers (1)

vachmara
vachmara

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

Related Questions