浏览代码

improve method for determining position

compares the centroid to a history of bounding boxes
Blake Blackshear 3 年之前
父节点
当前提交
338e4004d4
共有 1 个文件被更改,包括 60 次插入4 次删除
  1. 60 4
      frigate/objects.py

+ 60 - 4
frigate/objects.py

@@ -20,7 +20,9 @@ class ObjectTracker:
     def __init__(self, config: DetectConfig):
         self.tracked_objects = {}
         self.disappeared = {}
+        self.positions = {}
         self.max_disappeared = config.max_disappeared
+        self.detect_config = config
 
     def register(self, index, obj):
         rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
@@ -30,17 +32,71 @@ class ObjectTracker:
         obj["motionless_count"] = 0
         self.tracked_objects[id] = obj
         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,
+        }
 
     def deregister(self, id):
         del self.tracked_objects[id]
         del self.disappeared[id]
 
-    def update(self, id, new_obj):
-        self.disappeared[id] = 0
+    # tracks the current position of the object based on the last 10 bounding boxes
+    # returns False if the object has moved outside its previous position
+    def update_position(self, id, box):
+        position = self.positions[id]
+        xmin, ymin, xmax, ymax = box
+
+        # get the centroid
+        x = (xmax + xmin) / 2
+        y = (ymax + ymin) / 2
+
+        # if the centroid of this box is outside the computed bounding box
+        # assume the object has moved to a new position and reset the computed box
+        # TODO: should this only happen if there are a few boxes?
         if (
-            intersection_over_union(self.tracked_objects[id]["box"], new_obj["box"])
-            > 0.9
+            x < position["xmin"]
+            or x > position["xmax"]
+            or y < position["ymin"]
+            or y > position["ymax"]
         ):
+            position = {
+                "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
+
+    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
         else:
             self.tracked_objects[id]["motionless_count"] = 0