From 3645cf82b64faa3b673d2f05f48e9612386cd37c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:50:27 +0200 Subject: [PATCH 1/9] asso func created initalized on first frame --- boxmot/trackers/basetracker.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/boxmot/trackers/basetracker.py b/boxmot/trackers/basetracker.py index 5a73ae5e75..4987ae7ba6 100644 --- a/boxmot/trackers/basetracker.py +++ b/boxmot/trackers/basetracker.py @@ -4,6 +4,7 @@ import colorsys from abc import ABC, abstractmethod from boxmot.utils import logger as LOGGER +from boxmot.utils.iou import AssociationFunction class BaseTracker(ABC): @@ -15,7 +16,8 @@ def __init__( iou_threshold: float = 0.3, max_obs: int = 50, nr_classes: int = 80, - per_class: bool = False + per_class: bool = False, + asso_func: str = 'iou' ): """ Initialize the BaseTracker object with detection threshold, maximum age, minimum hits, @@ -39,10 +41,12 @@ def __init__( self.nr_classes = nr_classes self.iou_threshold = iou_threshold self.last_emb_size = None + self.asso_func_name = asso_func self.frame_count = 0 self.active_tracks = [] # This might be handled differently in derived classes self.per_class_active_tracks = None + self._first_frame_processed = False # Flag to track if the first frame has been processed # Initialize per-class active tracks if self.per_class: @@ -92,6 +96,29 @@ def get_class_dets_n_embs(self, dets, embs, cls_id): class_embs = None return class_dets, class_embs + @staticmethod + def on_first_frame_setup(method): + """ + Decorator to perform setup on the first frame only. + This ensures that initialization tasks (like setting the association function) only + happen once, on the first frame, and are skipped on subsequent frames. + """ + def wrapper(self, *args, **kwargs): + # If setup hasn't been done yet, perform it + if not self._first_frame_processed: + img = args[1] + self.h, self.w = img.shape[0:2] + self.asso_func = AssociationFunction(w=self.w, h=self.h, asso_mode=self.asso_func_name).asso_func + + # Mark that the first frame setup has been done + self._first_frame_processed = True + + # Call the original method (e.g., update) + return method(self, *args, **kwargs) + + return wrapper + + @staticmethod def per_class_decorator(update_method): """ From ec472f564ebc1db65cd34d6988437f03d17bac23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:55:34 +0200 Subject: [PATCH 2/9] asso func created initalized on first frame --- boxmot/utils/iou.py | 535 +++++++++++++++++++++++--------------------- 1 file changed, 279 insertions(+), 256 deletions(-) diff --git a/boxmot/utils/iou.py b/boxmot/utils/iou.py index 9affd81cc9..185d0672ee 100644 --- a/boxmot/utils/iou.py +++ b/boxmot/utils/iou.py @@ -1,261 +1,284 @@ -# Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license - import numpy as np - -def iou_batch(bboxes1, bboxes2) -> np.ndarray: - """ - From SORT: Computes IOU between two bboxes in the form [x1,y1,x2,y2] - """ - bboxes2 = np.expand_dims(bboxes2, 0) - bboxes1 = np.expand_dims(bboxes1, 1) - - xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) - yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) - yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - w = np.maximum(0.0, xx2 - xx1) - h = np.maximum(0.0, yy2 - yy1) - wh = w * h - o = wh / ( - (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + - (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - - wh - ) - return o - - -def hmiou_batch(bboxes1, bboxes2): - """ - Computes Hybrid Mass Intersection over Union (HM IoU) between two sets of bounding boxes. - :param bboxes1: Predicted bounding boxes (N, 4) in the format [x1, y1, x2, y2]. - :param bboxes2: Ground truth bounding boxes (M, 4) in the format [x1, y1, x2, y2]. - :return: HM IoU matrix of shape (N, M). - """ - # Ensure bounding boxes are in correct format - bboxes1 = np.expand_dims(bboxes1, 1) - bboxes2 = np.expand_dims(bboxes2, 0) +class AssociationFunction: + def __init__(self, w, h, asso_mode="iou"): + """ + Initializes the AssociationFunction class with the necessary parameters for bounding box operations. + The association function is selected based on the `asso_mode` string provided during class creation. + + Parameters: + w (int): The width of the frame, used for normalizing centroid distance. + h (int): The height of the frame, used for normalizing centroid distance. + asso_mode (str): The association function to use (e.g., "iou", "giou", "centroid", etc.). + """ + self.w = w + self.h = h + self.asso_mode = asso_mode + self.asso_func = self._get_asso_func(asso_mode) + + @staticmethod + def iou_batch(bboxes1, bboxes2) -> np.ndarray: + bboxes2 = np.expand_dims(bboxes2, 0) + bboxes1 = np.expand_dims(bboxes1, 1) + + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + wh = w * h + o = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - + wh + ) + return o + + @staticmethod + def hmiou_batch(bboxes1, bboxes2): + bboxes1 = np.expand_dims(bboxes1, 1) + bboxes2 = np.expand_dims(bboxes2, 0) + yy11 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + yy12 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + yy21 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) + yy22 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + o = np.maximum(0, yy12 - yy11) / np.maximum(1e-10, yy22 - yy21) + + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0., xx2 - xx1) + h = np.maximum(0., yy2 - yy1) + wh = w * h + iou = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh + ) + iou *= o + return iou + + @staticmethod + def giou_batch(bboxes1, bboxes2) -> np.ndarray: + """ + :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) + :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) + :return: + """ + # for details should go to https://arxiv.org/pdf/1902.09630.pdf + # ensure predict's bbox form + bboxes2 = np.expand_dims(bboxes2, 0) + bboxes1 = np.expand_dims(bboxes1, 1) + + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + wh = w * h + iou = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - + wh + ) + + xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) + yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) + xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) + yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + wc = xxc2 - xxc1 + hc = yyc2 - yyc1 + assert (wc > 0).all() and (hc > 0).all() + area_enclose = wc * hc + giou = iou - (area_enclose - wh) / area_enclose + giou = (giou + 1.0) / 2.0 # resize from (-1,1) to (0,1) + return giou + + + def centroid_batch(self, bboxes1, bboxes2) -> np.ndarray: + centroids1 = np.stack(((bboxes1[..., 0] + bboxes1[..., 2]) / 2, + (bboxes1[..., 1] + bboxes1[..., 3]) / 2), axis=-1) + centroids2 = np.stack(((bboxes2[..., 0] + bboxes2[..., 2]) / 2, + (bboxes2[..., 1] + bboxes2[..., 3]) / 2), axis=-1) + + centroids1 = np.expand_dims(centroids1, 1) + centroids2 = np.expand_dims(centroids2, 0) + + distances = np.sqrt(np.sum((centroids1 - centroids2) ** 2, axis=-1)) + norm_factor = np.sqrt(self.w ** 2 + self.h ** 2) + normalized_distances = distances / norm_factor + + return 1 - normalized_distances + - # Calculate vertical overlap term `o` - yy11 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - yy12 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - yy21 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) - yy22 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + def ciou_batch(bboxes1, bboxes2) -> np.ndarray: + """ + :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) + :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) + :return: + """ + # for details should go to https://arxiv.org/pdf/1902.09630.pdf + # ensure predict's bbox form + bboxes2 = np.expand_dims(bboxes2, 0) + bboxes1 = np.expand_dims(bboxes1, 1) + + # calculate the intersection box + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + wh = w * h + iou = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - + wh + ) + + centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0 + centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0 + centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0 + centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0 + + inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2 + + xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) + yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) + xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) + yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + + outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2 + + w1 = bboxes1[..., 2] - bboxes1[..., 0] + h1 = bboxes1[..., 3] - bboxes1[..., 1] + w2 = bboxes2[..., 2] - bboxes2[..., 0] + h2 = bboxes2[..., 3] - bboxes2[..., 1] + + # prevent dividing over zero. add one pixel shift + h2 = h2 + 1.0 + h1 = h1 + 1.0 + arctan = np.arctan(w2 / h2) - np.arctan(w1 / h1) + v = (4 / (np.pi**2)) * (arctan**2) + S = 1 - iou + alpha = v / (S + v) + ciou = iou - inner_diag / outer_diag - alpha * v + + return (ciou + 1) / 2.0 - o = np.maximum(0, yy12 - yy11) / np.maximum(1e-10, yy22 - yy21) # Avoid division by zero - - # Calculate standard IoU - xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) - yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) - yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - w = np.maximum(0., xx2 - xx1) - h = np.maximum(0., yy2 - yy1) - wh = w * h - iou = wh / ( - (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + - (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh - ) + def diou_batch(bboxes1, bboxes2) -> np.ndarray: + """ + :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) + :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) + :return: + """ + # for details should go to https://arxiv.org/pdf/1902.09630.pdf + # ensure predict's bbox form + bboxes2 = np.expand_dims(bboxes2, 0) + bboxes1 = np.expand_dims(bboxes1, 1) + + # calculate the intersection box + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + wh = w * h + iou = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - + wh + ) + + centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0 + centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0 + centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0 + centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0 + + inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2 + + xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) + yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) + xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) + yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + + outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2 + diou = iou - inner_diag / outer_diag + + return (diou + 1) / 2.0 - iou *= o # Multiply IoU by vertical overlap term - return iou - - -def giou_batch(bboxes1, bboxes2) -> np.ndarray: - """ - :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) - :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) - :return: - """ - # for details should go to https://arxiv.org/pdf/1902.09630.pdf - # ensure predict's bbox form - bboxes2 = np.expand_dims(bboxes2, 0) - bboxes1 = np.expand_dims(bboxes1, 1) - - xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) - yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) - yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - w = np.maximum(0.0, xx2 - xx1) - h = np.maximum(0.0, yy2 - yy1) - wh = w * h - iou = wh / ( - (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + - (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - - wh - ) - - xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) - yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) - xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) - yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) - wc = xxc2 - xxc1 - hc = yyc2 - yyc1 - assert (wc > 0).all() and (hc > 0).all() - area_enclose = wc * hc - giou = iou - (area_enclose - wh) / area_enclose - giou = (giou + 1.0) / 2.0 # resize from (-1,1) to (0,1) - return giou - - -def diou_batch(bboxes1, bboxes2) -> np.ndarray: - """ - :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) - :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) - :return: - """ - # for details should go to https://arxiv.org/pdf/1902.09630.pdf - # ensure predict's bbox form - bboxes2 = np.expand_dims(bboxes2, 0) - bboxes1 = np.expand_dims(bboxes1, 1) - - # calculate the intersection box - xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) - yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) - yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - w = np.maximum(0.0, xx2 - xx1) - h = np.maximum(0.0, yy2 - yy1) - wh = w * h - iou = wh / ( - (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + - (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - - wh - ) - - centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0 - centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0 - centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0 - centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0 - - inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2 - - xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) - yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) - xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) - yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) - - outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2 - diou = iou - inner_diag / outer_diag - - return (diou + 1) / 2.0 # resize from (-1,1) to (0,1) - - -def ciou_batch(bboxes1, bboxes2) -> np.ndarray: - """ - :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) - :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) - :return: - """ - # for details should go to https://arxiv.org/pdf/1902.09630.pdf - # ensure predict's bbox form - bboxes2 = np.expand_dims(bboxes2, 0) - bboxes1 = np.expand_dims(bboxes1, 1) - - # calculate the intersection box - xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) - yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) - xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) - yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) - w = np.maximum(0.0, xx2 - xx1) - h = np.maximum(0.0, yy2 - yy1) - wh = w * h - iou = wh / ( - (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + - (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - - wh - ) - - centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0 - centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0 - centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0 - centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0 - - inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2 - - xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) - yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) - xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) - yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) - - outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2 - - w1 = bboxes1[..., 2] - bboxes1[..., 0] - h1 = bboxes1[..., 3] - bboxes1[..., 1] - w2 = bboxes2[..., 2] - bboxes2[..., 0] - h2 = bboxes2[..., 3] - bboxes2[..., 1] - - # prevent dividing over zero. add one pixel shift - h2 = h2 + 1.0 - h1 = h1 + 1.0 - arctan = np.arctan(w2 / h2) - np.arctan(w1 / h1) - v = (4 / (np.pi**2)) * (arctan**2) - S = 1 - iou - alpha = v / (S + v) - ciou = iou - inner_diag / outer_diag - alpha * v - - return (ciou + 1) / 2.0 # resize from (-1,1) to (0,1) - - -def centroid_batch(bboxes1, bboxes2, w, h) -> np.ndarray: - """ - Computes the normalized centroid distance between two sets of bounding boxes. - Bounding boxes are in the format [x1, y1, x2, y2]. - `normalize_scale` is a tuple (width, height) to normalize the distance. - """ - - # Calculate centroids - centroids1 = np.stack(((bboxes1[..., 0] + bboxes1[..., 2]) / 2, - (bboxes1[..., 1] + bboxes1[..., 3]) / 2), axis=-1) - centroids2 = np.stack(((bboxes2[..., 0] + bboxes2[..., 2]) / 2, - (bboxes2[..., 1] + bboxes2[..., 3]) / 2), axis=-1) - - # Expand dimensions for broadcasting - centroids1 = np.expand_dims(centroids1, 1) - centroids2 = np.expand_dims(centroids2, 0) - - # Calculate Euclidean distances - distances = np.sqrt(np.sum((centroids1 - centroids2) ** 2, axis=-1)) - - # Normalize distances - norm_factor = np.sqrt(w**2 + h**2) - normalized_distances = distances / norm_factor - - return 1 - normalized_distances - - -def run_asso_func(func, *args): - """ - Wrapper function that checks the inputs to the association functions - and then call either one of the iou association functions or centroid. - - Parameters: - func: The batch function to call (either *iou*_batch or centroid_batch). - *args: Variable length argument list, containing either bounding boxes and optionally size parameters. - """ - if func not in [iou_batch, giou_batch, diou_batch, ciou_batch, centroid_batch, hmiou_batch]: - raise ValueError("Invalid function specified. Must be either '(g,d,c, )iou_batch' or 'centroid_batch'.") - - if func in (iou_batch, giou_batch, diou_batch, ciou_batch, centroid_batch, hmiou_batch): - if len(args) != 4 or not all(isinstance(arg, (list, np.ndarray)) for arg in args[0:2]): - raise ValueError("Invalid arguments for iou_batch. Expected two bounding boxes.") - return func(*args[0:2]) - elif func is centroid_batch: - if len(args) != 4 or not all(isinstance(arg, (list, np.ndarray)) for arg in args[:2]) or not all(isinstance(arg, (int)) for arg in args[2:]): - raise ValueError("Invalid arguments for centroid_batch. Expected two bounding boxes and two size parameters.") - return func(*args) - else: - raise ValueError("No such association method") - - -def get_asso_func(asso_mode): - ASSO_FUNCS = { - "iou": iou_batch, - "hmiou": hmiou_batch, - "giou": giou_batch, - "ciou": ciou_batch, - "diou": diou_batch, - "centroid": centroid_batch - } - - return ASSO_FUNCS[asso_mode] + + def giou_batch(self, bboxes1, bboxes2) -> np.ndarray: + """ + :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2) + :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2) + :return: + """ + # for details should go to https://arxiv.org/pdf/1902.09630.pdf + # ensure predict's bbox form + bboxes2 = np.expand_dims(bboxes2, 0) + bboxes1 = np.expand_dims(bboxes1, 1) + + xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0]) + yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1]) + xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2]) + yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3]) + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + wh = w * h + iou = wh / ( + (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) + + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - + wh + ) + + xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0]) + yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1]) + xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2]) + yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3]) + wc = xxc2 - xxc1 + hc = yyc2 - yyc1 + assert (wc > 0).all() and (hc > 0).all() + area_enclose = wc * hc + giou = iou - (area_enclose - wh) / area_enclose + giou = (giou + 1.0) / 2.0 # resize from (-1,1) to (0,1) + return giou + + + @staticmethod + def run_asso_func(self, bboxes1, bboxes2): + """ + Runs the selected association function (based on the initialization string) on the input bounding boxes. + + Parameters: + bboxes1: First set of bounding boxes. + bboxes2: Second set of bounding boxes. + """ + return self.asso_func(bboxes1, bboxes2) + + def _get_asso_func(self, asso_mode): + """ + Returns the corresponding association function based on the provided mode string. + + Parameters: + asso_mode (str): The association function to use (e.g., "iou", "giou", "centroid", etc.). + + Returns: + function: The appropriate function for the association calculation. + """ + ASSO_FUNCS = { + "iou": AssociationFunction.iou_batch, + "hmiou": AssociationFunction.hmiou_batch, + "giou": AssociationFunction.giou_batch, + "ciou": AssociationFunction.ciou_batch, + "diou": AssociationFunction.diou_batch, + "centroid": self.centroid_batch # only not being staticmethod + } + + if self.asso_mode not in ASSO_FUNCS: + raise ValueError(f"Invalid association mode: {self.asso_mode}. Choose from {list(ASSO_FUNCS.keys())}") + + return ASSO_FUNCS[self.asso_mode] From a4f12ab8c48ba4baf2b0aac246a3878bed76262f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:55:55 +0200 Subject: [PATCH 3/9] asso func created initalized on first frame --- boxmot/configs/ocsort.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/boxmot/configs/ocsort.yaml b/boxmot/configs/ocsort.yaml index 93e79e989b..b80f888feb 100644 --- a/boxmot/configs/ocsort.yaml +++ b/boxmot/configs/ocsort.yaml @@ -21,7 +21,7 @@ delta_t: asso_func: type: choice default: iou # from the default parameters - options: ['iou', 'giou', 'centroid'] + options: ['iou', 'giou'] use_byte: type: choice From 817f0970e13f0413c1671b10259b5b749cfe7164 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:58:46 +0200 Subject: [PATCH 4/9] asso func created initalized on first frame --- boxmot/utils/matching.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/boxmot/utils/matching.py b/boxmot/utils/matching.py index 830ef18621..81e3a11bf1 100644 --- a/boxmot/utils/matching.py +++ b/boxmot/utils/matching.py @@ -5,7 +5,8 @@ import scipy import torch from scipy.spatial.distance import cdist -from boxmot.utils.iou import iou_batch, diou_batch +from boxmot.utils.iou import AssociationFunction + """ Table for the 0.95 quantile of the chi-square distribution with N degrees of @@ -111,7 +112,7 @@ def d_iou_distance(atracks, btracks): ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32) if ious.size == 0: return ious - _ious = diou_batch(atlbrs, btlbrs) + _ious = AssociationFunction.diou_batch(atlbrs, btlbrs) cost_matrix = 1 - _ious @@ -138,7 +139,7 @@ def iou_distance(atracks, btracks): ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32) if ious.size == 0: return ious - _ious = iou_batch(atlbrs, btlbrs) + _ious = AssociationFunction.iou_batch(atlbrs, btlbrs) cost_matrix = 1 - _ious From 175ab98666db79962d452c50069d799d103b62a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:59:06 +0200 Subject: [PATCH 5/9] asso func created initalized on first frame --- boxmot/utils/association.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/boxmot/utils/association.py b/boxmot/utils/association.py index 4e1a5a965f..729104f0d0 100644 --- a/boxmot/utils/association.py +++ b/boxmot/utils/association.py @@ -2,7 +2,7 @@ import numpy as np -from boxmot.utils.iou import iou_batch, centroid_batch, run_asso_func +from boxmot.utils.iou import AssociationFunction def speed_direction_batch(dets, tracks): @@ -40,7 +40,7 @@ def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3): np.empty((0, 5), dtype=int), ) - iou_matrix = iou_batch(detections, trackers) + iou_matrix = AssociationFunction.iou_batch(detections, trackers) if min(iou_matrix.shape) > 0: a = (iou_matrix > iou_threshold).astype(np.int32) @@ -143,7 +143,7 @@ def associate( valid_mask = np.ones(previous_obs.shape[0]) valid_mask[np.where(previous_obs[:, 4] < 0)] = 0 - iou_matrix = run_asso_func(asso_func, detections, trackers, w, h) + iou_matrix = asso_func(detections, trackers) #iou_matrix = iou_batch(detections, trackers) scores = np.repeat(detections[:, -1][:, np.newaxis], trackers.shape[0], axis=1) # iou_matrix = iou_matrix * scores # a trick sometiems works, we don't encourage this @@ -235,7 +235,7 @@ def associate_kitti( """ Cost from IoU """ - iou_matrix = iou_batch(detections, trackers) + iou_matrix = AssociationFunction.iou_batch(detections, trackers) """ With multiple categories, generate the cost for catgory mismatch From 2df2bf6407e66bd6d780be996eaad1086d888a91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 16:59:40 +0200 Subject: [PATCH 6/9] asso func created initalized on first frame --- boxmot/trackers/ocsort/ocsort.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/boxmot/trackers/ocsort/ocsort.py b/boxmot/trackers/ocsort/ocsort.py index feb05cb906..f2151a5442 100644 --- a/boxmot/trackers/ocsort/ocsort.py +++ b/boxmot/trackers/ocsort/ocsort.py @@ -9,8 +9,6 @@ from boxmot.motion.kalman_filters.xysr_kf import KalmanFilterXYSR from boxmot.utils.association import associate, linear_assignment -from boxmot.utils.iou import get_asso_func -from boxmot.utils.iou import run_asso_func from boxmot.trackers.basetracker import BaseTracker from boxmot.utils.ops import xyxy2xysr @@ -212,7 +210,7 @@ def __init__( Q_xy_scaling: float = 0.01, Q_s_scaling: float = 0.0001 ): - super().__init__(max_age=max_age, per_class=per_class) + super().__init__(max_age=max_age, per_class=per_class, asso_func=asso_func) """ Sets key parameters for SORT """ @@ -223,13 +221,13 @@ def __init__( self.frame_count = 0 self.det_thresh = det_thresh self.delta_t = delta_t - self.asso_func = get_asso_func(asso_func) self.inertia = inertia self.use_byte = use_byte self.Q_xy_scaling = Q_xy_scaling self.Q_s_scaling = Q_s_scaling KalmanBoxTracker.count = 0 + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> np.ndarray: """ @@ -327,7 +325,7 @@ def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> if unmatched_dets.shape[0] > 0 and unmatched_trks.shape[0] > 0: left_dets = dets[unmatched_dets] left_trks = last_boxes[unmatched_trks] - iou_left = run_asso_func(self.asso_func, left_dets, left_trks, w, h) + iou_left = self.asso_func(left_dets, left_trks) iou_left = np.array(iou_left) if iou_left.max() > self.asso_threshold: """ From 9ca78cbdca2e581805248d8b682117749391bab0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 17:09:10 +0200 Subject: [PATCH 7/9] asso func created initalized on first frame --- boxmot/trackers/botsort/botsort.py | 1 + boxmot/trackers/bytetrack/byte_tracker.py | 1 + boxmot/trackers/deepocsort/deep_ocsort.py | 6 +++--- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/boxmot/trackers/botsort/botsort.py b/boxmot/trackers/botsort/botsort.py index f1ad58d6b8..9687f3c2f6 100644 --- a/boxmot/trackers/botsort/botsort.py +++ b/boxmot/trackers/botsort/botsort.py @@ -84,6 +84,7 @@ def __init__( self.cmc = get_cmc_method('ecc')() self.fuse_first_associate = fuse_first_associate + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> np.ndarray: self.check_inputs(dets, img) diff --git a/boxmot/trackers/bytetrack/byte_tracker.py b/boxmot/trackers/bytetrack/byte_tracker.py index f48ade4a59..32f2af99b3 100644 --- a/boxmot/trackers/bytetrack/byte_tracker.py +++ b/boxmot/trackers/bytetrack/byte_tracker.py @@ -150,6 +150,7 @@ def __init__( self.max_time_lost = self.buffer_size self.kalman_filter = KalmanFilterXYAH() + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray = None, embs: np.ndarray = None) -> np.ndarray: diff --git a/boxmot/trackers/deepocsort/deep_ocsort.py b/boxmot/trackers/deepocsort/deep_ocsort.py index 0efc16929a..aa8237f68e 100644 --- a/boxmot/trackers/deepocsort/deep_ocsort.py +++ b/boxmot/trackers/deepocsort/deep_ocsort.py @@ -10,7 +10,6 @@ from boxmot.motion.kalman_filters.xysr_kf import KalmanFilterXYSR from boxmot.motion.kalman_filters.xywh_kf import KalmanFilterXYWH from boxmot.utils.association import associate, linear_assignment -from boxmot.utils.iou import get_asso_func from boxmot.trackers.basetracker import BaseTracker from boxmot.utils.ops import xyxy2xysr @@ -272,7 +271,7 @@ def __init__( Q_s_scaling: float = 0.0001, **kwargs: dict ): - super().__init__(max_age=max_age, per_class=per_class) + super().__init__(max_age=max_age, per_class=per_class, asso_func=asso_func) """ Sets key parameters for SORT """ @@ -281,7 +280,7 @@ def __init__( self.iou_threshold = iou_threshold self.det_thresh = det_thresh self.delta_t = delta_t - self.asso_func = get_asso_func(asso_func) + self.asso_func = asso_func self.inertia = inertia self.w_association_emb = w_association_emb self.alpha_fixed_emb = alpha_fixed_emb @@ -300,6 +299,7 @@ def __init__( self.cmc_off = cmc_off self.aw_off = aw_off + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> np.ndarray: """ From 4873d4b310116532a8370d60aacb1796cc23c83d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 17:10:17 +0200 Subject: [PATCH 8/9] asso func created initalized on first frame --- boxmot/trackers/hybridsort/hybridsort.py | 5 ++--- boxmot/trackers/imprassoc/impr_assoc_tracker.py | 1 + 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/boxmot/trackers/hybridsort/hybridsort.py b/boxmot/trackers/hybridsort/hybridsort.py index d6440325ec..257fadf54c 100644 --- a/boxmot/trackers/hybridsort/hybridsort.py +++ b/boxmot/trackers/hybridsort/hybridsort.py @@ -13,7 +13,6 @@ from boxmot.trackers.hybridsort.association import ( associate_4_points_with_score, associate_4_points_with_score_with_reid, cal_score_dif_batch_two_score, embedding_distance, linear_assignment) -from boxmot.utils.iou import get_asso_func from boxmot.trackers.basetracker import BaseTracker @@ -353,7 +352,7 @@ class HybridSORT(BaseTracker): """ def __init__(self, reid_weights, device, half, det_thresh, per_class=False, max_age=30, min_hits=3, iou_threshold=0.3, delta_t=3, asso_func="iou", inertia=0.2, longterm_reid_weight=0, TCM_first_step_weight=0, use_byte=False): - super().__init__(max_age=max_age, per_class=per_class) + super().__init__(max_age=max_age, per_class=per_class, asso_func=asso_func) """ Sets key parameters for SORT @@ -365,7 +364,6 @@ def __init__(self, reid_weights, device, half, det_thresh, per_class=False, max_ self.frame_count: int = 0 self.det_thresh: float = det_thresh self.delta_t: int = delta_t - self.asso_func: str = get_asso_func(asso_func) # assuming get_asso_func returns a callable function self.inertia: float = inertia self.use_byte: bool = use_byte self.low_thresh: float = 0.1 @@ -394,6 +392,7 @@ def camera_update(self, trackers, warp_matrix): for tracker in trackers: tracker.camera_update(warp_matrix) + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> np.ndarray: """ diff --git a/boxmot/trackers/imprassoc/impr_assoc_tracker.py b/boxmot/trackers/imprassoc/impr_assoc_tracker.py index 022044fd80..3bff68fda0 100644 --- a/boxmot/trackers/imprassoc/impr_assoc_tracker.py +++ b/boxmot/trackers/imprassoc/impr_assoc_tracker.py @@ -264,6 +264,7 @@ def __init__( self.cmc = SOF() + @BaseTracker.on_first_frame_setup @BaseTracker.per_class_decorator def update(self, dets: np.ndarray, img: np.ndarray, embs: np.ndarray = None) -> np.ndarray: self.check_inputs(dets, img) From 38da1b69988cca5250c8509aa16885585fc76e08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikel=20Brostr=C3=B6m?= Date: Thu, 26 Sep 2024 17:43:30 +0200 Subject: [PATCH 9/9] add cmc method to config --- boxmot/configs/botsort.yaml | 4 ++-- boxmot/trackers/botsort/botsort.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/boxmot/configs/botsort.yaml b/boxmot/configs/botsort.yaml index a65187dbf8..4ff1115e49 100644 --- a/boxmot/configs/botsort.yaml +++ b/boxmot/configs/botsort.yaml @@ -35,5 +35,5 @@ appearance_thresh: cmc_method: type: choice - default: sof # from the default parameters - options: [sof, cmc] \ No newline at end of file + default: ecc # from the default parameters + options: [sof, ecc] \ No newline at end of file diff --git a/boxmot/trackers/botsort/botsort.py b/boxmot/trackers/botsort/botsort.py index 9687f3c2f6..40d878189a 100644 --- a/boxmot/trackers/botsort/botsort.py +++ b/boxmot/trackers/botsort/botsort.py @@ -52,7 +52,7 @@ def __init__( match_thresh: float = 0.8, proximity_thresh: float = 0.5, appearance_thresh: float = 0.25, - cmc_method: str = "sof", + cmc_method: str = "ecc", frame_rate=30, fuse_first_associate: bool = False, with_reid: bool = True, @@ -81,7 +81,7 @@ def __init__( weights=reid_weights, device=device, half=half ).model - self.cmc = get_cmc_method('ecc')() + self.cmc = get_cmc_method(cmc_method)() self.fuse_first_associate = fuse_first_associate @BaseTracker.on_first_frame_setup