Format: text/x-python
⬇ Download File (1.8 KB)
   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