def load_map_data(dataroot, location):
# Load the NuScenes map object
nusc_map = NuScenesMap(dataroot, location)
map_data = OrderedDict()
for layer in nusc_utils.STATIC_CLASSES:
# Retrieve all data associated with the current layer
records = getattr(nusc_map, layer)
polygons = list()
# Drivable area records can contain multiple polygons
if layer == 'drivable_area':
for record in records:
# Convert each entry in the record into a shapely object
for token in record['polygon_tokens']:
poly = nusc_map.extract_polygon(token)
if poly.is_valid:
polygons.append(poly)
else:
for record in records:
# Convert each entry in the record into a shapely object
poly = nusc_map.extract_polygon(record['polygon_token'])
if poly.is_valid:
polygons.append(poly)
# Store as an R-Tree for fast intersection queries
map_data[layer] = STRtree(polygons)
return map_data
def get_drivable_area_mask(nusc_map: NuScenesMap, scale_h: int = 2, scale_w: int = 2) -> Tuple[np.ndarray, List[float]]:
if nusc_map.map_name == 'singapore-onenorth':
map_dims = [1586, 2026]
elif nusc_map.map_name == 'singapore-hollandvillage':
map_dims = [2810, 3000]
elif nusc_map.map_name == 'singapore-queenstown':
map_dims = [3230, 3688]
elif nusc_map.map_name == 'boston-seaport':
map_dims = [2980, 2120]
else:
raise Exception('Error: Invalid map!')
patch_box = [map_dims[0] / 2, map_dims[1] / 2, map_dims[1], map_dims[0]]
patch = get_patch(patch_box, 0.0)
map_scale = 2
canvas_size = (patch_box[2]*scale_h, patch_box[3]*scale_w)
map_mask = np.zeros(canvas_size, np.uint8)
records = nusc_map.drivable_area
for record in records:
polygons = [nusc_map.extract_polygon(polygon_token) for polygon_token in record['polygon_tokens']]
for polygon in polygons:
new_polygon = polygon.intersection(patch)
new_polygon = affinity.scale(new_polygon, xfact=map_scale, yfact=map_scale, origin=(0, 0))
if not new_polygon.is_empty:
if new_polygon.geom_type is 'Polygon':
new_polygon = MultiPolygon([new_polygon])
map_mask = mask_for_polygons(new_polygon, map_mask)
return map_mask
