1 import math 2 3 4 class EuclideanDistTracker: 5 def __init__(self): 6 # Store the center positions of the objects 7 self.center_points = {} 8 # Keep the count of the IDs 9 # each time a new object id detected, the count will increase by one 10 self.id_count = 0 11 12 13 def update(self, objects_rect): 14 # Objects boxes and ids 15 objects_bbs_ids = [] 16 17 # Get center point of new object 18 for rect in objects_rect: 19 x, y, w, h = rect 20 cx = (x + x + w) // 2 21 cy = (y + y + h) // 2 22 23 # Find out if that object was detected already 24 same_object_detected = False 25 for id, pt in self.center_points.items(): 26 dist = math.hypot(cx - pt[0], cy - pt[1]) 27 28 if dist < 25: 29 self.center_points[id] = (cx, cy) 30 print(self.center_points) 31 objects_bbs_ids.append([x, y, w, h, id]) 32 same_object_detected = True 33 break 34 35 # New object is detected we assign the ID to that object 36 if same_object_detected is False: 37 self.center_points[self.id_count] = (cx, cy) 38 objects_bbs_ids.append([x, y, w, h, self.id_count]) 39 self.id_count += 1 40 41 # Clean the dictionary by center points to remove IDS not used anymore 42 new_center_points = {} 43 for obj_bb_id in objects_bbs_ids: 44 _, _, _, _, object_id = obj_bb_id 45 center = self.center_points[object_id] 46 new_center_points[object_id] = center 47 48 # Update dictionary with IDs not used removed 49 self.center_points = new_center_points.copy() 50 return objects_bbs_ids 51 52 53 54