added basic tracker. It takes way too much time.
This commit is contained in:
1 parent
e6c609fb08
commit
be436022a4
4 files changed
+21616
-11
No files matched your search
@@ -1,10 +1,9 @@
|
|||||||
import json
|
import json
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from ultralytics import YOLO
|
from ultralytics import YOLO
|
||||||
import cv2
|
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
|
from deep_sort_realtime.deepsort_tracker import DeepSort
|
||||||
import matplotlib
|
import matplotlib
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -74,11 +73,8 @@ class DodgeballTracker:
|
|||||||
|
|
||||||
def print_box(self, data, frame_idx, frame_ts):
|
def print_box(self, data, frame_idx, frame_ts):
|
||||||
xmin, ymin, xmax, ymax = self.extract_box(data)
|
xmin, ymin, xmax, ymax = self.extract_box(data)
|
||||||
v = np.array([
|
xavg = (xmax + xmin) / 2
|
||||||
(xmax + xmin) / 2,
|
v = np.array([xavg, ymax, 1])
|
||||||
ymax, # (ymax + ymin) / 2,
|
|
||||||
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})
|
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)
|
xmin, ymin, xmax, ymax = self.extract_box(data)
|
||||||
self.print_box(data, frame_idx, frame_ts)
|
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
|
return [[xmin, ymin, xmax - xmin, ymax - ymin], confidence, int(data[5])]
|
||||||
|
|
||||||
def run_model_on_frame(self, model, frame):
|
def run_model_on_frame(self, model, frame):
|
||||||
return model(frame)[0]
|
return model(frame)[0]
|
||||||
@@ -118,6 +114,7 @@ class DodgeballTracker:
|
|||||||
|
|
||||||
video_cap = cv2.VideoCapture(str(self.input_path))
|
video_cap = cv2.VideoCapture(str(self.input_path))
|
||||||
writer = create_video_writer(video_cap, str(self.output_path))
|
writer = create_video_writer(video_cap, str(self.output_path))
|
||||||
|
tracker = DeepSort(max_age=50)
|
||||||
|
|
||||||
# load the pre-trained YOLOv8n model
|
# load the pre-trained YOLOv8n model
|
||||||
model = YOLO("yolov8n.pt")
|
model = YOLO("yolov8n.pt")
|
||||||
@@ -125,8 +122,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
|
||||||
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():
|
||||||
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
|
# show the frame to our screen
|
||||||
self.write_frame(writer, frame)
|
self.write_frame(writer, frame)
|
||||||
|
|||||||
+1
-1
@@ -14,7 +14,7 @@
|
|||||||
<h1>Hello, world!</h1>
|
<h1>Hello, world!</h1>
|
||||||
|
|
||||||
<h3 id="first-ts"></h3>
|
<h3 id="first-ts"></h3>
|
||||||
<svg width="450" height="900" id="first" style="background-color: lightgray">
|
<svg width="900" height="450" id="first" style="background-color: lightgray">
|
||||||
</svg>
|
</svg>
|
||||||
|
|
||||||
<br>
|
<br>
|
||||||
|
|||||||
+1
-1
@@ -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.x, cy: point.y, r: 3}));
|
lst.push(m("circle", {cx: point.y, cy: point.x, r: 3}));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return m("g", lst);
|
return m("g", lst);
|
||||||
|
|||||||
+21598
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user