added a tracker that works and does not take too much time.

dodgeballtracker/logic.py

@@ -79,7 +79,8 @@ class DodgeballTracker:
         v = np.array([xavg, ymax, 1])
         ret = self.transformation_matrix @ v
         x, y = ret[0] / ret[2], ret[1] / ret[2]
-        self.points.append({"x": x, "y": y, "frameIndex": frame_idx, "frameTimestamp": frame_ts})
+        object_id = data[-1]
+        self.points.append({"x": x, "y": y, "frameIndex": frame_idx, "frameTimestamp": frame_ts, "objectId": object_id})
 
     def plot_boxe_to_frame(self, data, frame, frame_idx, frame_ts):
         confidence = float(data[4])
@@ -87,7 +88,6 @@ class DodgeballTracker:
 
             if detected_box_object.min_threshold < confidence:
                 xmin, ymin, xmax, ymax = self.extract_box(data)
-                self.print_box(data, frame_idx, frame_ts)
                 cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
                 return data  # [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])]
 
@@ -108,6 +108,13 @@ class DodgeballTracker:
         cv2.polylines(frame, pts, True, (0, 255, 255), 8)
 
     def persist_points(self):
+        nbr_id = max((p["objectId"] for p in self.points))
+        cm = matplotlib.colormaps['viridis']
+        for p in self.points:
+            value = p["objectId"] / nbr_id
+            color: Color = tuple((int(255 * cm(value)[i]) for i in range(3)))
+            p["color"] = f"rgb({color[0]}, {color[2]}, {color[1]})"
+
         with open(self.output_json_path, "wt") as file_:
             json.dump({"points": self.points}, file_, indent=4)
 
@@ -127,14 +134,18 @@ class DodgeballTracker:
 
         for i, (frame, ts) in enumerate(iter_frames(video_cap)):
             # run the YOLO model on the frame
+            res = []
             detections = self.run_model_on_frame(model, frame)
-            results = []
             for data in detections.boxes.data.tolist():
-                results.append(self.plot_boxe_to_frame(data, frame, i, ts))
-
-            tracks = tracker.update(np.array([r for r in results if r is not None]), img=frame)
+                selected_data = self.plot_boxe_to_frame(data, frame, i, ts)
+                if selected_data:
+                    res.append(selected_data)
 
             # show the frame to our screen
+            data_ = tracker.update(np.array(res), frame)
+            for d in data_:
+                self.print_box(d, i, ts)
+
             self.write_frame(writer, frame)
 
             if cv2.waitKey(1) == ord("q"):

frontend/main.js

@@ -70,7 +70,7 @@ function generatePoints(points, frames){
     const lst = [];
     for(const point of points.points){
         if(frames.indexOf(point.frameIndex) !== -1){
-            lst.push(m("circle", {cx: point.y, cy: point.x, r: 3}));
+            lst.push(m("circle", {cx: point.y, cy: point.x, r: 3, fill: point.color}));
         }
     }
     return m("g", lst);