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

This commit is contained in:
didier amyot committed 2024-01-24 08:01:58 -05:00
1 parent eb581763f9
commit 59a67afc2d
3 files changed
+20347 -14634

No files matched your search

+17 -6
View File
@@ -79,7 +79,8 @@ class DodgeballTracker:
v = np.array([xavg, ymax, 1]) v = np.array([xavg, ymax, 1])
ret = self.transformation_matrix @ v ret = self.transformation_matrix @ v
x, y = ret[0] / ret[2], ret[1] / ret[2] 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): def plot_boxe_to_frame(self, data, frame, frame_idx, frame_ts):
confidence = float(data[4]) confidence = float(data[4])
@@ -87,7 +88,6 @@ class DodgeballTracker:
if detected_box_object.min_threshold < confidence: if detected_box_object.min_threshold < confidence:
xmin, ymin, xmax, ymax = self.extract_box(data) 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) cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
return data # [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])] 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) cv2.polylines(frame, pts, True, (0, 255, 255), 8)
def persist_points(self): 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_: with open(self.output_json_path, "wt") as file_:
json.dump({"points": self.points}, file_, indent=4) json.dump({"points": self.points}, file_, indent=4)
@@ -127,14 +134,18 @@ class DodgeballTracker:
for i, (frame, ts) in enumerate(iter_frames(video_cap)): for i, (frame, ts) in enumerate(iter_frames(video_cap)):
# run the YOLO model on the frame # run the YOLO model on the frame
res = []
detections = self.run_model_on_frame(model, frame) detections = self.run_model_on_frame(model, frame)
results = []
for data in detections.boxes.data.tolist(): for data in detections.boxes.data.tolist():
results.append(self.plot_boxe_to_frame(data, frame, i, ts)) selected_data = self.plot_boxe_to_frame(data, frame, i, ts)
if selected_data:
tracks = tracker.update(np.array([r for r in results if r is not None]), img=frame) res.append(selected_data)
# show the frame to our screen # 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) self.write_frame(writer, frame)
if cv2.waitKey(1) == ord("q"): if cv2.waitKey(1) == ord("q"):
+1 -1
View File
@@ -70,7 +70,7 @@ function generatePoints(points, frames){
const lst = []; const lst = [];
for(const point of points.points){ for(const point of points.points){
if(frames.indexOf(point.frameIndex) !== -1){ 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); return m("g", lst);
+20329 -14627
View File
File diff suppressed because it is too large. Load diff