Skip to content

Commit d8e4424

Browse files
author
spencer@primus
committed
Add track visualization and add notebook formatting
1 parent 23c5ecd commit d8e4424

4 files changed

Lines changed: 1402 additions & 1031 deletions

File tree

‎Makefile‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
NAME := avapi
22
INSTALL_STAMP := .install.stamp
33
POETRY := $(shell command -v poetry 2> /dev/null)
4+
PYFOLDERS := avapi tests
45
.DEFAULT_GOAL := help
56

67
.PHONY: help
@@ -18,7 +19,7 @@ help:
1819
install: $(INSTALL_STAMP)
1920
$(INSTALL_STAMP): pyproject.toml poetry.lock
2021
@if [ -z $(POETRY) ]; then echo "Poetry could not be found. See https://python-poetry.org/docs/"; exit 2; fi
21-
$(POETRY) install
22+
$(POETRY) install --all-extras
2223
touch $(INSTALL_STAMP)
2324

2425
.PHONY: clean
@@ -28,18 +29,18 @@ clean:
2829

2930
.PHONY: lint
3031
lint: $(INSTALL_STAMP)
31-
$(POETRY) run isort --profile=black --lines-after-imports=2 --check-only ./tests/ $(NAME)
32-
$(POETRY) run black --check ./tests/ $(NAME) --diff
33-
$(POETRY) run flake8 --ignore=W503,E501 ./tests/ $(NAME)
34-
$(POETRY) run mypy ./tests/ $(NAME) --ignore-missing-imports
32+
$(POETRY) run isort --profile=black --lines-after-imports=2 --check-only $(PYFOLDERS)
33+
$(POETRY) run black --check $(PYFOLDERS) --diff
34+
$(POETRY) run flake8 --ignore=W503,E501 $(PYFOLDERS)
35+
$(POETRY) run mypy $(PYFOLDERS) --ignore-missing-imports
3536
$(POETRY) run bandit -r $(NAME) -s B608
3637

3738
.PHONY: format
3839
format: $(INSTALL_STAMP)
39-
$(POETRY) run autoflake --remove-all-unused-imports -i -r ./tests/ $(NAME) --exclude=__init__.py
40-
$(POETRY) run isort --profile=black --lines-after-imports=2 ./tests/ $(NAME)
41-
$(POETRY) run black ./tests/ $(NAME)
40+
$(POETRY) run autoflake --remove-all-unused-imports -i -r $(PYFOLDERS) --exclude=__init__.py
41+
$(POETRY) run isort --profile=black --lines-after-imports=2 $(PYFOLDERS)
42+
$(POETRY) run black $(PYFOLDERS)
4243

4344
.PHONY: test
4445
test: $(INSTALL_STAMP)
45-
$(POETRY) run pytest ./tests/ --cov-report term-missing --cov-fail-under 0 --cov $(NAME)
46+
$(POETRY) run pytest ./tests/ --cov-report term-missing --cov-fail-under 0 --cov $(NAME)

‎avapi/visualize/tracking.py‎

Lines changed: 163 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,147 @@
11
import numpy as np
2+
from avstack.geometry import (
3+
Attitude,
4+
Box3D,
5+
GlobalOrigin3D,
6+
Position,
7+
transform_orientation,
8+
)
29
from PIL import ImageDraw
310

411

