diff --git a/setup.py b/setup.py index a26c358..878bc5a 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setuptools.setup( name="centroid_tracker", - version="0.0.3", + version="0.0.9", author="Ross Leitch", author_email="ross@end-game.com", description="A small centroid tracker library based off https://www.pyimagesearch.com/2018/07/23/simple-object-tracking-with-opencv/", diff --git a/tracker/centroidtracker.py b/tracker/centroidtracker.py index c58b67b..ae2631b 100644 --- a/tracker/centroidtracker.py +++ b/tracker/centroidtracker.py @@ -7,16 +7,19 @@ class CentroidTracker(): def __init__(self, maxDisappeared=50): self.nextObjectID = 0 self.objects = OrderedDict() + self.originRects = OrderedDict() self.disappeared = OrderedDict() self.maxDisappeared = maxDisappeared - def register(self, centroid): + def register(self, centroid, rect): + self.originRects[self.nextObjectID] = rect self.objects[self.nextObjectID] = centroid self.disappeared[self.nextObjectID] = 0 self.nextObjectID += 1 def deregister(self, objectID): + del self.originRects[objectID] del self.objects[objectID] del self.disappeared[objectID] @@ -42,12 +45,12 @@ def get_id(self, rect): def update(self, rects): if(len(rects) == 0): - for objectID in self.disappeared.keys(): + for objectID in list(self.disappeared.keys()): self.disappeared[objectID] += 1 if(self.disappeared[objectID] > self.maxDisappeared): self.deregister(objectID) - return self.objects + return self.objects, self.originRects inputCentroids = np.zeros((len(rects), 2), dtype="int") @@ -58,7 +61,9 @@ def update(self, rects): if(len(self.objects) == 0): for i in range(0, len(inputCentroids)): - self.register(inputCentroids[i]) + centroid = inputCentroids[i] + rect = rects[i] + self.register(centroid, rect) else: objectIDs = list(self.objects.keys()) @@ -79,6 +84,7 @@ def update(self, rects): objectID = objectIDs[row] self.objects[objectID] = inputCentroids[col] + self.originRects[objectID] = rects[col] self.disappeared[objectID] = 0 usedRows.add(row) @@ -99,6 +105,8 @@ def update(self, rects): else: for col in unusedCols: - self.register(inputCentroids[col]) + centroid = inputCentroids[col] + rect = rects[col] + self.register(centroid, rect) - return self.objects + return self.objects, self.originRects