nuScenes map visualization

nuscenes.py

import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
from matplotlib.patches import Polygon

from nuscenes.map_expansion.map_api import NuScenesMap


cmap_neighbor = plt.get_cmap("Dark2")
COLOR = [cmap_neighbor(i) for i in range(cmap_neighbor.N)]

DATA_ROOT = "nuscenes/v1.0"
NUSC_MAP = {
    map_name: NuScenesMap(dataroot=DATA_ROOT, map_name=map_name)
    for map_name in ["boston-seaport", "singapore-hollandvillage", "singapore-onenorth", "singapore-queenstown"]
}
COLOR_MAP = dict(
        # drivable_area='#a6cee3',
        road_segment="#c6ddec", #'#1f78b4',
        # road_block='#b2df8a',
        lane="#c6ddec", #'#33a02c',
        ped_crossing="#d3ccd7", #'#e31a1c', #'#fb9a99',
        # walkway='#e31a1c',
        # stop_line='#fdbf6f',
        # carpark_area='#ff7f00',
        road_divider="#94a5b0", #'#cab2d6',
        lane_divider="#94a5b0", #'#6a3d9a',
        # traffic_light='#7e772e'
)

def render_map_patch(nusc_map, box_coords, alpha=0.5, figsize=(15, 15)):
    x_min, y_min, x_max, y_max = box_coords
    local_width = x_max - x_min
    local_height = y_max - y_min
    assert local_height > 0, 'Error: Map patch has 0 height!'
    local_aspect_ratio = local_width / local_height
    x_margin = np.minimum(local_width / 4, 20)
    y_margin = np.minimum(local_height / 4, 10)
    x_min -= x_margin
    x_max += x_margin
    y_min -= y_margin
    y_max += y_margin

    fig = plt.figure(figsize=figsize)
    ax = fig.add_axes([0, 0, 1, 1 / local_aspect_ratio])

    records = nusc_map.road_segment
    for record in records:
        polygon = nusc_map.extract_polygon(record["polygon_token"])
        ax.add_patch(Polygon(polygon.exterior.coords, ec=COLOR_MAP['road_divider'], fc=COLOR_MAP["road_segment"], alpha=alpha))
    records = nusc_map.lane
    for record in records:
        polygon = nusc_map.extract_polygon(record["polygon_token"])
        ax.add_patch(Polygon(polygon.exterior.coords, ec=COLOR_MAP['lane_divider'], fc=COLOR_MAP["lane"], alpha=alpha))
    records = nusc_map.ped_crossing
    for record in records:
        polygon = nusc_map.extract_polygon(record["polygon_token"])
        ax.add_patch(Polygon(polygon.exterior.coords, ec="#93858c", fc=COLOR_MAP["ped_crossing"], alpha=alpha))

    records = nusc_map.road_divider
    for record in records:
        line = nusc_map.extract_line(record["line_token"])
        if line.is_empty: continue
        ax.plot(*line.xy, color=COLOR_MAP['road_divider'], alpha=alpha)
    records = nusc_map.lane_divider
    for record in records:
        line = nusc_map.extract_line(record["line_token"])
        if line.is_empty: continue
        ax.plot(*line.xy, color=COLOR_MAP['lane_divider'], alpha=alpha)

    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)
    ax.grid(True)
    return fig, ax

def draw(path, filename, tar_dir, nusc_map):
    with open(os.path.join(path, filename), "rb") as f:
        result = pickle.load(f)
    hist = result["hist"][:, 0, :2]
    fut = result["fut"][:, 0]
    neighbor = result["neighbor"][:, 0, :, :2]
    neighbor_type = result["neighbor_type"]
    pred2000 = result["pred2000"][:,:,0]
    pred5 = result["pred5"][:,:,0]
    map_name = result["map"]

    valid_neighbor = np.all(neighbor[...,:2] < 1e5, axis=-1)

    min_x = min(min(hist[:,0]), min(fut[:,0]), min(np.reshape(neighbor[valid_neighbor], (-1, 2))[:, 0]))
    max_x = max(max(hist[:,0]), max(fut[:,0]), max(np.reshape(neighbor[valid_neighbor], (-1, 2))[:, 0]))
    min_y = min(min(hist[:,1]), min(fut[:,1]), min(np.reshape(neighbor[valid_neighbor], (-1, 2))[:, 1]))
    max_y = max(max(hist[:,1]), max(fut[:,1]), max(np.reshape(neighbor[valid_neighbor], (-1, 2))[:, 1]))

    fig, ax = render_map_patch(nusc_map[map_name], (min_x, min_y, max_x, max_y), figsize=(10, 10), alpha=1.0)

    ax.plot(hist[:,0], hist[:,1],'w-o',
        linewidth=4,
        markersize=3,
        zorder=650,
        path_effects=[pe.Stroke(linewidth=5, foreground='k'), pe.Normal()])
    ax.plot(fut[:,0], fut[:,1],'r-o',
        linewidth=4,
        markersize=3,
        zorder=650,
        path_effects=[pe.Stroke(linewidth=5, foreground='k'), pe.Normal()])
    
    for traj in pred5:
        ax.plot(traj[:,0], traj[:,1],'y-o', alpha=0.2,
            linewidth=4,
            markersize=3,
            zorder=650,
            path_effects=[pe.Stroke(linewidth=5, foreground='k', alpha=0.2), pe.Normal()])

    for i in range(neighbor.shape[1]):
        marker = "o" if "VEHICLE" in neighbor_type[i] or "EGO" in neighbor_type[i] else "s"
        traj = neighbor[:,i]
        valid = np.all(traj<1e5, axis=-1)
        traj = traj[valid]
        ax.plot(traj[:,0], traj[:,1],'-'+marker, color=COLOR[i%len(COLOR)],
            linewidth=4,
            markersize=3,
            zorder=640,
            path_effects=[pe.Stroke(linewidth=5, foreground='k'), pe.Normal()])


    fig.savefig(os.path.join(tar_dir, "{}.png".format(os.path.splitext(filename)[0])), dpi=300, bbox_inches="tight")
    plt.close()


if __name__ == "__main__":
    DATA_PATH = "/home/peix/nuscenes_pred"
    TAR_PATH = "/home/peix/nuscenes_plot"
    for root, _, files in os.walk(DATA_PATH):
        for filename in files:
            tar_dir = os.path.join(TAR_PATH, os.path.relpath(root, DATA_PATH))
            os.makedirs(tar_dir, exist_ok=True)
            draw(root, filename, tar_dir, NUSC_MAP)