VerokeAI/Object_tracking_boxmot
0
1import numpy as np2import cv2 as cv3 4def iou_obb_pair(i, j, bboxes1, bboxes2):5 """6 Compute IoU for the rotated rectangles at index i and j in the batches `bboxes1`, `bboxes2` .7 """8 rect1 = bboxes1[int(i)]9 rect2 = bboxes2[int(j)]10 11 (cx1, cy1, w1, h1, angle1) = rect1[0:5]12 (cx2, cy2, w2, h2, angle2) = rect2[0:5]13 14 15 r1 = ((cx1, cy1), (w1, h1), angle1)16 r2 = ((cx2, cy2), (w2, h2), angle2)17 18 # Compute intersection19 ret, intersect = cv.rotatedRectangleIntersection(r1, r2)20 if ret == 0 or intersect is None:21 return 0.0 # No intersection22 23 # Calculate intersection area24 intersection_area = cv.contourArea(intersect)25 26 # Calculate union area27 area1 = w1 * h128 area2 = w2 * h229 union_area = area1 + area2 - intersection_area30 31 # Compute IoU32 return intersection_area / union_area if union_area > 0 else 0.033 34class AssociationFunction:35 def __init__(self, w, h, asso_mode="iou"):36 """37 Initializes the AssociationFunction class with the necessary parameters for bounding box operations.38 The association function is selected based on the `asso_mode` string provided during class creation.39 40 Parameters:41 w (int): The width of the frame, used for normalizing centroid distance.42 h (int): The height of the frame, used for normalizing centroid distance.43 asso_mode (str): The association function to use (e.g., "iou", "giou", "centroid", etc.).44 """45 self.w = w46 self.h = h47 self.asso_mode = asso_mode48 self.asso_func = self._get_asso_func(asso_mode)49 50 @staticmethod51 def iou_batch(bboxes1, bboxes2) -> np.ndarray:52 bboxes2 = np.expand_dims(bboxes2, 0)53 bboxes1 = np.expand_dims(bboxes1, 1)54 55 xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])56 yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])57 xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])58 yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])59 w = np.maximum(0.0, xx2 - xx1)60 h = np.maximum(0.0, yy2 - yy1)61 wh = w * h62 o = wh / (63 (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) +64 (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) -65 wh66 )67 return o68 69 @staticmethod70 def iou_batch_obb(bboxes1, bboxes2) -> np.ndarray:71 72 N, M = len(bboxes1), len(bboxes2)73 74 def wrapper(i, j):75 return iou_obb_pair(i, j, bboxes1, bboxes2)76 77 iou_matrix = np.fromfunction(np.vectorize(wrapper), shape=(N, M), dtype=int)78 return iou_matrix79 80 @staticmethod81 def hmiou_batch(bboxes1, bboxes2):82 """83 Compute a modified Intersection over Union (hIoU) between two batches of bounding boxes,84 incorporating a vertical overlap ratio.85 86 Parameters:87 - bboxes1: (N, 4) array of bounding boxes [x1, y1, x2, y2]88 - bboxes2: (M, 4) array of bounding boxes [x1, y1, x2, y2]89 90 Returns:91 - hmiou: (N, M) array where hmiou[i, j] is the modified IoU between bboxes1[i] and bboxes2[j]92 """93 # Expand dimensions for broadcasting94 bboxes1 = np.expand_dims(bboxes1, axis=1) # Shape: (N, 1, 4)95 bboxes2 = np.expand_dims(bboxes2, axis=0) # Shape: (1, M, 4)96 97 # Compute vertical overlap ratio 'o'98 intersect_y1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])99 intersect_y2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])100 intersection_height = np.maximum(0.0, intersect_y2 - intersect_y1)101 102 union_y1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])103 union_y2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])104 union_height = np.maximum(1e-10, union_y2 - union_y1)105 106 o = intersection_height / union_height107 108 # Compute standard IoU109 inter_x1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])110 inter_y1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])111 inter_x2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])112 inter_y2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])113 114 inter_w = np.maximum(0.0, inter_x2 - inter_x1)115 inter_h = np.maximum(0.0, inter_y2 - inter_y1)116 inter_area = inter_w * inter_h117 118 area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) # Shape: (N, 1)119 area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) # Shape: (1, M)120 121 union_area = area1 + area2 - inter_area122 123 iou = inter_area / (union_area + 1e-10)124 125 # Modify IoU with vertical overlap ratio126 hmiou = iou * o127 128 return hmiou129 130 @staticmethod131 def giou_batch(bboxes1, bboxes2) -> np.ndarray:132 """133 :param bboxes1: predict of bbox(N,4)(x1,y1,x2,y2)134 :param bboxes2: groundtruth of bbox(N,4)(x1,y1,x2,y2)135 :return:136 """137 # Ensure predict's bbox form138 bboxes2 = np.expand_dims(bboxes2, 0)139 bboxes1 = np.expand_dims(bboxes1, 1)140 141 xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])142 yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])143 xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])144 yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])145 w = np.maximum(0.0, xx2 - xx1)146 h = np.maximum(0.0, yy2 - yy1)147 wh = w * h # Intersection area148 149 # Compute areas of individual boxes150 area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])151 area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1])152 153 # Union area154 union_area = area1 + area2 - wh155 156 iou = wh / union_area157 158 xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])159 yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])160 xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])161 yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])162 wc = xxc2 - xxc1163 hc = yyc2 - yyc1164 assert (wc > 0).all() and (hc > 0).all()165 area_enclose = wc * hc # Area of the smallest enclosing box166 167 # Corrected GIoU computation168 giou = iou - (area_enclose - union_area) / area_enclose169 giou = (giou + 1.0) / 2.0 # Resize from (-1,1) to (0,1)170 return giou171 172 173 def centroid_batch(self, bboxes1, bboxes2) -> np.ndarray:174 centroids1 = np.stack(((bboxes1[..., 0] + bboxes1[..., 2]) / 2,175 (bboxes1[..., 1] + bboxes1[..., 3]) / 2), axis=-1)176 centroids2 = np.stack(((bboxes2[..., 0] + bboxes2[..., 2]) / 2,177 (bboxes2[..., 1] + bboxes2[..., 3]) / 2), axis=-1)178 179 centroids1 = np.expand_dims(centroids1, 1)180 centroids2 = np.expand_dims(centroids2, 0)181 182 distances = np.sqrt(np.sum((centroids1 - centroids2) ** 2, axis=-1))183 norm_factor = np.sqrt(self.w ** 2 + self.h ** 2)184 normalized_distances = distances / norm_factor185 186 return 1 - normalized_distances187 188 def centroid_batch_obb(self, bboxes1, bboxes2) -> np.ndarray:189 centroids1 = np.stack((bboxes1[..., 0], bboxes1[..., 1]),axis=-1)190 centroids2 = np.stack((bboxes2[..., 0], bboxes2[..., 1]),axis=-1)191 192 centroids1 = np.expand_dims(centroids1, 1)193 centroids2 = np.expand_dims(centroids2, 0)194 195 distances = np.sqrt(np.sum((centroids1 - centroids2) ** 2, axis=-1))196 norm_factor = np.sqrt(self.w ** 2 + self.h ** 2)197 normalized_distances = distances / norm_factor198 199 return 1 - normalized_distances200 201 202 @staticmethod203 def ciou_batch(bboxes1, bboxes2) -> np.ndarray:204 """205 Calculate Complete Intersection over Union (CIoU) for batches of bounding boxes.206 207 :param bboxes1: Predicted bounding boxes of shape (N, 4) as (x1, y1, x2, y2)208 :param bboxes2: Ground truth bounding boxes of shape (N, 4) as (x1, y1, x2, y2)209 :return: CIoU scores scaled between 0 and 1210 """211 epsilon = 1e-7 # Small value to prevent division by zero212 213 # Expand dimensions for broadcasting214 bboxes2 = np.expand_dims(bboxes2, 0) # Shape: (1, M, 4)215 bboxes1 = np.expand_dims(bboxes1, 1) # Shape: (N, 1, 4)216 217 # Calculate the intersection box218 xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])219 yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])220 xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])221 yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])222 w = np.maximum(0.0, xx2 - xx1)223 h = np.maximum(0.0, yy2 - yy1)224 wh = w * h225 226 # Calculate IoU227 area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])228 area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1])229 iou = wh / (area1 + area2 - wh + epsilon)230 231 # Calculate center points232 centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0233 centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0234 centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0235 centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0236 237 # Calculate squared center distance238 inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2239 240 # Calculate smallest enclosing box diagonal241 xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])242 yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])243 xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])244 yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])245 outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2 + epsilon246 247 # Calculate aspect ratio consistency248 w1 = bboxes1[..., 2] - bboxes1[..., 0]249 h1 = bboxes1[..., 3] - bboxes1[..., 1]250 w2 = bboxes2[..., 2] - bboxes2[..., 0]251 h2 = bboxes2[..., 3] - bboxes2[..., 1]252 253 # Prevent division by zero254 h2 = h2 + epsilon255 h1 = h1 + epsilon256 arctan_diff = np.arctan(w2 / h2) - np.arctan(w1 / h1)257 v = (4 / (np.pi ** 2)) * (arctan_diff ** 2)258 259 # Calculate alpha260 S = 1 - iou261 alpha = v / (S + v + epsilon)262 263 # Compute CIoU264 ciou = iou - (inner_diag / outer_diag) + (alpha * v)265 266 # Scale CIoU to [0, 1]267 return (ciou + 1) / 2.0268 269 270 def diou_batch(bboxes1, bboxes2) -> np.ndarray:271 """272 :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2)273 :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2)274 :return:275 """276 # for details should go to https://arxiv.org/pdf/1902.09630.pdf277 # ensure predict's bbox form278 bboxes2 = np.expand_dims(bboxes2, 0)279 bboxes1 = np.expand_dims(bboxes1, 1)280 281 # calculate the intersection box282 xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])283 yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])284 xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])285 yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])286 w = np.maximum(0.0, xx2 - xx1)287 h = np.maximum(0.0, yy2 - yy1)288 wh = w * h289 iou = wh / (290 (bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1]) +291 (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) -292 wh293 )294 295 centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0296 centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0297 centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0298 centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0299 300 inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2301 302 xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])303 yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])304 xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])305 yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])306 307 outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2308 diou = iou - inner_diag / outer_diag309 310 return (diou + 1) / 2.0 311 312 313 @staticmethod314 def run_asso_func(self, bboxes1, bboxes2):315 """316 Runs the selected association function (based on the initialization string) on the input bounding boxes.317 318 Parameters:319 bboxes1: First set of bounding boxes.320 bboxes2: Second set of bounding boxes.321 """322 return self.asso_func(bboxes1, bboxes2)323 324 def _get_asso_func(self, asso_mode):325 """326 Returns the corresponding association function based on the provided mode string.327 328 Parameters:329 asso_mode (str): The association function to use (e.g., "iou", "giou", "centroid", etc.).330 331 Returns:332 function: The appropriate function for the association calculation.333 """334 ASSO_FUNCS = {335 "iou": AssociationFunction.iou_batch,336 "iou_obb": AssociationFunction.iou_batch_obb,337 "hmiou": AssociationFunction.hmiou_batch,338 "giou": AssociationFunction.giou_batch,339 "ciou": AssociationFunction.ciou_batch,340 "diou": AssociationFunction.diou_batch,341 "centroid": self.centroid_batch, # only not being staticmethod342 "centroid_obb": self.centroid_batch_obb343 }344 345 if self.asso_mode not in ASSO_FUNCS:346 raise ValueError(f"Invalid association mode: {self.asso_mode}. Choose from {list(ASSO_FUNCS.keys())}")347 348 return ASSO_FUNCS[self.asso_mode]349