Tools for blacklisting tracks
This commit is contained in:
parent
30648b9bb8
commit
a590a0dc35
10 changed files with 388 additions and 96 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
@ -1,6 +1,7 @@
|
||||||
.idea/
|
.idea/
|
||||||
OUT/
|
OUT/
|
||||||
EXPERIMENTS/
|
EXPERIMENTS/
|
||||||
|
runs/
|
||||||
|
|
||||||
## Core latex/pdflatex auxiliary files:
|
## Core latex/pdflatex auxiliary files:
|
||||||
*.aux
|
*.aux
|
||||||
|
|
16
poetry.lock
generated
16
poetry.lock
generated
|
@ -1015,6 +1015,20 @@ files = [
|
||||||
[package.extras]
|
[package.extras]
|
||||||
dev = ["hypothesis"]
|
dev = ["hypothesis"]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "jsonlines"
|
||||||
|
version = "4.0.0"
|
||||||
|
description = "Library with helpers for the jsonlines file format"
|
||||||
|
optional = false
|
||||||
|
python-versions = ">=3.8"
|
||||||
|
files = [
|
||||||
|
{file = "jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55"},
|
||||||
|
{file = "jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74"},
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.dependencies]
|
||||||
|
attrs = ">=19.2.0"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "jsonpointer"
|
name = "jsonpointer"
|
||||||
version = "2.4"
|
version = "2.4"
|
||||||
|
@ -3725,4 +3739,4 @@ watchdog = ["watchdog (>=2.3)"]
|
||||||
[metadata]
|
[metadata]
|
||||||
lock-version = "2.0"
|
lock-version = "2.0"
|
||||||
python-versions = "^3.10,<3.12,"
|
python-versions = "^3.10,<3.12,"
|
||||||
content-hash = "bf4feafd4afa6ceb39a1c599e3e7cdc84afbe11ab1672b49e5de99ad44568b08"
|
content-hash = "71868029d1943c412082bcb705dd76711f105afa8a93326a45e129230de8ffa9"
|
||||||
|
|
|
@ -10,6 +10,8 @@ trapserv = "trap.plumber:start"
|
||||||
tracker = "trap.tools:tracker_preprocess"
|
tracker = "trap.tools:tracker_preprocess"
|
||||||
compare = "trap.tools:tracker_compare"
|
compare = "trap.tools:tracker_compare"
|
||||||
process_data = "trap.process_data:main"
|
process_data = "trap.process_data:main"
|
||||||
|
blacklist = "trap.tools:blacklist_tracks"
|
||||||
|
rewrite_tracks = "trap.tools:rewrite_raw_track_files"
|
||||||
|
|
||||||
|
|
||||||
[tool.poetry.dependencies]
|
[tool.poetry.dependencies]
|
||||||
|
@ -39,6 +41,7 @@ pyglet-cornerpin = "^0.3.0"
|
||||||
opencv-python = {file="./opencv_python-4.10.0.84-cp310-cp310-linux_x86_64.whl"}
|
opencv-python = {file="./opencv_python-4.10.0.84-cp310-cp310-linux_x86_64.whl"}
|
||||||
setproctitle = "^1.3.3"
|
setproctitle = "^1.3.3"
|
||||||
bytetracker = { git = "https://github.com/rubenvandeven/bytetrack-pip" }
|
bytetracker = { git = "https://github.com/rubenvandeven/bytetrack-pip" }
|
||||||
|
jsonlines = "^4.0.0"
|
||||||
|
|
||||||
[build-system]
|
[build-system]
|
||||||
requires = ["poetry-core"]
|
requires = ["poetry-core"]
|
||||||
|
|
|
@ -86,7 +86,7 @@ class CameraAction(argparse.Action):
|
||||||
# 'camera_matrix': np.array(data['camera_matrix']),
|
# 'camera_matrix': np.array(data['camera_matrix']),
|
||||||
# 'dist_coeff': np.array(data['dist_coeff']),
|
# 'dist_coeff': np.array(data['dist_coeff']),
|
||||||
# }
|
# }
|
||||||
camera = Camera(np.array(data['camera_matrix']), np.array(data['dist_coeff']), data['dim']['width'], data['dim']['height'], namespace.H)
|
camera = Camera(np.array(data['camera_matrix']), np.array(data['dist_coeff']), data['dim']['width'], data['dim']['height'], namespace.H, namespace.camera_fps)
|
||||||
|
|
||||||
setattr(namespace, 'camera', camera)
|
setattr(namespace, 'camera', camera)
|
||||||
|
|
||||||
|
@ -276,6 +276,10 @@ frame_emitter_parser.add_argument("--video-loop",
|
||||||
|
|
||||||
# Tracker
|
# Tracker
|
||||||
|
|
||||||
|
tracker_parser.add_argument("--camera-fps",
|
||||||
|
help="Camera FPS",
|
||||||
|
type=int,
|
||||||
|
default=12)
|
||||||
tracker_parser.add_argument("--homography",
|
tracker_parser.add_argument("--homography",
|
||||||
help="File with homography params",
|
help="File with homography params",
|
||||||
type=Path,
|
type=Path,
|
||||||
|
@ -314,6 +318,9 @@ tracker_parser.add_argument("--smooth-tracks",
|
||||||
|
|
||||||
|
|
||||||
# Renderer
|
# Renderer
|
||||||
|
# render_parser.add_argument("--disable-renderer",
|
||||||
|
# help="Disable the renderer all together. Usefull when using an external renderer",
|
||||||
|
# action="store_true")
|
||||||
|
|
||||||
render_parser.add_argument("--render-file",
|
render_parser.add_argument("--render-file",
|
||||||
help="Render a video file previewing the prediction, and its delay compared to the current frame",
|
help="Render a video file previewing the prediction, and its delay compared to the current frame",
|
||||||
|
|
|
@ -1,7 +1,11 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
from argparse import Namespace
|
from argparse import Namespace
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
import dataclasses
|
||||||
from enum import IntFlag
|
from enum import IntFlag
|
||||||
from itertools import cycle
|
from itertools import cycle
|
||||||
|
import json
|
||||||
import logging
|
import logging
|
||||||
from multiprocessing import Event
|
from multiprocessing import Event
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
@ -11,21 +15,36 @@ import time
|
||||||
from typing import Iterable, List, Optional
|
from typing import Iterable, List, Optional
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import cv2
|
import cv2
|
||||||
|
import pandas as pd
|
||||||
import zmq
|
import zmq
|
||||||
import os
|
import os
|
||||||
from deep_sort_realtime.deep_sort.track import Track as DeepsortTrack
|
from deep_sort_realtime.deep_sort.track import Track as DeepsortTrack
|
||||||
from deep_sort_realtime.deep_sort.track import TrackState as DeepsortTrackState
|
from deep_sort_realtime.deep_sort.track import TrackState as DeepsortTrackState
|
||||||
from bytetracker.byte_tracker import STrack as ByteTrackTrack
|
from bytetracker.byte_tracker import STrack as ByteTrackTrack
|
||||||
from bytetracker.basetrack import TrackState as ByteTrackTrackState
|
from bytetracker.basetrack import TrackState as ByteTrackTrackState
|
||||||
|
from trajectron.environment import Environment, Node, Scene
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
from trap.utils import lerp
|
||||||
|
|
||||||
logger = logging.getLogger('trap.frame_emitter')
|
logger = logging.getLogger('trap.frame_emitter')
|
||||||
|
|
||||||
|
class DataclassJSONEncoder(json.JSONEncoder):
|
||||||
|
def default(self, o):
|
||||||
|
if isinstance(o, np.ndarray):
|
||||||
|
return o.tolist()
|
||||||
|
if dataclasses.is_dataclass(o):
|
||||||
|
d = dataclasses.asdict(o)
|
||||||
|
if isinstance(o, Frame):
|
||||||
|
# Don't send images over JSON
|
||||||
|
del d['img']
|
||||||
|
return d
|
||||||
|
return super().default(o)
|
||||||
|
|
||||||
|
|
||||||
class UrlOrPath():
|
class UrlOrPath():
|
||||||
def __init__(self, str):
|
def __init__(self, string):
|
||||||
self.url = urlparse(str)
|
self.url = urlparse(str(string))
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return self.url.geturl()
|
return self.url.geturl()
|
||||||
|
@ -64,14 +83,29 @@ class DetectionState(IntFlag):
|
||||||
return cls.Confirmed
|
return cls.Confirmed
|
||||||
raise RuntimeError("Should not run into Deleted entries here")
|
raise RuntimeError("Should not run into Deleted entries here")
|
||||||
|
|
||||||
|
@dataclass
|
||||||
class Camera:
|
class Camera:
|
||||||
def __init__(self, mtx, dist, w, h, H):
|
mtx: cv2.Mat
|
||||||
self.mtx = mtx
|
dist: cv2.Mat
|
||||||
self.dist = dist
|
w: float
|
||||||
self.w = w
|
h: float
|
||||||
self.h = h
|
H: cv2.Mat # homography
|
||||||
self.newcameramtx, self.roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))
|
|
||||||
self.H = H # homography
|
newcameramtx: cv2.Mat = field(init=False)
|
||||||
|
roi: cv2.typing.Rect = field(init=False)
|
||||||
|
fps: float
|
||||||
|
|
||||||
|
def __post_init__(self):
|
||||||
|
self.newcameramtx, self.roi = cv2.getOptimalNewCameraMatrix(self.mtx, self.dist, (self.w,self.h), 1, (self.w,self.h))
|
||||||
|
|
||||||
|
|
||||||
|
# def __init__(self, mtx, dist, w, h, H):
|
||||||
|
# self.mtx = mtx
|
||||||
|
# self.dist = dist
|
||||||
|
# self.w = w
|
||||||
|
# self.h = h
|
||||||
|
# self.newcameramtx, self.roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))
|
||||||
|
# self.H = H # homography
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
@ -131,6 +165,7 @@ class Track:
|
||||||
history: List[Detection] = field(default_factory=lambda: [])
|
history: List[Detection] = field(default_factory=lambda: [])
|
||||||
predictor_history: Optional[list] = None # in image space
|
predictor_history: Optional[list] = None # in image space
|
||||||
predictions: Optional[list] = None
|
predictions: Optional[list] = None
|
||||||
|
fps: int = 12
|
||||||
|
|
||||||
def get_projected_history(self, H, camera: Optional[Camera]= None) -> np.array:
|
def get_projected_history(self, H, camera: Optional[Camera]= None) -> np.array:
|
||||||
foot_coordinates = [d.get_foot_coords() for d in self.history]
|
foot_coordinates = [d.get_foot_coords() for d in self.history]
|
||||||
|
@ -138,7 +173,7 @@ class Track:
|
||||||
if len(foot_coordinates):
|
if len(foot_coordinates):
|
||||||
if camera:
|
if camera:
|
||||||
coords = cv2.undistortPoints(np.array([foot_coordinates]).astype('float32'), camera.mtx, camera.dist, None, camera.newcameramtx)
|
coords = cv2.undistortPoints(np.array([foot_coordinates]).astype('float32'), camera.mtx, camera.dist, None, camera.newcameramtx)
|
||||||
coords = cv2.perspectiveTransform(np.array(coords),H)
|
coords = cv2.perspectiveTransform(np.array(coords),camera.H)
|
||||||
return coords.reshape((coords.shape[0],2))
|
return coords.reshape((coords.shape[0],2))
|
||||||
else:
|
else:
|
||||||
coords = cv2.perspectiveTransform(np.array([foot_coordinates]),H)
|
coords = cv2.perspectiveTransform(np.array([foot_coordinates]),H)
|
||||||
|
@ -149,8 +184,64 @@ class Track:
|
||||||
coords = self.get_projected_history(H, camera)
|
coords = self.get_projected_history(H, camera)
|
||||||
return [{"x":c[0], "y":c[1]} for c in coords]
|
return [{"x":c[0], "y":c[1]} for c in coords]
|
||||||
|
|
||||||
|
def get_with_interpolated_history(self) -> Track:
|
||||||
|
# new_history = [Detection(d.track_id, l, t, w, h, d.conf, d.state, d.frame_nr, d.det_class) for l, t, w, h, d in zip(ls,ts,ws,hs, track.history)]
|
||||||
|
# new_track = Track(track.track_id, new_history, track.predictor_history, track.predictions)
|
||||||
|
new_history = []
|
||||||
|
for j in range(len(self.history)-1):
|
||||||
|
a = self.history[j]
|
||||||
|
b = self.history[j+1]
|
||||||
|
gap = b.frame_nr - a.frame_nr
|
||||||
|
new_history.append(Detection(a.track_id, a.l, a.t, a.w, a.h, a.conf, a.state, a.frame_nr, a.det_class))
|
||||||
|
if gap < 1:
|
||||||
|
logger.error(f"WARNING, gap between frames {a.frame_nr} -> {b.frame_nr} is negative?")
|
||||||
|
if gap > 1:
|
||||||
|
for g in range(1, gap):
|
||||||
|
l = lerp(a.l, b.l, g/gap)
|
||||||
|
t = lerp(a.t, b.t, g/gap)
|
||||||
|
w = lerp(a.w, b.w, g/gap)
|
||||||
|
h = lerp(a.h, b.h, g/gap)
|
||||||
|
conf = 0
|
||||||
|
state = DetectionState.Lost
|
||||||
|
frame_nr = a.frame_nr + g
|
||||||
|
new_history.append(Detection(a.track_id, l, t, w, h, conf, state, frame_nr, a.det_class))
|
||||||
|
|
||||||
|
return Track(
|
||||||
|
self.track_id,
|
||||||
|
new_history,
|
||||||
|
self.predictor_history,
|
||||||
|
self.predictions,
|
||||||
|
self.fps)
|
||||||
|
|
||||||
|
|
||||||
|
def to_trajectron_node(self, camera: Camera, env: Environment) -> Node:
|
||||||
|
positions = self.get_projected_history(None, camera)
|
||||||
|
velocity = np.gradient(positions, self.fps, axis=0)
|
||||||
|
acceleration = np.gradient(velocity, self.fps, axis=0)
|
||||||
|
|
||||||
|
new_first_idx = self.history[0].frame_nr
|
||||||
|
|
||||||
|
data_columns = pd.MultiIndex.from_product([['position', 'velocity', 'acceleration'], ['x', 'y']])
|
||||||
|
|
||||||
|
|
||||||
|
# vx = derivative_of(x, scene.dt)
|
||||||
|
# vy = derivative_of(y, scene.dt)
|
||||||
|
# ax = derivative_of(vx, scene.dt)
|
||||||
|
# ay = derivative_of(vy, scene.dt)
|
||||||
|
|
||||||
|
data_dict = {
|
||||||
|
('position', 'x'): positions[:,0],
|
||||||
|
('position', 'y'): positions[:,1],
|
||||||
|
('velocity', 'x'): velocity[:,0],
|
||||||
|
('velocity', 'y'): velocity[:,1],
|
||||||
|
('acceleration', 'x'): acceleration[:,0],
|
||||||
|
('acceleration', 'y'): acceleration[:,1]
|
||||||
|
}
|
||||||
|
|
||||||
|
node_data = pd.DataFrame(data_dict, columns=data_columns)
|
||||||
|
|
||||||
|
return Node(node_type=env.NodeType.PEDESTRIAN, node_id=self.track_id, data=node_data, first_timestep=new_first_idx)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
|
@ -86,7 +86,7 @@ def start():
|
||||||
|
|
||||||
# instantiating process with arguments
|
# instantiating process with arguments
|
||||||
procs = [
|
procs = [
|
||||||
ExceptionHandlingProcess(target=run_ws_forwarder, kwargs={'config': args, 'is_running': isRunning}, name='forwarder'),
|
# ExceptionHandlingProcess(target=run_ws_forwarder, kwargs={'config': args, 'is_running': isRunning}, name='forwarder'),
|
||||||
ExceptionHandlingProcess(target=run_frame_emitter, kwargs={'config': args, 'is_running': isRunning}, name='frame_emitter'),
|
ExceptionHandlingProcess(target=run_frame_emitter, kwargs={'config': args, 'is_running': isRunning}, name='frame_emitter'),
|
||||||
ExceptionHandlingProcess(target=run_tracker, kwargs={'config': args, 'is_running': isRunning}, name='tracker'),
|
ExceptionHandlingProcess(target=run_tracker, kwargs={'config': args, 'is_running': isRunning}, name='tracker'),
|
||||||
]
|
]
|
||||||
|
|
|
@ -26,7 +26,7 @@ import matplotlib.pyplot as plt
|
||||||
|
|
||||||
import zmq
|
import zmq
|
||||||
|
|
||||||
from trap.frame_emitter import Frame
|
from trap.frame_emitter import DataclassJSONEncoder, Frame
|
||||||
from trap.tracker import Track, Smoother
|
from trap.tracker import Track, Smoother
|
||||||
|
|
||||||
logger = logging.getLogger("trap.prediction")
|
logger = logging.getLogger("trap.prediction")
|
||||||
|
@ -160,8 +160,16 @@ class PredictionServer:
|
||||||
|
|
||||||
self.prediction_socket: zmq.Socket = context.socket(zmq.PUB)
|
self.prediction_socket: zmq.Socket = context.socket(zmq.PUB)
|
||||||
self.prediction_socket.bind(config.zmq_prediction_addr)
|
self.prediction_socket.bind(config.zmq_prediction_addr)
|
||||||
|
self.external_predictions = not self.config.zmq_prediction_addr.startswith("ipc://")
|
||||||
# print(self.prediction_socket)
|
# print(self.prediction_socket)
|
||||||
|
|
||||||
|
def send_frame(self, frame: Frame):
|
||||||
|
if self.external_predictions:
|
||||||
|
# data = json.dumps(frame, cls=DataclassJSONEncoder)
|
||||||
|
self.prediction_socket.send_json(frame, cls=DataclassJSONEncoder)
|
||||||
|
else:
|
||||||
|
self.prediction_socket.send_pyobj(frame)
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
|
|
||||||
if self.config.seed is not None:
|
if self.config.seed is not None:
|
||||||
|
@ -275,12 +283,14 @@ class PredictionServer:
|
||||||
if self.config.predict_training_data:
|
if self.config.predict_training_data:
|
||||||
input_dict = eval_scene.get_clipped_input_dict(timestep, hyperparams['state'])
|
input_dict = eval_scene.get_clipped_input_dict(timestep, hyperparams['state'])
|
||||||
else:
|
else:
|
||||||
|
# print('await', self.config.zmq_trajectory_addr)
|
||||||
zmq_ev = self.trajectory_socket.poll(timeout=2000)
|
zmq_ev = self.trajectory_socket.poll(timeout=2000)
|
||||||
if not zmq_ev:
|
if not zmq_ev:
|
||||||
# on no data loop so that is_running is checked
|
# on no data loop so that is_running is checked
|
||||||
continue
|
continue
|
||||||
|
|
||||||
data = self.trajectory_socket.recv()
|
data = self.trajectory_socket.recv()
|
||||||
|
# print('recv tracker frame')
|
||||||
frame: Frame = pickle.loads(data)
|
frame: Frame = pickle.loads(data)
|
||||||
# trajectory_data = {t.track_id: t.get_projected_history_as_dict(frame.H) for t in frame.tracks.values()}
|
# trajectory_data = {t.track_id: t.get_projected_history_as_dict(frame.H) for t in frame.tracks.values()}
|
||||||
# trajectory_data = json.loads(data)
|
# trajectory_data = json.loads(data)
|
||||||
|
@ -330,7 +340,7 @@ class PredictionServer:
|
||||||
first_timestep=timestep
|
first_timestep=timestep
|
||||||
)
|
)
|
||||||
|
|
||||||
input_dict[node] = np.array([x[-1],y[-1],vx[-1],vy[-1],ax[-1],ay[-1]])
|
input_dict[node] = np.array(object=[x[-1],y[-1],vx[-1],vy[-1],ax[-1],ay[-1]])
|
||||||
|
|
||||||
# print(input_dict)
|
# print(input_dict)
|
||||||
|
|
||||||
|
@ -340,7 +350,7 @@ class PredictionServer:
|
||||||
# And want to update the network
|
# And want to update the network
|
||||||
|
|
||||||
# data = json.dumps({})
|
# data = json.dumps({})
|
||||||
self.prediction_socket.send_pyobj(frame)
|
self.send_frame(frame)
|
||||||
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
@ -444,7 +454,7 @@ class PredictionServer:
|
||||||
if self.config.smooth_predictions:
|
if self.config.smooth_predictions:
|
||||||
frame = self.smoother.smooth_frame_predictions(frame)
|
frame = self.smoother.smooth_frame_predictions(frame)
|
||||||
|
|
||||||
self.prediction_socket.send_pyobj(frame)
|
self.send_frame(frame)
|
||||||
time.sleep(.5)
|
time.sleep(.5)
|
||||||
logger.info('Stopping')
|
logger.info('Stopping')
|
||||||
|
|
||||||
|
|
117
trap/tools.py
117
trap/tools.py
|
@ -1,15 +1,17 @@
|
||||||
from argparse import Namespace
|
from argparse import Namespace
|
||||||
|
import json
|
||||||
import math
|
import math
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import pickle
|
import pickle
|
||||||
from tempfile import mktemp
|
from tempfile import mktemp
|
||||||
|
|
||||||
|
import jsonlines
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
import trap.tracker
|
import trap.tracker
|
||||||
from trap.config import parser
|
from trap.config import parser
|
||||||
from trap.frame_emitter import Detection, DetectionState, video_src_from_config, Frame
|
from trap.frame_emitter import Detection, DetectionState, video_src_from_config, Frame
|
||||||
from trap.tracker import DETECTOR_YOLOv8, Smoother, _yolov8_track, Track, TrainingDataWriter, Tracker
|
from trap.tracker import DETECTOR_YOLOv8, Smoother, _yolov8_track, Track, TrainingDataWriter, Tracker, read_tracks_json
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
@ -20,6 +22,8 @@ from ultralytics import YOLO
|
||||||
from ultralytics.engine.results import Results as YOLOResult
|
from ultralytics.engine.results import Results as YOLOResult
|
||||||
import tqdm
|
import tqdm
|
||||||
|
|
||||||
|
from trap.utils import lerp
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger('tools')
|
logger = logging.getLogger('tools')
|
||||||
|
@ -98,6 +102,7 @@ def tracker_preprocess():
|
||||||
|
|
||||||
total = 0
|
total = 0
|
||||||
frames = FrameGenerator(config)
|
frames = FrameGenerator(config)
|
||||||
|
total_tracks = set()
|
||||||
for frame in frames:
|
for frame in frames:
|
||||||
bar.update()
|
bar.update()
|
||||||
|
|
||||||
|
@ -105,7 +110,6 @@ def tracker_preprocess():
|
||||||
total += len(detections)
|
total += len(detections)
|
||||||
# detections = _yolov8_track(frame, model, imgsz=1440, classes=[0])
|
# detections = _yolov8_track(frame, model, imgsz=1440, classes=[0])
|
||||||
|
|
||||||
bar.set_description(f"{frames.video_nr}/{len(frames.video_srcs)} [{frames.frame_idx}/{frames.frame_count}] {marquee_string(str(frames.video_path), 10, frames.n//2)} | dets {len(detections)}: {[d.track_id for d in detections]} (∑{total})")
|
|
||||||
|
|
||||||
for detection in detections:
|
for detection in detections:
|
||||||
track = tracks[detection.track_id]
|
track = tracks[detection.track_id]
|
||||||
|
@ -114,6 +118,9 @@ def tracker_preprocess():
|
||||||
|
|
||||||
active_track_ids = [d.track_id for d in detections]
|
active_track_ids = [d.track_id for d in detections]
|
||||||
active_tracks = {t.track_id: t for t in tracks.values() if t.track_id in active_track_ids}
|
active_tracks = {t.track_id: t for t in tracks.values() if t.track_id in active_track_ids}
|
||||||
|
total_tracks.update(active_track_ids)
|
||||||
|
|
||||||
|
bar.set_description(f"{frames.video_nr}/{len(frames.video_srcs)} [{frames.frame_idx}/{frames.frame_count}] {marquee_string(str(frames.video_path), 10, frames.n//2)} | dets {len(detections)}: {[d.track_id for d in detections]} (∑{total} → {len(total_tracks)})")
|
||||||
|
|
||||||
writer.add(frame, active_tracks.values())
|
writer.add(frame, active_tracks.values())
|
||||||
|
|
||||||
|
@ -122,12 +129,13 @@ def tracker_preprocess():
|
||||||
bgr_colors = [
|
bgr_colors = [
|
||||||
(255, 0, 0),
|
(255, 0, 0),
|
||||||
(0, 255, 0),
|
(0, 255, 0),
|
||||||
(0, 0, 255),
|
# (0, 0, 255),# red used for missing waypoints
|
||||||
(0, 255, 255),
|
(0, 255, 255),
|
||||||
]
|
]
|
||||||
|
|
||||||
def detection_color(detection: Detection, i):
|
def detection_color(detection: Detection, i, prev_detection: Optional[Detection] = None):
|
||||||
return bgr_colors[i % len(bgr_colors)] if detection.state != DetectionState.Lost else (100,100,100)
|
vague = detection.state == DetectionState.Lost or (prev_detection and detection.frame_nr - prev_detection.frame_nr > 1)
|
||||||
|
return bgr_colors[i % len(bgr_colors)] if not vague else (0,0,255)
|
||||||
|
|
||||||
def to_point(coord):
|
def to_point(coord):
|
||||||
return (int(coord[0]), int(coord[1]))
|
return (int(coord[0]), int(coord[1]))
|
||||||
|
@ -164,13 +172,7 @@ def tracker_compare():
|
||||||
for i, (tracker, detections) in enumerate(trackers_detections):
|
for i, (tracker, detections) in enumerate(trackers_detections):
|
||||||
|
|
||||||
for track_id in tracker.tracks:
|
for track_id in tracker.tracks:
|
||||||
history = tracker.tracks[track_id].history
|
draw_track(frame.img, tracker.tracks[track_id], i)
|
||||||
cv2.putText(frame.img, f"{track_id}", to_point(history[0].get_foot_coords()), cv2.FONT_HERSHEY_DUPLEX, 1, color=bgr_colors[i % len(bgr_colors)])
|
|
||||||
for j in range(len(history)-1):
|
|
||||||
a = history[j]
|
|
||||||
b = history[j+1]
|
|
||||||
color = detection_color(b, i)
|
|
||||||
cv2.line(frame.img, to_point(a.get_foot_coords()), to_point(b.get_foot_coords()), color, 1)
|
|
||||||
for detection in detections:
|
for detection in detections:
|
||||||
color = color = detection_color(detection, i)
|
color = color = detection_color(detection, i)
|
||||||
l, t, r, b = detection.to_ltrb()
|
l, t, r, b = detection.to_ltrb()
|
||||||
|
@ -183,6 +185,97 @@ def tracker_compare():
|
||||||
|
|
||||||
bar.set_description(f"[{frames.video_nr}/{len(frames.video_srcs)}] [{frames.frame_idx}/{frames.frame_count}] {str(frames.video_path)}")
|
bar.set_description(f"[{frames.video_nr}/{len(frames.video_srcs)}] [{frames.frame_idx}/{frames.frame_count}] {str(frames.video_path)}")
|
||||||
|
|
||||||
|
|
||||||
|
def draw_track(img: cv2.Mat, track: Track, color_index: int):
|
||||||
|
history = track.history
|
||||||
|
cv2.putText(img, f"{track.track_id} ({len(history)})", to_point(history[0].get_foot_coords()), cv2.FONT_HERSHEY_DUPLEX, 1, color=bgr_colors[color_index % len(bgr_colors)])
|
||||||
|
|
||||||
|
point_color = detection_color(history[0], color_index)
|
||||||
|
cv2.circle(img, to_point(history[0].get_foot_coords()), 3, point_color, 2)
|
||||||
|
|
||||||
|
for j in range(len(history)-1):
|
||||||
|
a = history[j]
|
||||||
|
b = history[j+1]
|
||||||
|
# TODO)) replace with Track.get_with_interpolated_history()
|
||||||
|
gap = b.frame_nr - a.frame_nr - 1
|
||||||
|
if gap < 0:
|
||||||
|
print(f"WARNING, gap between frames {a.frame_nr} -> {b.frame_nr} is negative?")
|
||||||
|
if gap > 0:
|
||||||
|
for g in range(gap):
|
||||||
|
p1 = a.get_foot_coords()
|
||||||
|
p2 = b.get_foot_coords()
|
||||||
|
point = (lerp(p1[0], p2[0], g/gap), lerp(p1[1], p2[1], g/gap))
|
||||||
|
|
||||||
|
cv2.circle(img, to_point(point), 3, (0,0,255), 1)
|
||||||
|
|
||||||
|
color = detection_color(b, color_index, a)
|
||||||
|
cv2.line(img, to_point(a.get_foot_coords()), to_point(b.get_foot_coords()), color, 1)
|
||||||
|
point_color = detection_color(b, color_index)
|
||||||
|
cv2.circle(img, to_point(b.get_foot_coords()), 3, point_color, 2)
|
||||||
|
|
||||||
|
def blacklist_tracks():
|
||||||
|
config = parser.parse_args()
|
||||||
|
|
||||||
|
cv2.namedWindow("frame", cv2.WND_PROP_FULLSCREEN)
|
||||||
|
cv2.setWindowProperty("frame",cv2.WND_PROP_FULLSCREEN,cv2.WINDOW_FULLSCREEN)
|
||||||
|
|
||||||
|
backdrop = cv2.imread('../DATASETS/hof3/output.png')
|
||||||
|
blacklist = []
|
||||||
|
path: Path = config.save_for_training
|
||||||
|
blacklist_file = path / "blacklist.jsonl"
|
||||||
|
whitelist_file = path / "whitelist.jsonl" # for skipping
|
||||||
|
tracks_file = path / "tracks.json"
|
||||||
|
|
||||||
|
FPS = 12 # TODO)) From config
|
||||||
|
|
||||||
|
if whitelist_file.exists():
|
||||||
|
# with whitelist_file.open('r') as fp:
|
||||||
|
with jsonlines.open(whitelist_file, 'r') as reader:
|
||||||
|
whitelist = [l for l in reader.iter(type=str)]
|
||||||
|
else:
|
||||||
|
whitelist = []
|
||||||
|
|
||||||
|
try:
|
||||||
|
for track in read_tracks_json(tracks_file, blacklist_file, FPS):
|
||||||
|
if track.track_id in whitelist:
|
||||||
|
logger.info(f'skip whitelisted {track.track_id}')
|
||||||
|
continue
|
||||||
|
|
||||||
|
img = backdrop.copy()
|
||||||
|
draw_track(img, track, 0)
|
||||||
|
|
||||||
|
imgS = cv2.resize(img, (1920, 1080))
|
||||||
|
cv2.imshow('frame', imgS)
|
||||||
|
while True:
|
||||||
|
k = cv2.waitKey(0)
|
||||||
|
if k==27: # Esc key to stop
|
||||||
|
raise StopIteration
|
||||||
|
elif k == ord('s'):
|
||||||
|
break # skip for now
|
||||||
|
elif k == ord('y'):
|
||||||
|
with jsonlines.open(whitelist_file, mode='a') as writer:
|
||||||
|
# skip next time around
|
||||||
|
writer.write(track.track_id)
|
||||||
|
break
|
||||||
|
elif k == ord('n'):
|
||||||
|
print('blacklist', track.track_id)
|
||||||
|
logger.info(f"Append {len(blacklist)} items to {str(blacklist_file)}")
|
||||||
|
with jsonlines.open(blacklist_file, mode='a') as writer:
|
||||||
|
writer.write(track.track_id)
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
# ignore all other keypresses
|
||||||
|
print(k) # else print its value
|
||||||
|
continue
|
||||||
|
except StopIteration as e:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def rewrite_raw_track_files():
|
||||||
|
logging.basicConfig(level=logging.DEBUG)
|
||||||
|
config = parser.parse_args()
|
||||||
|
trap.tracker.rewrite_raw_track_files(config.save_for_training)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def interpolate_missing_frames(data: pd.DataFrame):
|
def interpolate_missing_frames(data: pd.DataFrame):
|
||||||
|
|
197
trap/tracker.py
197
trap/tracker.py
|
@ -4,11 +4,13 @@ import csv
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
from math import nan
|
||||||
from multiprocessing import Event
|
from multiprocessing import Event
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import pickle
|
import pickle
|
||||||
import time
|
import time
|
||||||
from typing import Optional, List
|
from typing import Optional, List
|
||||||
|
import jsonlines
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
import torchvision
|
import torchvision
|
||||||
|
@ -28,7 +30,7 @@ from bytetracker import BYTETracker
|
||||||
|
|
||||||
from tsmoothie.smoother import KalmanSmoother, ConvolutionSmoother
|
from tsmoothie.smoother import KalmanSmoother, ConvolutionSmoother
|
||||||
import tsmoothie.smoother
|
import tsmoothie.smoother
|
||||||
from datetime import datetime
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
# Detection = [int, int, int, int, float, int]
|
# Detection = [int, int, int, int, float, int]
|
||||||
# Detections = [Detection]
|
# Detections = [Detection]
|
||||||
|
@ -89,6 +91,43 @@ class Multifile():
|
||||||
def readline(self):
|
def readline(self):
|
||||||
return self.g.__next__()
|
return self.g.__next__()
|
||||||
|
|
||||||
|
FIELDNAMES = ['frame_id', 'track_id', 'l', 't', 'w', 'h', 'x', 'y', 'state']
|
||||||
|
|
||||||
|
def read_tracks_json(path: Path, blacklist_path: Path, fps):
|
||||||
|
"""
|
||||||
|
Reader for tracks.json produced by TrainingDataWriter
|
||||||
|
"""
|
||||||
|
with path.open('r') as fp:
|
||||||
|
tracks_dict: dict = json.load(fp)
|
||||||
|
|
||||||
|
if blacklist_path.exists():
|
||||||
|
with jsonlines.open(blacklist_path, 'r') as reader:
|
||||||
|
blacklist = [track_id for track_id in reader.iter(type=str)]
|
||||||
|
|
||||||
|
else:
|
||||||
|
blacklist = []
|
||||||
|
|
||||||
|
for track_id, detection_values in tracks_dict.items():
|
||||||
|
if track_id in blacklist:
|
||||||
|
continue
|
||||||
|
|
||||||
|
history = []
|
||||||
|
# for detection_values in
|
||||||
|
for detection_items in detection_values:
|
||||||
|
d = dict(zip(FIELDNAMES, detection_items))
|
||||||
|
history.append(Detection(
|
||||||
|
d['track_id'],
|
||||||
|
d['l'],
|
||||||
|
d['t'],
|
||||||
|
d['w'],
|
||||||
|
d['h'],
|
||||||
|
nan,
|
||||||
|
d['state'],
|
||||||
|
d['frame_id'],
|
||||||
|
1
|
||||||
|
))
|
||||||
|
|
||||||
|
yield Track(track_id, history, fps=fps)
|
||||||
|
|
||||||
class TrainingDataWriter:
|
class TrainingDataWriter:
|
||||||
def __init__(self, training_path: Optional[Path]):
|
def __init__(self, training_path: Optional[Path]):
|
||||||
|
@ -114,7 +153,7 @@ class TrainingDataWriter:
|
||||||
self.training_fp = open(self.path / f'all-{d}.txt', 'w')
|
self.training_fp = open(self.path / f'all-{d}.txt', 'w')
|
||||||
logger.debug(f"Writing tracker data to {self.training_fp.name}")
|
logger.debug(f"Writing tracker data to {self.training_fp.name}")
|
||||||
# following https://github.com/StanfordASL/Trajectron-plus-plus/blob/master/experiments/pedestrians/process_data.py
|
# following https://github.com/StanfordASL/Trajectron-plus-plus/blob/master/experiments/pedestrians/process_data.py
|
||||||
self.csv = csv.DictWriter(self.training_fp, fieldnames=['frame_id', 'track_id', 'l', 't', 'w', 'h', 'x', 'y', 'state'], delimiter='\t', quoting=csv.QUOTE_NONE)
|
self.csv = csv.DictWriter(self.training_fp, fieldnames=FIELDNAMES, delimiter='\t', quoting=csv.QUOTE_NONE)
|
||||||
self.count = 0
|
self.count = 0
|
||||||
return self
|
return self
|
||||||
|
|
||||||
|
@ -146,58 +185,80 @@ class TrainingDataWriter:
|
||||||
return
|
return
|
||||||
|
|
||||||
self.training_fp.close()
|
self.training_fp.close()
|
||||||
|
rewrite_raw_track_files(self.path)
|
||||||
source_files = list(self.path.glob("*.txt")) # we loop twice, so need a list instead of generator
|
|
||||||
total = 0
|
|
||||||
sources = Multifile(source_files)
|
|
||||||
for line in sources:
|
|
||||||
if len(line) > 3: # make sure not to count empty lines
|
|
||||||
total += 1
|
|
||||||
|
|
||||||
|
|
||||||
lines = {
|
|
||||||
'train': int(total * .8),
|
|
||||||
'val': int(total * .12),
|
|
||||||
'test': int(total * .08),
|
|
||||||
}
|
|
||||||
|
|
||||||
logger.info(f"Splitting gathered data from {sources.name}")
|
|
||||||
# for source_file in source_files:
|
|
||||||
|
|
||||||
tracks_file = self.path / 'tracks.json'
|
|
||||||
tracks = defaultdict(lambda: [])
|
|
||||||
|
|
||||||
for name, line_nrs in lines.items():
|
|
||||||
dir_path = self.path / name
|
|
||||||
dir_path.mkdir(exist_ok=True)
|
|
||||||
file = dir_path / 'tracked.txt'
|
|
||||||
logger.debug(f"- Write {line_nrs} lines to {file}")
|
|
||||||
with file.open('w') as target_fp:
|
|
||||||
max_track_id = 0
|
|
||||||
offset = 0
|
|
||||||
prev_file = None
|
|
||||||
for i in range(line_nrs):
|
|
||||||
line = sources.readline()
|
|
||||||
current_file = sources.current_file
|
|
||||||
if prev_file != current_file:
|
|
||||||
offset = max_track_id
|
|
||||||
|
|
||||||
logger.debug(f'{name} - update offset {offset} ({sources.current_file})')
|
|
||||||
prev_file = current_file
|
|
||||||
|
|
||||||
parts = line.split('\t')
|
def rewrite_raw_track_files(path: Path):
|
||||||
track_id = int(parts[1]) + offset
|
source_files = list(sorted(path.glob("*.txt"))) # we loop twice, so need a list instead of generator
|
||||||
|
total = 0
|
||||||
if track_id > max_track_id:
|
sources = Multifile(source_files)
|
||||||
max_track_id = track_id
|
for line in sources:
|
||||||
|
if len(line) > 3: # make sure not to count empty lines
|
||||||
parts[1] = str(track_id)
|
total += 1
|
||||||
target_fp.write("\t".join(parts))
|
|
||||||
tracks[track_id].append(parts)
|
|
||||||
|
destinations = {
|
||||||
with tracks_file.open('w') as fp:
|
'train': int(total * .8),
|
||||||
json.dump(tracks, fp)
|
'val': int(total * .12),
|
||||||
|
'test': int(total * .08),
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.info(f"Splitting gathered data from {source_files}")
|
||||||
|
# for source_file in source_files:
|
||||||
|
|
||||||
|
tracks_file = path / 'tracks.json'
|
||||||
|
tracks = defaultdict(lambda: [])
|
||||||
|
|
||||||
|
|
||||||
|
offset = 0
|
||||||
|
max_track_id = 0
|
||||||
|
prev_file = None
|
||||||
|
|
||||||
|
# all-2024-11-12T13:30.txt
|
||||||
|
file_date = None
|
||||||
|
|
||||||
|
for name, line_nrs in destinations.items():
|
||||||
|
dir_path = path / name
|
||||||
|
dir_path.mkdir(exist_ok=True)
|
||||||
|
file = dir_path / 'tracked.txt'
|
||||||
|
logger.debug(f"- Write {line_nrs} lines to {file}")
|
||||||
|
with file.open('w') as target_fp:
|
||||||
|
|
||||||
|
for i in range(line_nrs):
|
||||||
|
line = sources.readline()
|
||||||
|
current_file = sources.current_file
|
||||||
|
if prev_file != current_file:
|
||||||
|
offset: int = max_track_id
|
||||||
|
|
||||||
|
logger.info(f'{name} - update offset {offset} ({sources.current_file})')
|
||||||
|
prev_file = current_file
|
||||||
|
|
||||||
|
file_date = datetime.strptime(current_file.name, 'all-%Y-%m-%dT%H:%M.txt')
|
||||||
|
|
||||||
|
if file_date:
|
||||||
|
frame_date = file_date + timedelta(seconds = int(parts[0])//10)
|
||||||
|
else:
|
||||||
|
frame_date = None
|
||||||
|
|
||||||
|
parts = line.split('\t')
|
||||||
|
track_id = int(parts[1]) + offset
|
||||||
|
|
||||||
|
if track_id > max_track_id:
|
||||||
|
max_track_id = track_id
|
||||||
|
|
||||||
|
parts[1] = str(track_id)
|
||||||
|
target_fp.write("\t".join(parts))
|
||||||
|
|
||||||
|
parts = [float(p) for p in parts]
|
||||||
|
tracks[track_id].append([
|
||||||
|
int(parts[0] / 10),
|
||||||
|
track_id,
|
||||||
|
] + parts[2:8] + [int(parts[8])])
|
||||||
|
|
||||||
|
with tracks_file.open('w') as fp:
|
||||||
|
logger.info(f"Write {len(tracks)} tracks to {str(tracks_file)}")
|
||||||
|
json.dump(tracks, fp)
|
||||||
|
|
||||||
|
|
||||||
class TrackerWrapper():
|
class TrackerWrapper():
|
||||||
|
@ -317,6 +378,7 @@ class Tracker:
|
||||||
for detection in detections:
|
for detection in detections:
|
||||||
track = self.tracks[detection.track_id]
|
track = self.tracks[detection.track_id]
|
||||||
track.track_id = detection.track_id # for new tracks
|
track.track_id = detection.track_id # for new tracks
|
||||||
|
track.fps = self.config.camera.fps # for new tracks
|
||||||
|
|
||||||
track.history.append(detection) # add to history
|
track.history.append(detection) # add to history
|
||||||
|
|
||||||
|
@ -441,6 +503,7 @@ class Tracker:
|
||||||
if self.config.smooth_tracks:
|
if self.config.smooth_tracks:
|
||||||
frame = self.smoother.smooth_frame_tracks(frame)
|
frame = self.smoother.smooth_frame_tracks(frame)
|
||||||
|
|
||||||
|
# print(f"send to {self.trajectory_socket}, {self.config.zmq_trajectory_addr}")
|
||||||
self.trajectory_socket.send_pyobj(frame)
|
self.trajectory_socket.send_pyobj(frame)
|
||||||
|
|
||||||
end_time = time.time()
|
end_time = time.time()
|
||||||
|
@ -549,25 +612,27 @@ class Smoother:
|
||||||
def smooth(self, points: List[float]):
|
def smooth(self, points: List[float]):
|
||||||
self.smoother.smooth(points)
|
self.smoother.smooth(points)
|
||||||
return self.smoother.smooth_data[0]
|
return self.smoother.smooth_data[0]
|
||||||
|
|
||||||
|
def smooth_track(self, track: Track) -> Track:
|
||||||
|
ls = [d.l for d in track.history]
|
||||||
|
ts = [d.t for d in track.history]
|
||||||
|
ws = [d.w for d in track.history]
|
||||||
|
hs = [d.h for d in track.history]
|
||||||
|
self.smoother.smooth(ls)
|
||||||
|
ls = self.smoother.smooth_data[0]
|
||||||
|
self.smoother.smooth(ts)
|
||||||
|
ts = self.smoother.smooth_data[0]
|
||||||
|
self.smoother.smooth(ws)
|
||||||
|
ws = self.smoother.smooth_data[0]
|
||||||
|
self.smoother.smooth(hs)
|
||||||
|
hs = self.smoother.smooth_data[0]
|
||||||
|
new_history = [Detection(d.track_id, l, t, w, h, d.conf, d.state, d.frame_nr, d.det_class) for l, t, w, h, d in zip(ls,ts,ws,hs, track.history)]
|
||||||
|
new_track = Track(track.track_id, new_history, track.predictor_history, track.predictions, track.fps)
|
||||||
|
|
||||||
def smooth_frame_tracks(self, frame: Frame) -> Frame:
|
def smooth_frame_tracks(self, frame: Frame) -> Frame:
|
||||||
new_tracks = []
|
new_tracks = []
|
||||||
for track in frame.tracks.values():
|
for track in frame.tracks.values():
|
||||||
ls = [d.l for d in track.history]
|
new_track = self.smooth_track(track)
|
||||||
ts = [d.t for d in track.history]
|
|
||||||
ws = [d.w for d in track.history]
|
|
||||||
hs = [d.h for d in track.history]
|
|
||||||
self.smoother.smooth(ls)
|
|
||||||
ls = self.smoother.smooth_data[0]
|
|
||||||
self.smoother.smooth(ts)
|
|
||||||
ts = self.smoother.smooth_data[0]
|
|
||||||
self.smoother.smooth(ws)
|
|
||||||
ws = self.smoother.smooth_data[0]
|
|
||||||
self.smoother.smooth(hs)
|
|
||||||
hs = self.smoother.smooth_data[0]
|
|
||||||
new_history = [Detection(d.track_id, l, t, w, h, d.conf, d.state, d.frame_nr, d.det_class) for l, t, w, h, d in zip(ls,ts,ws,hs, track.history)]
|
|
||||||
new_track = Track(track.track_id, new_history, track.predictor_history, track.predictions)
|
|
||||||
new_tracks.append(new_track)
|
new_tracks.append(new_track)
|
||||||
frame.tracks = {t.track_id: t for t in new_tracks}
|
frame.tracks = {t.track_id: t for t in new_tracks}
|
||||||
return frame
|
return frame
|
||||||
|
|
8
trap/utils.py
Normal file
8
trap/utils.py
Normal file
|
@ -0,0 +1,8 @@
|
||||||
|
def lerp(a: float, b: float, t: float) -> float:
|
||||||
|
"""Linear interpolate on the scale given by a to b, using t as the point on that scale.
|
||||||
|
Examples
|
||||||
|
--------
|
||||||
|
50 == lerp(0, 100, 0.5)
|
||||||
|
4.2 == lerp(1, 5, 0.8)
|
||||||
|
"""
|
||||||
|
return (1 - t) * a + t * b
|
Loading…
Reference in a new issue