group by label before tracking objects

This commit is contained in:
Blake Blackshear 2020-01-09 06:52:28 -06:00
parent 32b212c7b6
commit be5a114f6a

View File

@ -251,6 +251,7 @@ class ObjectTracker(threading.Thread):
def register(self, index, obj): def register(self, index, obj):
id = f"{str(obj['frame_time'])}-{index}" id = f"{str(obj['frame_time'])}-{index}"
obj['id'] = id
self.tracked_objects[id] = obj self.tracked_objects[id] = obj
self.disappeared[id] = 0 self.disappeared[id] = 0
@ -281,20 +282,30 @@ class ObjectTracker(threading.Thread):
# to update # to update
return return
# compute centroids # group by name
new_object_groups = defaultdict(lambda: [])
for obj in new_objects: for obj in new_objects:
new_object_groups[obj['name']].append(obj)
# track objects for each label type
# TODO: this is going to miss deregistering objects that are not in the new groups
for label, group in new_object_groups.items():
current_objects = [o for o in self.tracked_objects.values() if o['name'] == label]
current_ids = [o['id'] for o in current_objects]
current_centroids = np.array([o['centroid'] for o in current_objects])
# compute centroids
for obj in group:
centroid_x = int((obj['box']['xmin']+obj['box']['xmax']) / 2.0) centroid_x = int((obj['box']['xmin']+obj['box']['xmax']) / 2.0)
centroid_y = int((obj['box']['ymin']+obj['box']['ymax']) / 2.0) centroid_y = int((obj['box']['ymin']+obj['box']['ymax']) / 2.0)
obj['centroid'] = (centroid_x, centroid_y) obj['centroid'] = (centroid_x, centroid_y)
if len(self.tracked_objects) == 0: if len(current_objects) == 0:
for index, obj in enumerate(new_objects): for index, obj in enumerate(group):
self.register(index, obj) self.register(index, obj)
return return
new_centroids = np.array([o['centroid'] for o in new_objects]) new_centroids = np.array([o['centroid'] for o in group])
current_ids = list(self.tracked_objects.keys())
current_centroids = np.array([o['centroid'] for o in self.tracked_objects.values()])
# compute the distance between each pair of tracked # compute the distance between each pair of tracked
# centroids and new centroids, respectively -- our # centroids and new centroids, respectively -- our