1 # import the necessary packages 2 from scipy.spatial import distance as dist 3 from collections import OrderedDict 4 import numpy as np 5 6 7 class CentroidTracker: 8 def __init__(self, maxDisappeared=50, maxDistance=50): 9 # initialize the next unique object ID along with two ordered 10 # dictionaries used to keep track of mapping a given object 11 # ID to its centroid and number of consecutive frames it has 12 # been marked as "disappeared", respectively 13 self.nextObjectID = 0 14 self.objects = OrderedDict() 15 self.disappeared = OrderedDict() 16 self.bbox = OrderedDict() # CHANGE 17 18 # store the number of maximum consecutive frames a given 19 # object is allowed to be marked as "disappeared" until we 20 # need to deregister the object from tracking 21 self.maxDisappeared = maxDisappeared 22 23 # store the maximum distance between centroids to associate 24 # an object -- if the distance is larger than this maximum 25 # distance we'll start to mark the object as "disappeared" 26 self.maxDistance = maxDistance 27 28 def register(self, centroid, inputRect): 29 # when registering an object we use the next available object 30 # ID to store the centroid 31 self.objects[self.nextObjectID] = centroid 32 self.bbox[self.nextObjectID] = inputRect # CHANGE 33 self.disappeared[self.nextObjectID] = 0 34 self.nextObjectID += 1 35 36 def deregister(self, objectID): 37 # to deregister an object ID we delete the object ID from 38 # both of our respective dictionaries 39 del self.objects[objectID] 40 del self.disappeared[objectID] 41 del self.bbox[objectID] # CHANGE 42 43 def update(self, rects): 44 # check to see if the list of input bounding box rectangles 45 # is empty 46 if len(rects) == 0: 47 # loop over any existing tracked objects and mark them 48 # as disappeared 49 for objectID in list(self.disappeared.keys()): 50 self.disappeared[objectID] += 1 51 52 # if we have reached a maximum number of consecutive 53 # frames where a given object has been marked as 54 # missing, deregister it 55 if self.disappeared[objectID] > self.maxDisappeared: 56 self.deregister(objectID) 57 58 # return early as there are no centroids or tracking info 59 # to update 60 # return self.objects 61 return self.bbox 62 63 # initialize an array of input centroids for the current frame 64 inputCentroids = np.zeros((len(rects), 2), dtype="int") 65 inputRects = [] 66 # loop over the bounding box rectangles 67 for (i, (startX, startY, endX, endY)) in enumerate(rects): 68 # use the bounding box coordinates to derive the centroid 69 cX = int((startX + endX) / 2.0) 70 cY = int((startY + endY) / 2.0) 71 inputCentroids[i] = (cX, cY) 72 inputRects.append(rects[i]) # CHANGE 73 74 # if we are currently not tracking any objects take the input 75 # centroids and register each of them 76 if len(self.objects) == 0: 77 for i in range(0, len(inputCentroids)): 78 self.register(inputCentroids[i], inputRects[i]) # CHANGE 79 80 # otherwise, are are currently tracking objects so we need to 81 # try to match the input centroids to existing object 82 # centroids 83 else: 84 # grab the set of object IDs and corresponding centroids 85 objectIDs = list(self.objects.keys()) 86 objectCentroids = list(self.objects.values()) 87 88 # compute the distance between each pair of object 89 # centroids and input centroids, respectively -- our 90 # goal will be to match an input centroid to an existing 91 # object centroid 92 D = dist.cdist(np.array(objectCentroids), inputCentroids) 93 94 # in order to perform this matching we must (1) find the 95 # smallest value in each row and then (2) sort the row 96 # indexes based on their minimum values so that the row 97 # with the smallest value as at the *front* of the index 98 # list 99 rows = D.min(axis=1).argsort() 100 101 # next, we perform a similar process on the columns by 102 # finding the smallest value in each column and then 103 # sorting using the previously computed row index list 104 cols = D.argmin(axis=1)[rows] 105 106 # in order to determine if we need to update, register, 107 # or deregister an object we need to keep track of which 108 # of the rows and column indexes we have already examined 109 usedRows = set() 110 usedCols = set() 111 112 # loop over the combination of the (row, column) index 113 # tuples 114 for (row, col) in zip(rows, cols): 115 # if we have already examined either the row or 116 # column value before, ignore it 117 if row in usedRows or col in usedCols: 118 continue 119 120 # if the distance between centroids is greater than 121 # the maximum distance, do not associate the two 122 # centroids to the same object 123 if D[row, col] > self.maxDistance: 124 continue 125 126 # otherwise, grab the object ID for the current row, 127 # set its new centroid, and reset the disappeared 128 # counter 129 objectID = objectIDs[row] 130 self.objects[objectID] = inputCentroids[col] 131 self.bbox[objectID] = inputRects[col] # CHANGE 132 self.disappeared[objectID] = 0 133 134 # indicate that we have examined each of the row and 135 # column indexes, respectively 136 usedRows.add(row) 137 usedCols.add(col) 138 139 # compute both the row and column index we have NOT yet 140 # examined 141 unusedRows = set(range(0, D.shape[0])).difference(usedRows) 142 unusedCols = set(range(0, D.shape[1])).difference(usedCols) 143 144 # in the event that the number of object centroids is 145 # equal or greater than the number of input centroids 146 # we need to check and see if some of these objects have 147 # potentially disappeared 148 if D.shape[0] >= D.shape[1]: 149 # loop over the unused row indexes 150 for row in unusedRows: 151 # grab the object ID for the corresponding row 152 # index and increment the disappeared counter 153 objectID = objectIDs[row] 154 self.disappeared[objectID] += 1 155 156 # check to see if the number of consecutive 157 # frames the object has been marked "disappeared" 158 # for warrants deregistering the object 159 if self.disappeared[objectID] > self.maxDisappeared: 160 self.deregister(objectID) 161 162 # otherwise, if the number of input centroids is greater 163 # than the number of existing object centroids we need to 164 # register each new input centroid as a trackable object 165 else: 166 for col in unusedCols: 167 self.register(inputCentroids[col], inputRects[col]) 168 169 # return the set of trackable objects 170 # return self.objects 171 return self.bbox 172 173