blakeblackshear.frigate/frigate/objects.py

149 lines
5.6 KiB
Python
Raw Normal View History

2020-11-04 13:31:25 +01:00
import copy
2019-02-26 03:27:02 +01:00
import datetime
2019-12-31 21:59:22 +01:00
import itertools
2020-11-04 13:31:25 +01:00
import multiprocessing as mp
import random
import string
2020-11-04 13:31:25 +01:00
import threading
import time
from collections import defaultdict
2020-11-04 13:31:25 +01:00
import cv2
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.util import calculate_region, draw_box_with_label
2019-02-26 03:27:02 +01:00
2020-02-16 04:07:54 +01:00
class ObjectTracker():
def __init__(self, max_disappeared):
2019-12-31 21:59:22 +01:00
self.tracked_objects = {}
2020-02-16 04:07:54 +01:00
self.disappeared = {}
self.max_disappeared = max_disappeared
2019-12-31 21:59:22 +01:00
def register(self, index, obj):
rand_id = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
id = f"{obj['frame_time']}-{rand_id}"
2020-01-09 13:52:28 +01:00
obj['id'] = id
2020-07-09 13:57:16 +02:00
obj['start_time'] = obj['frame_time']
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
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]
2019-12-31 21:59:22 +01:00
def update(self, id, new_obj):
2020-02-16 04:07:54 +01:00
self.disappeared[id] = 0
2020-01-08 03:43:25 +01:00
self.tracked_objects[id].update(new_obj)
2019-12-31 21:59:22 +01:00
2020-02-16 04:07:54 +01:00
def match_and_update(self, frame_time, new_objects):
2020-01-09 13:52:28 +01:00
# group by name
new_object_groups = defaultdict(lambda: [])
2019-12-31 21:59:22 +01:00
for obj in new_objects:
2020-02-16 04:07:54 +01:00
new_object_groups[obj[0]].append({
'label': obj[0],
'score': obj[1],
'box': obj[2],
'area': obj[3],
'region': obj[4],
'frame_time': frame_time
})
2019-12-31 21:59:22 +01:00
# 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()):
if not obj['label'] in new_object_groups:
2020-02-23 18:18:00 +01:00
if self.disappeared[obj['id']] >= self.max_disappeared:
self.deregister(obj['id'])
else:
2020-02-23 18:18:00 +01:00
self.disappeared[obj['id']] += 1
if len(new_objects) == 0:
return
2020-01-09 13:52:28 +01:00
# track objects for each label type
for label, group in new_object_groups.items():
2020-02-16 04:07:54 +01:00
current_objects = [o for o in self.tracked_objects.values() if o['label'] == label]
2020-01-09 13:52:28 +01:00
current_ids = [o['id'] for o in current_objects]
current_centroids = np.array([o['centroid'] for o in current_objects])
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:
2020-02-16 04:07:54 +01:00
centroid_x = int((obj['box'][0]+obj['box'][2]) / 2.0)
centroid_y = int((obj['box'][1]+obj['box'][3]) / 2.0)
2020-01-09 13:52:28 +01:00
obj['centroid'] = (centroid_x, centroid_y)
if len(current_objects) == 0:
for index, obj in enumerate(group):
self.register(index, obj)
return
new_centroids = np.array([o['centroid'] for o in group])
# compute the distance between each pair of tracked
# centroids and new centroids, respectively -- our
# goal will be to match each new centroid to an existing
# 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 and then (2) sort the row
# indexes based on their minimum values so that the row
# with the smallest value is at the *front* of the index
# list
rows = D.min(axis=1).argsort()
# next, we perform a similar process on the columns by
# finding the smallest value in each column and then
# sorting using the previously computed row index list
cols = D.argmin(axis=1)[rows]
# in order to determine if we need to update, register,
# or deregister an object we need to keep track of which
# of the rows and column indexes we have already examined
usedRows = set()
usedCols = set()
# loop over the combination of the (row, column) index
# tuples
for (row, col) in zip(rows, cols):
# if we have already examined either the row or
# column value before, ignore it
if row in usedRows or col in usedCols:
continue
2019-12-31 21:59:22 +01:00
2020-01-09 13:52:28 +01:00
# otherwise, 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
# indicate that we have examined each of the row and
# column indexes, respectively
usedRows.add(row)
usedCols.add(col)
2020-01-11 20:22:56 +01:00
# compute the column index we have NOT yet examined
2020-02-16 04:07:54 +01:00
unusedRows = set(range(0, D.shape[0])).difference(usedRows)
2020-01-09 13:52:28 +01:00
unusedCols = set(range(0, D.shape[1])).difference(usedCols)
2020-02-16 04:07:54 +01:00
# in the event that the number of object centroids is
# equal or greater than the number of input centroids
# we need to check and see if some of these objects have
# potentially disappeared
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])