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