File size: 12,239 Bytes
663494c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 |
"""
Written by Jinhyung Park
Simple BEV visualization for 3D points & boxes.
"""
import cv2
import matplotlib
import numpy as np
# np.set_printoptions(precision=3, suppress=True)
class Canvas_BEV(object):
def __init__(
self,
canvas_shape=(2000, 2000),
canvas_x_range=(-50.0, 50.0),
canvas_y_range=(-50.0, 50.0),
canvas_bg_color=(0, 0, 0),
):
"""
Args:
canvas_shape (Tuple[int]): Shape of BEV Canvas image. First element
corresponds to X range, the second element to Y range.
canvas_x_range (Tuple[int]): Range of X-coords to visualize. X is
vertical: negative ~ positive is top ~ down.
canvas_y_range (Tuple[int]): Range of Y-coords to visualize. Y is
horizontal: negative ~ positive is left ~ right.
canvas_bg_color (Tuple[int]): RGB (0 ~ 255) of Canvas background
color.
"""
# Sanity check ratios
if (canvas_shape[0] / canvas_shape[1]) != (
(canvas_x_range[0] - canvas_x_range[1])
/ (canvas_y_range[0] - canvas_y_range[1])
):
print(
"Not an error, but the x & y ranges are not "
"proportional to canvas height & width."
)
self.canvas_shape = canvas_shape
self.canvas_x_range = canvas_x_range
self.canvas_y_range = canvas_y_range
self.canvas_bg_color = canvas_bg_color
self.clear_canvas()
def get_canvas(self):
return self.canvas
def clear_canvas(self):
self.canvas = np.zeros((*self.canvas_shape, 3), dtype=np.uint8)
self.canvas[..., :] = self.canvas_bg_color
def get_canvas_coords(self, xy):
"""
Args:
xy (ndarray): (N, 2+) array of coordinates. Additional columns
beyond the first two are ignored.
Returns:
canvas_xy (ndarray): (N, 2) array of xy scaled into canvas
coordinates. Invalid locations of canvas_xy are clipped into
range. "x" is dim0, "y" is dim1 of canvas.
valid_mask (ndarray): (N,) boolean mask indicating which of
canvas_xy fits into canvas.
"""
xy = np.copy(xy) # prevent in-place modifications
# np.set_printoptions(precision=3, suppress=True)
# print(xy.shape)
# print(xy[0:100, 0])
# print(xy[0:100, 1])
# print(xy[0:100, 2])
# # print(xy[40:44])
x = xy[:, 0]
y = xy[:, 1]
# Get valid mask
valid_mask = (
(x > self.canvas_x_range[0])
& (x < self.canvas_x_range[1])
& (y > self.canvas_y_range[0])
& (y < self.canvas_y_range[1])
)
# print(np.max(x))
# print(np.min(x))
# print(np.max(y))
# print(np.min(y))
# print(self.canvas_x_range)
# print(self.canvas_y_range)
# print(self.canvas_shape)
# zxc
# Rescale points
x = (x - self.canvas_x_range[0]) / (
self.canvas_x_range[1] - self.canvas_x_range[0]
)
# print(x[40:44])
x = x * self.canvas_shape[0]
# print(x[40:44])
x = np.clip(np.around(x), 0, self.canvas_shape[0] - 1).astype(np.int32)
# print(x[40:44])
y = (y - self.canvas_y_range[0]) / (
self.canvas_y_range[1] - self.canvas_y_range[0]
)
y = y * self.canvas_shape[1]
y = np.clip(np.around(y), 0, self.canvas_shape[1] - 1).astype(np.int32)
# Return
canvas_xy = np.stack([x, y], axis=1)
# print(canvas_xy.shape)
# print(canvas_xy[40:44])
# # zxc
return canvas_xy, valid_mask
def draw_canvas_points(
self, canvas_xy, radius=-1, colors=None, colors_operand=None
):
"""
Draws canvas_xy onto self.canvas.
Args:
canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates.
"x" is dim0, "y" is dim1 of canvas.
radius (Int):
-1: Each point is visualized as a single pixel.
r: Each point is visualized as a circle with radius r.
colors:
None: colors all points white.
Tuple: RGB (0 ~ 255), indicating a single color for all points.
ndarray: (N, 3) array of RGB values for each point.
String: Such as "Spectral", uses a matplotlib cmap, with the
operand (the value cmap is called on for each point) being
colors_operand. If colors_operand is None, uses normalized
distance from (0, 0) of XY point coords.
colors_operand (ndarray | None): (N,) array of values cooresponding
to canvas_xy, to be used only if colors is a cmap.
"""
if len(canvas_xy) == 0:
return
if colors is None:
colors = np.full((len(canvas_xy), 3), fill_value=255, dtype=np.uint8)
elif isinstance(colors, tuple):
assert len(colors) == 3
colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8)
colors_tmp[..., :] = np.array(colors)
colors = colors_tmp
elif isinstance(colors, np.ndarray):
assert len(colors) == len(canvas_xy)
colors = colors.astype(np.uint8)
elif isinstance(colors, str):
colors = matplotlib.cm.get_cmap(colors)
if colors_operand is None:
# Get distances from (0, 0) (albeit potentially clipped)
origin_center = self.get_canvas_coords(np.zeros((1, 2)))[0][0]
colors_operand = np.sqrt(((canvas_xy - origin_center) ** 2).sum(axis=1))
# Normalize 0 ~ 1 for cmap
colors_operand = colors_operand - colors_operand.min()
colors_operand = colors_operand / colors_operand.max()
# Get cmap colors - note that cmap returns (*input_shape, 4), with
# colors scaled 0 ~ 1
colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8)
else:
raise Exception(
"colors type {} was not an expected type".format(type(colors))
)
# direct draw on the canvas, x-> height, y->width
if radius == -1:
# self.canvas[canvas_xy[:, 0], canvas_xy[:, 1], :] = colors
self.canvas[canvas_xy[:, 1], canvas_xy[:, 0], :] = colors
# draw with cv2 requires x->horizontal (width), y-> vertical height
# change from xy coordinate to yx coordinate
else:
for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()):
# self.canvas = cv2.circle(
# self.canvas, (y, x), radius, color, -1, lineType=cv2.LINE_AA
# )
self.canvas = cv2.circle(
self.canvas, (x, y), radius, color, -1, lineType=cv2.LINE_AA
)
def draw_boxes(
self,
boxes=None,
corners=None,
colors=None,
texts=None,
box_line_thickness=2,
box_text_size=0.5,
text_corner=0,
):
"""
Draws a set of boxes onto the canvas.
Args:
boxes (ndarray): Can either be of shape:
(N, 7): Then, assumes (x, y, z, x_size, y_size, z_size, yaw)
(N, 5): Then, assumes (x, y, x_size, y_size, yaw)
Everything is in the same coordinate system as points
(not canvas coordinates)
colors:
None: colors all points white.
Tuple: RGB (0 ~ 255), indicating a single color for all points.
ndarray: (N, 3) array of RGB values for each point.
texts (List[String]): Length N; text to write next to boxes.
box_line_thickness (int): cv2 line/text thickness
box_text_size (float): cv2 putText size
text_corner (int): 0 ~ 3. Which corner of 3D box to write text at.
"""
num_boxes = len(boxes) if boxes is not None else len(corners)
# Setup colors
if colors is None:
colors = np.full((num_boxes, 3), fill_value=255, dtype=np.uint8)
elif isinstance(colors, tuple):
assert len(colors) == 3
colors_tmp = np.zeros((num_boxes, 3), dtype=np.uint8)
colors_tmp[..., : len(colors)] = np.array(colors)
colors = colors_tmp
elif isinstance(colors, np.ndarray):
assert len(colors) == num_boxes
colors = colors.astype(np.uint8)
else:
raise Exception(
"colors type {} was not an expected type".format(type(colors))
)
if boxes is not None:
boxes = np.copy(boxes) # prevent in-place modifications
assert len(boxes.shape) == 2
if boxes.shape[-1] == 7:
boxes = boxes[:, [0, 1, 3, 4, 6]]
else:
assert boxes.shape[-1] == 5
## Get the BEV four corners
# Get BEV 4 corners in box canonical coordinates
bev_corners = (
np.array([[[0.5, 0.5], [-0.5, 0.5], [-0.5, -0.5], [0.5, -0.5]]])
* boxes[:, None, [2, 3]]
) # N x 4 x 2
# print(bev_corners[10])
# Get rotation matrix from yaw
rot_sin = np.sin(boxes[:, -1])
rot_cos = np.cos(boxes[:, -1])
rot_matrix = np.stack(
[[rot_cos, -rot_sin], [rot_sin, rot_cos]]
) # 2 x 2 x N
# Rotate BEV 4 corners. Result: N x 4 x 2
bev_corners = np.einsum("aij,jka->aik", bev_corners, rot_matrix)
# print(bev_corners[10])
# Translate BEV 4 corners
bev_corners = bev_corners + boxes[:, None, [0, 1]]
elif corners is not None:
bev_corners = corners
# print(bev_corners.shape)
# print(bev_corners[10:13])
## Transform BEV 4 corners to canvas coords
bev_corners_canvas, valid_mask = self.get_canvas_coords(
bev_corners.reshape(-1, 2)
)
# print(bev_corners_canvas[40:52])
bev_corners_canvas = bev_corners_canvas.reshape(*bev_corners.shape)
valid_mask = valid_mask.reshape(*bev_corners.shape[:-1])
# print(bev_corners_canvas[10:13])
# At least 1 corner in canvas to draw.
valid_mask = valid_mask.sum(axis=1) > 0
bev_corners_canvas = bev_corners_canvas[valid_mask]
if texts is not None:
texts = np.array(texts)[valid_mask]
if colors is not None:
colors = np.array(colors)[valid_mask]
## Draw onto canvas
# Draw the outer boundaries
idx_draw_pairs = [(0, 1), (1, 2), (2, 3), (3, 0)]
for i, (color, curr_box_corners) in enumerate(
zip(colors.tolist(), bev_corners_canvas)
):
curr_box_corners = curr_box_corners.astype(np.int32)
# if i >= 10 and i < 13:
# print(curr_box_corners[:, ::-1])
# change xy to yx coordinate as the cv2 requires
for start, end in idx_draw_pairs:
self.canvas = cv2.line(
self.canvas,
tuple(curr_box_corners[start].tolist()),
tuple(curr_box_corners[end].tolist()),
# tuple(curr_box_corners[start][::-1].tolist()),
# tuple(curr_box_corners[end][::-1].tolist()),
color=color,
thickness=box_line_thickness,
)
if texts is not None:
self.canvas = cv2.putText(
self.canvas,
str(texts[i]),
# tuple(curr_box_corners[text_corner][::-1].tolist()),
tuple(curr_box_corners[text_corner].tolist()),
cv2.FONT_HERSHEY_SIMPLEX,
box_text_size,
color=color,
thickness=box_line_thickness,
)
|