def create_crosswalks_model(map_api: MapAPI) -> CrosswalkModel:
vertices = []
vertexType = []
for element in map_api:
element_id = MapAPI.id_as_str(element.id)
if map_api.is_crosswalk(element):
crosswalk = map_api.get_crosswalk_coords(element_id)
coords = earclip(crosswalk["xyz"])
for triangle in coords:
for vertex in triangle:
vertices.append(vertex[0])
vertices.append(vertex[1])
# z-axis (0 in 2d)
vertices.append(0.001)
# White
vertexType.append(1)
vertices = np.array(vertices, dtype=np.float32)
vertexType = np.array(vertexType, dtype=np.float32)
return CrosswalkModel(vertices, len(vertices) // 3, colors=vertexType)
class TLSemanticRasterizer(Rasterizer):
"""
Rasteriser for the vectorised semantic map (generally loaded from json files).
"""
@staticmethod
def from_cfg(cfg, data_manager):
raster_cfg = cfg["raster_params"]
# map_type = raster_cfg["map_type"]
dataset_meta_key = raster_cfg["dataset_meta_key"]
render_context = RenderContext(
raster_size_px=np.array(raster_cfg["raster_size"]),
pixel_size_m=np.array(raster_cfg["pixel_size"]),
center_in_raster_ratio=np.array(raster_cfg["ego_center"]),
)
# filter_agents_threshold = raster_cfg["filter_agents_threshold"]
# history_num_frames = cfg["model_params"]["history_num_frames"]
semantic_map_filepath = data_manager.require(
raster_cfg["semantic_map_key"])
try:
dataset_meta = _load_metadata(dataset_meta_key, data_manager)
world_to_ecef = np.array(dataset_meta["world_to_ecef"],
dtype=np.float64)
except (KeyError, FileNotFoundError
): # TODO remove when new dataset version is available
world_to_ecef = get_hardcoded_world_to_ecef()
return TLSemanticRasterizer(render_context, semantic_map_filepath,
world_to_ecef)
def __init__(
self,
render_context: RenderContext,
semantic_map_path: str,
world_to_ecef: np.ndarray,
):
super(TLSemanticRasterizer, self).__init__()
self.render_context = render_context
self.raster_size = render_context.raster_size_px
self.pixel_size = render_context.pixel_size_m
self.ego_center = render_context.center_in_raster_ratio
self.world_to_ecef = world_to_ecef
self.proto_API = MapAPI(semantic_map_path, world_to_ecef)
self.bounds_info = self.get_bounds()
self.raster_channels = 3
# TODO is this the right place for this function?
def get_bounds(self) -> dict:
"""
For each elements of interest returns bounds [[min_x, min_y],[max_x, max_y]] and proto ids
Coords are computed by the MapAPI and, as such, are in the world ref system.
Returns:
dict: keys are classes of elements, values are dict with `bounds` and `ids` keys
"""
lanes_ids = []
crosswalks_ids = []
lanes_bounds = np.empty(
(0, 2, 2), dtype=np.float) # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]
crosswalks_bounds = np.empty(
(0, 2, 2), dtype=np.float) # [(X_MIN, Y_MIN), (X_MAX, Y_MAX)]
for element in self.proto_API:
element_id = MapAPI.id_as_str(element.id)
if self.proto_API.is_lane(element):
lane = self.proto_API.get_lane_coords(element_id)
x_min = min(np.min(lane["xyz_left"][:, 0]),
np.min(lane["xyz_right"][:, 0]))
y_min = min(np.min(lane["xyz_left"][:, 1]),
np.min(lane["xyz_right"][:, 1]))
x_max = max(np.max(lane["xyz_left"][:, 0]),
np.max(lane["xyz_right"][:, 0]))
y_max = max(np.max(lane["xyz_left"][:, 1]),
np.max(lane["xyz_right"][:, 1]))
lanes_bounds = np.append(lanes_bounds,
np.asarray([[[x_min, y_min],
[x_max, y_max]]]),
axis=0)
lanes_ids.append(element_id)
if self.proto_API.is_crosswalk(element):
crosswalk = self.proto_API.get_crosswalk_coords(element_id)
x_min = np.min(crosswalk["xyz"][:, 0])
y_min = np.min(crosswalk["xyz"][:, 1])
x_max = np.max(crosswalk["xyz"][:, 0])
y_max = np.max(crosswalk["xyz"][:, 1])
crosswalks_bounds = np.append(
crosswalks_bounds,
np.asarray([[[x_min, y_min], [x_max, y_max]]]),
axis=0,
)
crosswalks_ids.append(element_id)
return {
"lanes": {
"bounds": lanes_bounds,
"ids": lanes_ids
},
"crosswalks": {
"bounds": crosswalks_bounds,
"ids": crosswalks_ids
},
}
def rasterize(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
) -> np.ndarray:
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(history_frames[0]["ego_rotation"])
else:
ego_translation_m = np.append(
agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
raster_from_world = self.render_context.raster_from_world(
ego_translation_m, ego_yaw_rad)
world_from_raster = np.linalg.inv(raster_from_world)
# get XY of center pixel in world coordinates
center_in_raster_px = np.asarray(self.raster_size) * (0.5, 0.5)
center_in_world_m = transform_point(center_in_raster_px,
world_from_raster)
sem_im = self.render_semantic_map(center_in_world_m, raster_from_world,
history_tl_faces[0])
return sem_im.astype(np.float32) / 255
def render_semantic_map(self, center_in_world: np.ndarray,
raster_from_world: np.ndarray,
tl_faces: np.ndarray) -> np.ndarray:
"""Renders the semantic map at given x,y coordinates.
Args:
center_in_world (np.ndarray): XY of the image center in world ref system
raster_from_world (np.ndarray):
Returns:
np.ndarray: RGB raster
"""
img = 255 * np.ones(
shape=(self.raster_size[1], self.raster_size[0], 3),
dtype=np.uint8)
# filter using half a radius from the center
raster_radius = float(
np.linalg.norm(self.raster_size * self.pixel_size)) / 2
# get active traffic light faces
active_tl_ids = set(
filter_tl_faces_by_status(tl_faces, "ACTIVE")["face_id"].tolist())
# plot lanes
lanes_lines = defaultdict(list)
for idx in elements_within_bounds(center_in_world,
self.bounds_info["lanes"]["bounds"],
raster_radius):
lane = self.proto_API[self.bounds_info["lanes"]["ids"]
[idx]].element.lane
# get image coords
lane_coords = self.proto_API.get_lane_coords(
self.bounds_info["lanes"]["ids"][idx])
xy_left = cv2_subpixel(
transform_points(lane_coords["xyz_left"][:, :2],
raster_from_world))
xy_right = cv2_subpixel(
transform_points(lane_coords["xyz_right"][:, :2],
raster_from_world))
lanes_area = np.vstack(
(xy_left, np.flip(xy_right,
0))) # start->end left then end->start right
# Note(lberg): this called on all polygons skips some of them, don't know why
cv2.fillPoly(img, [lanes_area], (17, 17, 31),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
lane_type = "default" # no traffic light face is controlling this lane
lane_tl_ids = set(
[MapAPI.id_as_str(la_tc) for la_tc in lane.traffic_controls])
for tl_id in lane_tl_ids.intersection(active_tl_ids):
if self.proto_API.is_traffic_face_colour(tl_id, "red"):
lane_type = "red"
elif self.proto_API.is_traffic_face_colour(tl_id, "green"):
lane_type = "green"
elif self.proto_API.is_traffic_face_colour(tl_id, "yellow"):
lane_type = "yellow"
lanes_lines[lane_type].extend([xy_left, xy_right])
# For debugging
# import IPython; IPython.embed()
cv2.polylines(img,
lanes_lines["default"],
False, (255, 217, 82),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
cv2.polylines(img,
lanes_lines["green"],
False, (0, 255, 0),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
cv2.polylines(img,
lanes_lines["yellow"],
False, (255, 255, 0),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
cv2.polylines(img,
lanes_lines["red"],
False, (255, 0, 0),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
# plot crosswalks
crosswalks = []
for idx in elements_within_bounds(
center_in_world, self.bounds_info["crosswalks"]["bounds"],
raster_radius):
crosswalk = self.proto_API.get_crosswalk_coords(
self.bounds_info["crosswalks"]["ids"][idx])
xy_cross = cv2_subpixel(
transform_points(crosswalk["xyz"][:, :2], raster_from_world))
crosswalks.append(xy_cross)
cv2.polylines(img,
crosswalks,
True, (255, 117, 69),
lineType=cv2.LINE_AA,
shift=CV2_SHIFT)
return img
def to_rgb(self, in_im: np.ndarray, **kwargs: dict) -> np.ndarray:
return (in_im * 255).astype(np.uint8)
def _get_frame_data(mapAPI: MapAPI, frame: np.ndarray,
agents_frame: np.ndarray,
tls_frame: np.ndarray) -> FrameVisualization:
"""Get visualisation objects for the current frame.
:param mapAPI: mapAPI object (used for lanes, crosswalks etc..)
:param frame: the current frame (used for ego)
:param agents_frame: agents in this frame
:param tls_frame: the tls of this frame
:return: A FrameVisualization object. NOTE: trajectory are not included here
"""
ego_xy = frame["ego_translation"][:2]
#################
# plot lanes
lane_indices = indices_in_bounds(ego_xy,
mapAPI.bounds_info["lanes"]["bounds"], 50)
active_tl_ids = set(
filter_tl_faces_by_status(tls_frame, "ACTIVE")["face_id"].tolist())
lanes_vis: List[LaneVisualization] = []
for idx, lane_idx in enumerate(lane_indices):
lane_idx = mapAPI.bounds_info["lanes"]["ids"][lane_idx]
lane_tl_ids = set(mapAPI.get_lane_traffic_control_ids(lane_idx))
lane_colour = "gray"
for tl_id in lane_tl_ids.intersection(active_tl_ids):
lane_colour = COLORS[mapAPI.get_color_for_face(tl_id)]
lane_coords = mapAPI.get_lane_coords(lane_idx)
left_lane = lane_coords["xyz_left"][:, :2]
right_lane = lane_coords["xyz_right"][::-1, :2]
lanes_vis.append(
LaneVisualization(xs=np.hstack((left_lane[:, 0], right_lane[:,
0])),
ys=np.hstack((left_lane[:, 1], right_lane[:,
1])),
color=lane_colour))
#################
# plot crosswalks
crosswalk_indices = indices_in_bounds(
ego_xy, mapAPI.bounds_info["crosswalks"]["bounds"], 50)
crosswalks_vis: List[CWVisualization] = []
for idx in crosswalk_indices:
crosswalk = mapAPI.get_crosswalk_coords(
mapAPI.bounds_info["crosswalks"]["ids"][idx])
crosswalks_vis.append(
CWVisualization(xs=crosswalk["xyz"][:, 0],
ys=crosswalk["xyz"][:, 1],
color="yellow"))
#################
# plot ego and agents
agents_frame = np.insert(agents_frame, 0, get_ego_as_agent(frame))
box_world_coords = get_box_world_coords(agents_frame)
# ego
ego_vis = EgoVisualization(xs=box_world_coords[0, :, 0],
ys=box_world_coords[0, :, 1],
color="red",
center_x=agents_frame["centroid"][0, 0],
center_y=agents_frame["centroid"][0, 1])
# agents
agents_frame = agents_frame[1:]
box_world_coords = box_world_coords[1:]
agents_vis: List[AgentVisualization] = []
for agent, box_coord in zip(agents_frame, box_world_coords):
label_index = np.argmax(agent["label_probabilities"])
agent_type = PERCEPTION_LABELS[label_index]
agents_vis.append(
AgentVisualization(xs=box_coord[..., 0],
ys=box_coord[..., 1],
color="#1F77B4" if agent_type not in COLORS else
COLORS[agent_type],
track_id=agent["track_id"],
agent_type=PERCEPTION_LABELS[label_index],
prob=agent["label_probabilities"][label_index]))
return FrameVisualization(ego=ego_vis,
agents=agents_vis,
lanes=lanes_vis,
crosswalks=crosswalks_vis,
trajectories=[])