diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py
index 7b69227..30daa1c 100644
--- a/dodgeballtracker/logic.py
+++ b/dodgeballtracker/logic.py
@@ -1,10 +1,9 @@
import json
-import sys
from pathlib import Path
from ultralytics import YOLO
import cv2
-from dodgeballtracker.helper import create_video_writer, iter_frames, time_tracker
+from dodgeballtracker.helper import create_video_writer, iter_frames
from deep_sort_realtime.deepsort_tracker import DeepSort
import matplotlib
import numpy as np
@@ -74,11 +73,8 @@ class DodgeballTracker:
def print_box(self, data, frame_idx, frame_ts):
xmin, ymin, xmax, ymax = self.extract_box(data)
- v = np.array([
- (xmax + xmin) / 2,
- ymax, # (ymax + ymin) / 2,
- 1
- ])
+ xavg = (xmax + xmin) / 2
+ 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})
@@ -91,7 +87,7 @@ class DodgeballTracker:
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
+ return [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])]
def run_model_on_frame(self, model, frame):
return model(frame)[0]
@@ -118,6 +114,7 @@ class DodgeballTracker:
video_cap = cv2.VideoCapture(str(self.input_path))
writer = create_video_writer(video_cap, str(self.output_path))
+ tracker = DeepSort(max_age=50)
# load the pre-trained YOLOv8n model
model = YOLO("yolov8n.pt")
@@ -125,8 +122,18 @@ class DodgeballTracker:
for i, (frame, ts) in enumerate(iter_frames(video_cap)):
# run the YOLO model on the frame
detections = self.run_model_on_frame(model, frame)
+ results = []
for data in detections.boxes.data.tolist():
- self.plot_boxe_to_frame(data, frame, i, ts)
+ results.append(self.plot_boxe_to_frame(data, frame, i, ts))
+
+ tracks = tracker.update_tracks([r for r in results if r is not None], frame=frame)
+ for track in tracks:
+ if not track.is_confirmed():
+ continue
+ track_id = track.track_id
+ ltrb = track.to_ltrb()
+ print(f"{track_id} => { ltrb }")
+
# show the frame to our screen
self.write_frame(writer, frame)
diff --git a/frontend/index.html b/frontend/index.html
index adfca9e..f6b66be 100644
--- a/frontend/index.html
+++ b/frontend/index.html
@@ -14,7 +14,7 @@
Hello, world!
-