5-
def draw_detections(image, detections, show_class=False, show_score=False):
12+
def add_metadata_to_image(
13+
image, x0, y1, class_, score, show_class=False, show_score=False
14+
):
15+
draw = ImageDraw.Draw(image)
16+
if show_class and show_score:
17+
draw.text((x0, y1 + 2), "{}:{}".format(class_, score), fill=(0, 255, 0))
18+
elif show_class:
19+
draw.text((x0, y1 + 2), "{}".format(class_), fill=(0, 255, 0))
20+
elif show_score:
21+
draw.text((x0, y1 + 2), "{}".format(score), fill=(0, 255, 0))
22+
return image
23+
24+
25+
def draw_box_detections_3d(
26+
image,
27+
detections,
28+
extent=[(-100, 100), (-100, 100)],
29+
img_size=(600, 600),
30+
show_class=False,
31+
show_score=False,
32+
):
33+
# get pixel scaling
34+
pix_per_m = (
35+
img_size[0] / (extent[0][1] - extent[0][0]),
36+
img_size[1] / (extent[1][1] - extent[1][0]),
37+
)
38+
39+
draw = ImageDraw.Draw(image)
40+
for detection in detections:
41+
bev_corners = detection.box.corners[:4, :2]
42+
vertices = bev_corners * pix_per_m
43+
vertices[:, 0] += img_size[0] / 2
44+
vertices[:, 1] += img_size[1] / 2
45+
vertices = [(x, y) for x, y in vertices]
46+
draw.polygon(vertices, outline=(0, 255, 0), width=4)
47+
48+
class_ = detection.obj_type
49+
score = detection.score if detection.score else "N/A"
50+
image = add_metadata_to_image(
51+
image, vertices[0][0], vertices[0][1], class_, score, show_class, show_score
52+
)
53+
54+
del draw
55+
return image
56+
57+
58+
def draw_stonesoup_box_tracks_3d(
59+
image,
60+
tracks,
61+
extent=[(-100, 100), (-100, 100)],
62+
img_size=(600, 600),
63+
show_history=True,
64+
show_class=True,
65+
show_score=True,
66+
):
67+
"""Draw tracks on an image
68+
69+
Parameters
70+
----------
71+
image: :class:`PIL.Image`
72+
Image on which to draw the tracks
73+
detections: : set of :class:`~.Tracks`
74+
A set of tracks generated by our :class:`~.MultiTargetTracker`
75+
show_history: bool
76+
Whether to draw the trajectory of the track. Default is ``True``
77+
show_class: bool
78+
Whether to draw the class of the object. Default is ``True``
79+
show_score: bool
80+
Whether to draw the score of the object. Default is ``True``
81+
82+
Returns
83+
-------
84+
: :class:`PIL.Image`
85+
Image with tracks drawn
86+
87+
"""
88+
# get pixel scaling
89+
pix_per_m = (
90+
img_size[0] / (extent[0][1] - extent[0][0]),
91+
img_size[1] / (extent[1][1] - extent[1][0]),
92+
)
93+
94+
# np.array(track.states[0].state_vector[[0, 2, 4]])
95+
96+
draw = ImageDraw.Draw(image)
97+
for track in tracks:
98+
# import pdb; pdb.set_trace()
99+
bboxes = [
100+
Box3D(
101+
position=Position(
102+
np.array(state.state_vector[[0, 2, 4]]).reshape(3),
103+
GlobalOrigin3D,
104+
),
105+
attitude=Attitude(
106+
transform_orientation(
107+
[0, 0, state.state_vector[9]], "euler", "quat"
108+
),
109+
GlobalOrigin3D,
110+
),
111+
hwl=np.array(state.state_vector[[6, 7, 8]]).reshape(3),
112+
)
113+
for state in track.states
114+
]
115+
116+
# draw current box
117+
bev_corners = bboxes[-1].corners[:4, :2]
118+
vertices = bev_corners * pix_per_m
119+
vertices[:, 0] += img_size[0] / 2
120+
vertices[:, 1] += img_size[1] / 2
121+
vertices = [(x, y) for x, y in vertices]
122+
draw.polygon(vertices, outline=(255, 0, 0), width=2)
123+
124+
if show_history:
125+
pts = [
126+
(
127+
img_size[0] / 2 + box.t[0] * pix_per_m[0],
128+
img_size[1] / 2 + box.t[1] * pix_per_m[1],
129+
)
130+
for box in bboxes
131+
]
132+
draw.line(pts, fill=(255, 0, 0), width=2)
133+
134+
class_ = track.metadata["class"]["name"]
135+
score = round(
136+
float(track.metadata["score"]) if track.metadata["score"] else 0.0, 2
137+
)
138+
image = add_metadata_to_image(
139+
image, vertices[0][0], vertices[0][1], class_, score, show_class, show_score
140+
)
141+
return image
142+
143+
144+
def draw_box_detections_2d(image, detections, show_class=False, show_score=False):
6145
"""Draw detections on an image
7146
8147
Parameters
@@ -26,23 +165,26 @@ def draw_detections(image, detections, show_class=False, show_score=False):
26165
x0, y0, x1, y1 = detection.box.box2d
27166
# x0, y0, w, h = np.array(detection.state_vector).reshape(4)
28167
# x1, y1 = (x0 + w, y0 + h)
29-
draw.rectangle([x0, y0, x1, y1], outline=(0, 255, 0), width=1)
168+
# draw.rectangle([x0, y0, x1, y1], outline=(0, 255, 0), width=1)
169+
vertices = [
170+
(x0, y0),
171+
(x0, y1),
172+
(x1, y1),
173+
(x1, y0),
174+
]
175+
draw.polygon(vertices, outline=(0, 255, 0), width=1)
176+
30177
class_ = detection.obj_type
31178
score = detection.score if detection.score else "N/A"
32-
# class_ = detection.metadata['class']['name']
33-
# score = round(float(detection.metadata['score']),2)
34-
if show_class and show_score:
35-
draw.text((x0, y1 + 2), "{}:{}".format(class_, score), fill=(0, 255, 0))
36-
elif show_class:
37-
draw.text((x0, y1 + 2), "{}".format(class_), fill=(0, 255, 0))
38-
elif show_score:
39-
draw.text((x0, y1 + 2), "{}".format(score), fill=(0, 255, 0))
179+
image = add_metadata_to_image(
180+
image, x0, y1, class_, score, show_class, show_score
181+
)
40182

41183
del draw
42184
return image
43185

44186

45-
def draw_stonesoup_tracks(
187+
def draw_stonesoup_box_tracks_2d(
46188
image, tracks, show_history=True, show_class=True, show_score=True
47189
):
48190
"""Draw tracks on an image
@@ -77,7 +219,13 @@ def draw_stonesoup_tracks(
77219
x0, y0, w, h = bboxes[-1]
78220
x1 = x0 + w
79221
y1 = y0 + h
80-
draw.rectangle([x0, y0, x1, y1], outline=(255, 0, 0), width=2)
222+
vertices = [
223+
(x0, y0),
224+
(x0, y1),
225+
(x1, y1),
226+
(x1, y0),
227+
]
228+
draw.polygon(vertices, outline=(255, 0, 0), width=2)
81229

82230
if show_history:
83231
pts = [(box[0] + box[2] / 2, box[1] + box[3] / 2) for box in bboxes]
@@ -87,10 +235,7 @@ def draw_stonesoup_tracks(
87235
score = round(
88236
float(track.metadata["score"]) if track.metadata["score"] else 0.0, 2
89237
)
90-
if show_class and show_score:
91-
draw.text((x0, y1 + 2), "{}:{}".format(class_, score), fill=(255, 0, 0))
92-
elif show_class:
93-
draw.text((x0, y1 + 2), "{}".format(class_), fill=(255, 0, 0))
94-
elif show_score:
95-
draw.text((x0, y1 + 2), "{}".format(score), fill=(255, 0, 0))
238+
image = add_metadata_to_image(
239+
image, x0, y1, class_, score, show_class, show_score
240+
)
96241
return image

0 commit comments

Comments
 (0)