Switch tracking to opencv - much better

.vscode/settings.json

@@ -1,4 +1,7 @@
 {
     "python.pythonPath": "~/.virtualenvs/obj-tracking/bin/python",
-    "python.formatting.provider": "black"
+    "python.formatting.provider": "black",
+    "python.linting.pylintArgs": [
+        "--extension-pkg-whitelist=cv2"
+    ]
 }

examples/opencv_objs_tracking.py

@@ -66,6 +66,10 @@ while True:
         (x, y, w, h) = [int(v) for v in box]
         cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
 
+    # objs = trackers.getObjects()
+    # for o in objs:
+    #     print(type(o), o)
+
     # show the output frame
     cv2.imshow("Frame", frame)
     key = cv2.waitKey(1) & 0xFF
@@ -81,8 +85,9 @@ while True:
         # create a new object tracker for the bounding box and add it
         # to our multi-object tracker
         tracker = OPENCV_OBJECT_TRACKERS[args["tracker"]]()
-        trackers.add(tracker, frame, box)
-
+        # print('handpick box', type(box), box)
+        n = trackers.add(tracker, frame, box)
+        # print(type(n), n)
     # if the `q` key was pressed, break from the loop
     elif key == ord("q"):
         break

src/main.py

@@ -10,13 +10,12 @@ import argparse
 import imutils
 import time
 import cv2
-import os
-import dlib
-from utils import check_if_inside_the_boxes, is_it_the_same_obj, distance
+from utils import check_if_inside_the_boxes, is_it_the_same_obj, box_distance
 
 # tracking
-
 OPENCV_OBJECT_TRACKERS = {"csrt": cv2.TrackerCSRT_create}
+# initialize OpenCV's special multi-object tracker
+cv_trackers = cv2.MultiTracker_create()
 trackers = []
 finished = []
 
@@ -185,8 +184,53 @@ while True:
 
     _frame_count += 1
 
-    # for dlib
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    # grab the updated bounding box coordinates (if any) for each
+    # object that is being tracked
+    (success, boxes) = cv_trackers.update(frame)
+    print("success", success)
+    print("boxes", boxes)
+
+    untracking = []
+    # loop over the bounding boxes and draw then on the frame
+    if success:
+        obj_cnt = len(boxes)
+        print(f'obj_cnt: ', obj_cnt, boxes)
+        for idx in range(obj_cnt):
+            box = boxes[idx]
+            (x, y, w, h) = [int(v) for v in box]
+            _last_pos = trackers[idx]["curr_position"]
+            (_x, _y, _w, _h) = _last_pos
+            curr_distance = box_distance(box, _last_pos)
+            last_distance = trackers[idx]["distance"]
+            trackers[idx]["distance"] = curr_distance
+            trackers[idx]["curr_position"] = [int(v) for v in box]
+
+            STILL_DIST_PX = 2
+            if last_distance < STILL_DIST_PX and curr_distance < STILL_DIST_PX:
+                trackers[idx]["still"] += 1
+            else:
+                trackers[idx]["still"] = 0
+
+            if trackers[idx]["still"] > 30 or x < 5 or x > 1250:
+                untracking.append(trackers[idx])
+
+            # DRAW on FRAME
+            tk = trackers[idx]
+            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
+            # cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)
+            color = [int(c) for c in COLORS[0]]
+            print(
+                f"{tk['id']} - {tk['type']} - centroid: {x}, {y} - distance: [stl:{tk['still']}] {last_distance:.3f} -> {curr_distance:.3f}"
+            )
+            cv2.putText(
+                frame,
+                f"{tk['id']} - {tk['type']}",
+                (x, y - 5),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                color,
+                2,
+            )
 
     # only detect once a sec
     if _frame_count % 15 == 1:
@@ -205,17 +249,12 @@ while True:
             # (1) check whether it's the same object as one in trackers
             is_same = False
             for t in trackers:
-                tracker = t["tracker"]
                 if _class != t["type"]:
                     continue
 
-                pos = tracker.get_position()
-                i = int(pos.left())
-                j = int(pos.top())
-                _w = int(pos.right()) - i
-                _h = int(pos.bottom()) - j
+                _x, _y, _w, _h = t["curr_position"]
                 print(f"[{t['id']}] - {t['type']}")
