tensorturtle / classy-sort-yolov5

Ready-to-use realtime multi-object tracker that works for any object category. YOLOv5 + SORT implementation.
GNU General Public License v3.0
110 stars 29 forks source link

To get confidence score parameter in the sort output #11

Open hlmhlr opened 2 years ago

hlmhlr commented 2 years ago

Hi, Thanks for your repo with multi-class sort. Currently, sort outputs the following parameters: "frame_num, bbox_x1, bbox_y1, bbox_x2, bbox_y2, category, u_overdot, v_overdot, s_overdot, and identity".

How can I add confidence score value for the tracked objects? The output required is: "frame_num, bbox_x1, bbox_y1, bbox_x2, bbox_y2, category, u_overdot, v_overdot, s_overdot, identity, and confidence_score".

Your help would be highly appreciated. Thanks,

tensorturtle commented 2 years ago

Hi, do you still need assistance with this issue?

RUIXUEZHAOFENGNIAN commented 2 years ago

I have the same needs

tensorturtle commented 2 years ago

Hi @RUIXUEZHAOFENGNIAN , I won't be making that modification because it is a breaking change, but I can tell you how to do it yourself:

  1. In sort.py's KalmanBoxTracker class, find __init__() and update(). Put self.conf = bbox[4] in each.
  2. In the same class, under get_state(), add the confidence to the returned tensor with: arr_conf = np.expand_dims(np.array([self.conf]), 0). Put arr_conf in the final np.concatenate function, wherever you'd like.
  3. Make corresponding changes the code that follows.