vukasin-stanojevic / BoostTrack

MIT License
65 stars 5 forks source link

Use BoostTrack+ in YOLOv8-OBB #7

Closed yuanzd123 closed 1 month ago

yuanzd123 commented 1 month ago

Hello Vukasin,

Thank you for your excellent paper, and this open source project. I really want to try it on my YOLOv8-OBB. However, with built-in tracker of the YOLO model, there are only two choice BOTSORT and BYTETRACK. I cannot find any instruction to integrate a custom tracker into YOLOv8. May you kindly provide some steps of using BoostTrack+ in YOLOv8-OBB? Thank you!

currently I am using botsort to track:

model = YOLO("best.pt").to('cuda:0')
results = model.track(frame, persist=True, tracker="botsort.yaml", conf=0.691, iou=0.25)
if results[0].obb is not None:
            boxes = results[0].obb.xyxyxyxy.cpu().numpy()
            clss = results[0].obb.cls.cpu().tolist()
            track_ids = results[0].obb.id.int().cpu().tolist() if results[0].obb.id is not None else [None] * len(clss)
            confs = results[0].obb.conf.cpu().tolist()

            detections = []
            for box, cls, track_id, conf in zip(boxes, clss, track_ids, confs):
                # Process box data
                box = np.array(box).reshape(-1, 2)
                center_x = box[:, 0].mean()
                center_y = box[:, 1].mean()
                width = box[:, 0].max() - box[:, 0].min()
                height = box[:, 1].max() - box[:, 1].min()
                detections.append((int(center_x), int(center_y), int(width), int(height), int(cls), track_id, conf, box))

            frame = process_frame(frame, detections, track_ids)
        cv2.imshow("YOLOv10 OBB Detection", frame)
yuanzd123 commented 1 month ago

Hello @vukasin-stanojevic , can you provide some guidelines please? I would really appreciate

vukasin-stanojevic commented 1 month ago

Hello, If I understood your question correctly, you should only make the line 165 of main.py work:

targets = tracker.update(pred, img, np_img[0].numpy(), tag)

The arguments of update method are: dets - a Tensor outputted from the detector (its shape is number_of_detections $\cdot$ 5). It is assumed that dets is on the GPU (in line 184 if boost_track.py it is converted to numpy ndarray: dets = dets.cpu().detach().numpy()). Important: detected bounding boxes should be in the format $(x_1, y_1, x_2, y_2, conf)$ where $(x_1, y_1)$ are the top left coordinate and $(x_2, y_2)$ bottom right coordinate of the detected bounding box.

img_tensor - a image Tensor passed through the detector to generate detections (img_tensor is possibly scaled down version of the original image). The shape of img_tensor is 1 x 3 x h_s x w_s. img_numpy - numpy version of the original image (original scale, which means img_numpy can be larger than img_tensor). The shape of img_numpy is h x w x 3. tag - a string of the form sequence_name:frame_number (e.g. MOT20-01:1)

So, in essence, BoostTrack only needs the detected bounding boxes and the image to work - it does not matter if the detections are generated by YOLOX or any other model. @yuanzd123

yuanzd123 commented 1 month ago

Thanks for replying to me. The problem is that the YOLO-OBB bounding box format is [class_index, x1, y1, x2, y2, x3, y3, x4, y4], and there are 4 points instead of 2 points (that from normal YOLO model). So I assume that the tracker does not work on an OBB bounding box format?

vukasin-stanojevic commented 1 month ago

And what are these x and y values? I didn't work with YOLOv8-OBB, but it seems like there is noting special about it. For example, the page gives the following easy solution to getting the bounding box info: image

I think I provided you with enough info. The rest is yours research. Good luck!
@yuanzd123

yuanzd123 commented 1 month ago

Thank you very much!