-                is_same = is_it_the_same_obj(x, y, w, h, i, j, _w, _h, id=t["id"])
+                is_same = is_it_the_same_obj(x, y, w, h, _x, _y, _w, _h, id=t["id"])
                 if is_same:
                     break
 
@@ -225,82 +264,31 @@ while True:
                 # to our multi-object tracker
                 # tracker = OPENCV_OBJECT_TRACKERS[args["tracker"]]()
                 # trackers.add(tracker, frame, box)
-                tracker = dlib.correlation_tracker()
-                rect = dlib.rectangle(x, y, x + w, y + h)
-                print("NEW TRACKER rect", rect)
-
+                _tracker = OPENCV_OBJECT_TRACKERS["csrt"]()
+                bbox = (x, y, w, h)
+                cv_trackers.add(_tracker, frame, bbox)
                 t = {
                     "id": tracker_counter,
                     "type": _class,
-                    "tracker": tracker,
+                    "curr_position": bbox,
                     "direction": "",
-                    "last_distance": -1,
-                    "last_position": (x + w / 2, y + h / 2),
+                    "distance": -1,
+                    "last_position": bbox,
                     "still": 0,
                 }
                 tracker_counter += 1
-                tracker.start_track(frame_rgb, rect)
                 trackers.append(t)
-
+                print(f'trackers ADD - now total #{len(trackers)}')
                 print(f" i -> {i} ({x},{y}), {w},{h} ({x + w},{y + h})")
 
-            # # draw a bounding box rectangle and label on the frame
-            # color = [int(c) for c in COLORS[classIDs[i]]]
-            # cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
-            # text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])
-            # cv2.putText(
-            #     frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2
-            # )
-
         _what = ",".join([LABELS[c] for c in classIDs])
         print(f"[{_frame_count:08d}] :: {_what}")
 
-    untracking = []
-    for tk in trackers:
-        tk["tracker"].update(frame_rgb)
-        pos = tk["tracker"].get_position()
-
-        # unpack the position object
-        startX = int(pos.left())
-        startY = int(pos.top())
-        endX = int(pos.right())
-        endY = int(pos.bottom())
-
-        tcx, tcy = (startX + endX) / 2, (startY + endY) / 2
-        # calculate distance
-        _x, _y = tk["last_position"]
-        _d = distance(_x, _y, tcx, tcy)
-        _last_distance = tk["last_distance"]
-        tk["last_distance"] = _d
-        tk["last_position"] = (tcx, tcy)
-        STILL_DISTANCE_IN_PX = 2
-        if _last_distance < STILL_DISTANCE_IN_PX and _d < STILL_DISTANCE_IN_PX:
-            tk["still"] += 1
-        else:
-            tk["still"] = 0
-
-        if tk["still"] > 30 or tcx < 10 or tcx > 1200:
-            untracking.append(tk)
-
-        cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)
-        color = [int(c) for c in COLORS[0]]
-        print(
-            f"{tk['id']} - {tk['type']} - centroid: {tcx, tcy} - distance: [stl:{tk['still']}] {_last_distance:.3f} -> {_d:.3f}"
-        )
-        cv2.putText(
-            frame,
-            f"{tk['id']} - {tk['type']}",
-            (startX, startY - 5),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.5,
-            color,
-            2,
-        )
-
     # untracking
-    untracking_ids = [ut["id"] for ut in untracking]
-    trackers = [tk for tk in trackers if tk["id"] not in untracking_ids]
-    finished += untracking
+    print("untracking: ", [ut["id"] for ut in untracking])
+    # untracking_ids = [ut["id"] for ut in untracking]
+    # trackers = [tk for tk in trackers if tk["id"] not in untracking_ids]
+    # finished += untracking
 
     if args["live"]:
         cv2.imshow("Frame", frame)

src/utils.py

@@ -47,6 +47,14 @@ AREAS = [
 ]
 
 
+def box_distance(pos1, pos2):
+    x1, y1, w1, h1 = pos1
+    x2, y2, w2, h2 = pos2
+    tx1, tx2 = x1 + w1 / 2, x2 + w2 / 2
+    ty1, ty2 = y1 + h1 / 2, y2 + h2 / 2
+    return math.sqrt(math.pow(tx2 - tx1, 2) + math.pow(ty2 - ty1, 2))
+
+
 def distance(x2, y2, x1, y1):
     return math.sqrt(math.pow(x2 - x1, 2) + math.pow(y2 - y1, 2))