blakeblackshear.frigate/frigate/track/centroid_tracker.py

242 lines
9.3 KiB
Python
Raw Normal View History

import random
import string
from collections import defaultdict
2020-11-04 13:31:25 +01:00
import numpy as np
2019-12-31 21:59:22 +01:00
from scipy.spatial import distance as dist
2020-11-04 13:31:25 +01:00
from frigate.config import DetectConfig
2023-05-31 16:12:43 +02:00
from frigate.track import ObjectTracker
from frigate.util import intersection_over_union
2020-11-04 13:31:25 +01:00
2019-02-26 03:27:02 +01:00
2023-05-31 16:12:43 +02:00
class CentroidTracker(ObjectTracker):
def __init__(self, config: DetectConfig):
2019-12-31 21:59:22 +01:00
self.tracked_objects = {}
2023-10-24 02:50:22 +02:00
self.untracked_object_boxes = []
2020-02-16 04:07:54 +01:00
self.disappeared = {}
self.positions = {}
self.max_disappeared = config.max_disappeared
self.detect_config = config
2019-12-31 21:59:22 +01:00
def register(self, index, obj):
2021-02-17 14:23:32 +01:00
rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
id = f"{obj['frame_time']}-{rand_id}"
2021-02-17 14:23:32 +01:00
obj["id"] = id
obj["start_time"] = obj["frame_time"]
obj["motionless_count"] = 0
obj["position_changes"] = 0
2019-12-31 21:59:22 +01:00
self.tracked_objects[id] = obj
2020-02-16 04:07:54 +01:00
self.disappeared[id] = 0
self.positions[id] = {
"xmins": [],
"ymins": [],
"xmaxs": [],
"ymaxs": [],
"xmin": 0,
"ymin": 0,
"xmax": self.detect_config.width,
"ymax": self.detect_config.height,
}
2019-12-31 21:59:22 +01:00
def deregister(self, id):
del self.tracked_objects[id]
2020-02-16 04:07:54 +01:00
del self.disappeared[id]
2021-02-17 14:23:32 +01:00
2022-02-08 14:40:45 +01:00
# tracks the current position of the object based on the last N bounding boxes
# returns False if the object has moved outside its previous position
def update_position(self, id, box):
position = self.positions[id]
2022-02-05 15:29:01 +01:00
position_box = (
position["xmin"],
position["ymin"],
position["xmax"],
position["ymax"],
)
xmin, ymin, xmax, ymax = box
2022-02-05 15:29:01 +01:00
iou = intersection_over_union(position_box, box)
2022-02-05 15:29:01 +01:00
# if the iou drops below the threshold
# assume the object has moved to a new position and reset the computed box
2022-02-05 15:29:01 +01:00
if iou < 0.6:
self.positions[id] = {
"xmins": [xmin],
"ymins": [ymin],
"xmaxs": [xmax],
"ymaxs": [ymax],
"xmin": xmin,
"ymin": ymin,
"xmax": xmax,
"ymax": ymax,
}
return False
# if there are less than 10 entries for the position, add the bounding box
# and recompute the position box
if len(position["xmins"]) < 10:
position["xmins"].append(xmin)
position["ymins"].append(ymin)
position["xmaxs"].append(xmax)
position["ymaxs"].append(ymax)
# by using percentiles here, we hopefully remove outliers
position["xmin"] = np.percentile(position["xmins"], 15)
position["ymin"] = np.percentile(position["ymins"], 15)
position["xmax"] = np.percentile(position["xmaxs"], 85)
position["ymax"] = np.percentile(position["ymaxs"], 85)
return True
2022-02-13 17:20:38 +01:00
def is_expired(self, id):
obj = self.tracked_objects[id]
# get the max frames for this label type or the default
max_frames = self.detect_config.stationary.max_frames.objects.get(
obj["label"], self.detect_config.stationary.max_frames.default
)
# if there is no max_frames for this label type, continue
if max_frames is None:
return False
# if the object has exceeded the max_frames setting, deregister
if (
obj["motionless_count"] - self.detect_config.stationary.threshold
> max_frames
):
return True
2022-04-16 15:43:49 +02:00
return False
def update(self, id, new_obj):
self.disappeared[id] = 0
# update the motionless count if the object has not moved to a new position
if self.update_position(id, new_obj["box"]):
self.tracked_objects[id]["motionless_count"] += 1
2022-02-13 17:20:38 +01:00
if self.is_expired(id):
self.deregister(id)
return
else:
2022-02-10 04:27:33 +01:00
# register the first position change and then only increment if
# the object was previously stationary
if (
self.tracked_objects[id]["position_changes"] == 0
or self.tracked_objects[id]["motionless_count"]
>= self.detect_config.stationary.threshold
2022-02-10 04:27:33 +01:00
):
self.tracked_objects[id]["position_changes"] += 1
self.tracked_objects[id]["motionless_count"] = 0
2022-02-10 04:27:33 +01:00
2020-01-08 03:43:25 +01:00
self.tracked_objects[id].update(new_obj)
2019-12-31 21:59:22 +01:00
def update_frame_times(self, frame_time):
2022-02-13 17:20:38 +01:00
for id in list(self.tracked_objects.keys()):
self.tracked_objects[id]["frame_time"] = frame_time
2022-02-06 21:50:15 +01:00
self.tracked_objects[id]["motionless_count"] += 1
2022-02-13 17:20:38 +01:00
if self.is_expired(id):
self.deregister(id)
2023-05-31 16:12:43 +02:00
def match_and_update(self, frame_time, detections):
2020-01-09 13:52:28 +01:00
# group by name
2023-05-31 16:12:43 +02:00
detection_groups = defaultdict(lambda: [])
for obj in detections:
detection_groups[obj[0]].append(
2021-02-17 14:23:32 +01:00
{
"label": obj[0],
"score": obj[1],
"box": obj[2],
"area": obj[3],
"ratio": obj[4],
"region": obj[5],
2021-02-17 14:23:32 +01:00
"frame_time": frame_time,
}
)
# update any tracked objects with labels that are not
# seen in the current objects and deregister if needed
2020-02-23 18:18:00 +01:00
for obj in list(self.tracked_objects.values()):
2023-05-31 16:12:43 +02:00
if obj["label"] not in detection_groups:
2021-02-17 14:23:32 +01:00
if self.disappeared[obj["id"]] >= self.max_disappeared:
self.deregister(obj["id"])
else:
2021-02-17 14:23:32 +01:00
self.disappeared[obj["id"]] += 1
2023-05-31 16:12:43 +02:00
if len(detections) == 0:
return
2021-02-17 14:23:32 +01:00
2020-01-09 13:52:28 +01:00
# track objects for each label type
2023-05-31 16:12:43 +02:00
for label, group in detection_groups.items():
2021-02-17 14:23:32 +01:00
current_objects = [
o for o in self.tracked_objects.values() if o["label"] == label
]
current_ids = [o["id"] for o in current_objects]
current_centroids = np.array([o["centroid"] for o in current_objects])
2020-01-09 13:52:28 +01:00
2020-01-11 20:22:56 +01:00
# compute centroids of new objects
2020-01-09 13:52:28 +01:00
for obj in group:
2021-02-17 14:23:32 +01:00
centroid_x = int((obj["box"][0] + obj["box"][2]) / 2.0)
centroid_y = int((obj["box"][1] + obj["box"][3]) / 2.0)
obj["centroid"] = (centroid_x, centroid_y)
2020-01-09 13:52:28 +01:00
if len(current_objects) == 0:
for index, obj in enumerate(group):
self.register(index, obj)
continue
2021-02-17 14:23:32 +01:00
new_centroids = np.array([o["centroid"] for o in group])
2020-01-09 13:52:28 +01:00
# compute the distance between each pair of tracked
# centroids and new centroids, respectively -- our
# goal will be to match each current centroid to a new
2020-01-09 13:52:28 +01:00
# object centroid
D = dist.cdist(current_centroids, new_centroids)
# in order to perform this matching we must (1) find the smallest
# value in each row (i.e. the distance from each current object to
# the closest new object) and then (2) sort the row indexes based
# on their minimum values so that the row with the smallest
# distance (the best match) is at the *front* of the index list
2020-01-09 13:52:28 +01:00
rows = D.min(axis=1).argsort()
# next, we determine which new object each existing object matched
# against, and apply the same sorting as was applied previously
2020-01-09 13:52:28 +01:00
cols = D.argmin(axis=1)[rows]
# many current objects may register with each new object, so only
# match the closest ones. unique returns the indices of the first
# occurrences of each value, and because the rows are sorted by
# distance, this will be index of the closest match
_, index = np.unique(cols, return_index=True)
rows = rows[index]
cols = cols[index]
# loop over the combination of the (row, column) index tuples
for row, col in zip(rows, cols):
# grab the object ID for the current row, set its new centroid,
# and reset the disappeared counter
2019-12-31 21:59:22 +01:00
objectID = current_ids[row]
2020-01-11 20:22:56 +01:00
self.update(objectID, group[col])
2020-01-09 13:52:28 +01:00
# compute the row and column indices we have NOT yet examined
unusedRows = set(range(D.shape[0])).difference(rows)
unusedCols = set(range(D.shape[1])).difference(cols)
2020-01-09 13:52:28 +01:00
2020-02-16 04:07:54 +01:00
# in the event that the number of object centroids is
2021-02-17 14:23:32 +01:00
# equal or greater than the number of input centroids
# we need to check and see if some of these objects have
# potentially disappeared
2020-02-16 04:07:54 +01:00
if D.shape[0] >= D.shape[1]:
for row in unusedRows:
id = current_ids[row]
if self.disappeared[id] >= self.max_disappeared:
self.deregister(id)
else:
self.disappeared[id] += 1
2020-01-11 20:22:56 +01:00
# if the number of input centroids is greater
2020-01-09 13:52:28 +01:00
# than the number of existing object centroids we need to
# register each new input centroid as a trackable object
2020-02-16 04:07:54 +01:00
else:
for col in unusedCols:
self.register(col, group[col])