metropolis.py

# Copyright (c) Facebook, Inc. and its affiliates.

# Original copyright notice:
# nuScenes dev-kit.
# Code written by Oscar Beijbom, 2018.
import json
import sys
import time
from os import path
from typing import Tuple, List, Optional, Dict, Any, Union

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import (  # @manual=fbsource//third-party/pypi/matplotlib:matplotlib
    Axes,
)
from PIL import Image
from pyquaternion import Quaternion
from skimage.morphology import dilation
from skimage.segmentation import find_boundaries
from skimage.transform import warp

from .utils import pathmgr
from .utils.color_map import get_colormap, plot_deph_normalized_colormap
from .utils.data_classes import LidarPointCloud, Box, Box2d, EquiBox2d
from .utils.geo import TopocentricConverter
from .utils.geometry_utils import (
    view_points,
    view_points_eq,
    transform_matrix,
    box_in_image,
    inverse_map_eq,
)

# GDAL import is optional
try:
    from osgeo import gdal
except ModuleNotFoundError:
    gdal = None

PYTHON_VERSION = sys.version_info[0]

if not PYTHON_VERSION == 3:
    raise ValueError("Metropolis dev-kit only supports Python version 3.")


EYE4 = np.eye(4)


class Metropolis:
    def __init__(
        self,
        split: str,
        dataroot: str,
        verbose: bool = True,
        map_resolution: float = 0.1,
    ):
        self.split = split
        self.dataroot = dataroot
        self.verbose = verbose
        self.table_names = [
            "category",
            "attribute",
            "instance",
            "sensor",
            "calibrated_sensor",
            "ego_pose",
            "panoptic",
            "scene",
            "sample",
            "sample_data",
            "sample_annotation",
            "sample_annotation_2d",
        ]

        assert pathmgr.exists(
            self.table_root
        ), f"Database version not found: {self.table_root}"

        start_time = time.time()
        if verbose:
            print(f"======\nLoading Metropolis tables for split {self.split}...")

        # Explicitly assign tables to help the IDE determine valid class members.
        self.category = self.__load_table__("category")
        self.attribute = self.__load_table__("attribute")
        self.instance = self.__load_table__("instance")
        self.sensor = self.__load_table__("sensor")
        self.calibrated_sensor = self.__load_table__("calibrated_sensor")
        self.ego_pose = self.__load_table__("ego_pose")
        self.panoptic = self.__load_table__("panoptic")
        self.scene = self.__load_table__("scene")
        self.sample = self.__load_table__("sample")
        self.sample_data = self.__load_table__("sample_data")
        self.sample_annotation = self.__load_table__("sample_annotation")
        self.sample_annotation_2d = self.__load_table__("sample_annotation_2d")
        self.geo = self.__load_table__("geo")

        # Initialize the colormap which maps from class names to RGB values.
        self.colormap = get_colormap()

        if verbose:
            for table in self.table_names:
                print(f"{len(getattr(self, table))} {table},")
            print(f"Done loading in {time.time() - start_time:.3f} seconds.\n======")

        # Make reverse indexes for common lookups.
        self.__make_reverse_index__(verbose)

    @property
    def table_root(self) -> str:
        """Returns the folder where the tables are stored for the relevant version."""
        return path.join(self.dataroot, self.split)

    def __load_table__(self, table_name: str) -> List[Dict[str, Any]]:
        """Loads a table."""
        with pathmgr.open(path.join(self.table_root, f"{table_name}.json")) as f:
            table = json.load(f)
        return table

    def __make_reverse_index__(self, verbose: bool) -> None:  # noqa C901
        """De-normalizes database to create reverse indices for common cases.

        Args:
            verbose: Whether to print outputs.
        """

        start_time = time.time()
        if verbose:
            print("Reverse indexing ...")

        # Store the mapping from token to table index for each table.
        self._token2ind = {}
        for table in self.table_names:
            self._token2ind[table] = {}

            for ind, member in enumerate(getattr(self, table)):
                self._token2ind[table][member["token"]] = ind

        # Decorate (adds short-cut) sample_annotation table with for category name.
        for record in self.sample_annotation:
            inst = self.get("instance", record["instance_token"])
            record["category_name"] = self.get("category", inst["category_token"])[
                "name"
            ]

        # Do the same for sample_annotation_2d
        for record in self.sample_annotation_2d:
            inst = self.get("instance", record["instance_token"])
            record["category_name"] = self.get("category", inst["category_token"])[
                "name"
            ]

        # Decorate (adds short-cut) sample_data with sensor information.
        for record in self.sample_data:
            cs_record = self.get("calibrated_sensor", record["calibrated_sensor_token"])
            sensor_record = self.get("sensor", cs_record["sensor_token"])
            record["sensor_modality"] = sensor_record["modality"]
            record["channel"] = sensor_record["channel"]

        # Reverse-index samples with sample_data, annotations and panoptic
        for record in self.sample:
            record["data"] = {}
            record["anns"] = []
            record["anns_2d"] = []

        for record in self.sample_data:
            sample_record = self.get("sample", record["sample_token"])
            sample_record["data"][record["channel"]] = record["token"]

        for ann_record in self.sample_annotation:
            sample_record = self.get("sample", ann_record["sample_token"])
            sample_record["anns"].append(ann_record["token"])

        for ann_record_2d in self.sample_annotation_2d:
            sample_record = self.get("sample", ann_record_2d["sample_token"])
            sample_record["anns_2d"].append(ann_record_2d["token"])

        for pano_record in self.panoptic:
            sample_record = self.get("sample", pano_record["sample_token"])
            sample_record["panoptic_token"] = pano_record["token"]

        if verbose:
            print(
                f"Done reverse indexing in {time.time() - start_time:.1f} seconds.\n======"
            )

    def get(self, table_name: str, token: str) -> Dict[str, Any]:
        """Returns a record from table in constant runtime.

        Args:
            table_name: Table name.
            token: Token of the record.

        Returns:
            Table record. See README.md for record details for each table.
        """
        assert table_name in self.table_names, f"Table {table_name} not found"

        return getattr(self, table_name)[self.getind(table_name, token)]

    def getind(self, table_name: str, token: str) -> int:
        """This returns the index of the record in a table in constant runtime.

        Args:
            table_name: Table name.
            token: Token of the record.

        Returns:
            The index of the record in table, table is an array.
        """
        return self._token2ind[table_name][token]

    def field2token(self, table_name: str, field: str, query: Any) -> List[str]:
        """This function queries all records for a certain field value, and returns
        the tokens for the matching records.

        Warning: this runs in linear time.

        Args:
            table_name: Table name.
            field: Field name. See README.md for details.
            query: Query to match against. Needs to type match the content of the
                query field.

        Returns:
            List of tokens for the matching records.
        """
        matches = []
        for member in getattr(self, table_name):
            if member[field] == query:
                matches.append(member["token"])
        return matches

    def get_box(self, sample_annotation_token: str) -> Box:
        """Instantiates a Box class from a sample annotation record.

        Args:
            sample_annotation_token: Unique sample_annotation identifier.
        """
        record = self.get("sample_annotation", sample_annotation_token)
        return Box(
            record["translation"],
            record["size"],
            Quaternion(record["rotation"]),
            name=record["category_name"],
            token=record["token"],
        )

    def get_box_2d(self, sample_annotation_2d_token: str) -> Box2d:
        """Instantiates a Box2d class from a 2D sample annotation record.

        Args:
            sample_annotation_token: Unique sample_annotation identifier.

        Returns:
            The box object.
        """
        record = self.get("sample_annotation_2d", sample_annotation_2d_token)
        return Box2d(
            record["bounding_box"], name=record["category_name"], token=record["token"]
        )

    def get_boxes(
        self, sample_data_token: str, get_all_visible: bool = False
    ) -> List[Box]:
        """Instantiates Boxes for all annotation for a particular sample_data record

        Args:
            sample_data_token: Unique sample_data identifier.
            get_all_visible: If true, retrieve annotations for all objects that are
                potentially visible from this sample_data (i.e. those that have a
                2D annotation in the corresponding 360 image). Otherwise, only
                return objects that have been annotated in 3D directly on this sample.

        Return:
            A list of boxes.
        """
        sd_record = self.get("sample_data", sample_data_token)
        curr_sample_record = self.get("sample", sd_record["sample_token"])

        if get_all_visible:
            instance_tokens = {
                self.get("sample_annotation_2d", sa_2d_token)["instance_token"]
                for sa_2d_token in curr_sample_record["anns_2d"]
            }
            return list(
                map(
                    self.get_box,
                    (
                        sa["token"]
                        for sa in self.sample_annotation
                        if sa["instance_token"] in instance_tokens
                    ),
                )
            )
        else:
            return list(map(self.get_box, curr_sample_record["anns"]))

    def get_boxes_2d(self, sample_data_token: str) -> List[Box2d]:
        """Instantiates 2D Boxes for all annotation for a particular sample_data record

        Args:
            sample_data_token: Unique sample_data identifier.

        Return:
            A list of 2D Boxes.
        """
        sd_record = self.get("sample_data", sample_data_token)
        curr_sample_record = self.get("sample", sd_record["sample_token"])
        return list(map(self.get_box_2d, curr_sample_record["anns_2d"]))

    def get_color(self, category_name: str) -> Tuple[int, int, int]:
        """Provides the default colors based on the category names.

        Args:
            category_name: Name of the category.

        Returns:
            Color for the category, or (0, 0, 0) if the category is not found.
        """
        return self.colormap.get(category_name, (0, 0, 0))

    def get_sample_data_path(self, sample_data_token: str) -> str:
        """Returns the path to a sample_data."""
        sd_record = self.get("sample_data", sample_data_token)
        return path.join(self.dataroot, sd_record["filename"])

    def get_sample_data(
        self,
        sample_data_token: str,
        selected_anntokens: Optional[List[str]] = None,
        selected_2d_anntokens: Optional[List[str]] = None,
        use_flat_vehicle_coordinates: bool = False,
        get_all_visible_boxes: bool = False,
    ) -> Tuple[
        str,
        List[Box],
        Optional[Union[List[Box2d], List[EquiBox2d]]],
        Optional[np.array],
    ]:
        """Returns the data path as well as all annotations related to that sample_data.

        Note that the boxes are transformed into the current sensor's coordinate frame.

        Args:
            sample_data_token: Sample_data token.
            selected_anntokens: If provided only return the selected 3D annotation.
            selected_2d_anntokens: If provided only return the selected 2D annotation.
            use_flat_vehicle_coordinates: Instead of the current sensor's coordinate
                frame, use ego frame which is aligned to z-plane in the world.
            get_all_visible_boxes: If true, retrieve 3D boxes for all objects that are
                potentially visible from this sample_data (i.e. those that have a
                2D annotation in the corresponding 360 image). Otherwise, only
                return boxes that have been annotated in 3D directly on this sample.

        Returns:
            data_path: Path to the data file.
            boxes: 3D bounding boxes.
            boxes_2d: 2D bounding boxes (only returned for cameras).
            camera_intrinsic: Camera intrinsics.
        """

        # Retrieve sensor & pose records
        sd_record = self.get("sample_data", sample_data_token)
        cs_record = self.get("calibrated_sensor", sd_record["calibrated_sensor_token"])
        sensor_record = self.get("sensor", cs_record["sensor_token"])
        pose_record = self.get("ego_pose", sd_record["ego_pose_token"])

        is_camera_like = (
            sensor_record["modality"] == "depth"
            or sensor_record["modality"] == "camera"
        )

        data_path = self.get_sample_data_path(sample_data_token)

        if (
            sensor_record["modality"] == "camera"
            or sensor_record["modality"] == "depth"
        ):
            cam_intrinsic = np.array(cs_record["camera_intrinsic"])
            imsize = (sd_record["width"], sd_record["height"])
        else:
            cam_intrinsic = None
            imsize = None

        #### 3D annotations ####

        # Retrieve all sample annotations and map to sensor coordinate system.
        if selected_anntokens is not None:
            boxes = list(map(self.get_box, selected_anntokens))
        else:
            boxes = self.get_boxes(sample_data_token, get_all_visible_boxes)

        # Make list of Box objects including coord system transforms.
        box_list = []
        for box in boxes:
            if use_flat_vehicle_coordinates:
                # Move box to ego vehicle coord system parallel to world z plane.
                yaw = Quaternion(pose_record["rotation"]).yaw_pitch_roll[0]
                box.translate(-np.array(pose_record["translation"]))
                box.rotate(
                    Quaternion(
                        scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]
                    ).inverse
                )
            else:
                # Move box to ego vehicle coord system.
                box.translate(-np.array(pose_record["translation"]))
                box.rotate(Quaternion(pose_record["rotation"]).inverse)

                #  Move box to sensor coord system.
                box.translate(-np.array(cs_record["translation"]))
                box.rotate(Quaternion(cs_record["rotation"]).inverse)

            # For perspective camera-like sensors, check if the box is visible
            if is_camera_like and sensor_record["channel"] != "CAM_EQUIRECTANGULAR":
                if not box_in_image(box, cam_intrinsic, imsize):
                    continue

            box_list.append(box)

        #### 2D Annotations, only for cameras ####
        box_2d_list = None

        if is_camera_like:
            if selected_2d_anntokens is not None:
                boxes_2d = list(map(self.get_box_2d, selected_2d_anntokens))
            else:
                boxes_2d = self.get_boxes_2d(sample_data_token)

            if sensor_record["channel"] == "CAM_EQUIRECTANGULAR":
                # For equirectangular images just return all boxes
                box_2d_list = boxes_2d
            else:
                # For perspective images, project the 2D boxes from the corresponding
                # equirectangular image
                sample_record = self.get("sample", sd_record["sample_token"])

                box_2d_list = []
                if "CAM_EQUIRECTANGULAR" in sample_record["data"]:
                    sd_eq_record = self.get(
                        "sample_data", sample_record["data"]["CAM_EQUIRECTANGULAR"]
                    )
                    cs_eq_record = self.get(
                        "calibrated_sensor", sd_eq_record["calibrated_sensor_token"]
                    )

                    for box_2d in boxes_2d:
                        box_2d_eq = EquiBox2d.from_box_2d(
                            box_2d,
                            Quaternion(cs_eq_record["rotation"]),
                            Quaternion(cs_record["rotation"]),
                            cam_intrinsic,
                            (sd_eq_record["width"], sd_eq_record["height"]),
                            (sd_record["width"], sd_record["height"]),
                        )

                        if box_2d_eq is not None:
                            box_2d_list.append(box_2d_eq)

        return data_path, box_list, box_2d_list, cam_intrinsic

    def render_pointcloud_in_image(
        self,
        sample_token: str,
        dot_size: int = 0.5,
        downsample: int = 20,
        pointsensor_channel: str = "MVS",
        camera_channel: str = "CAM_FRONT",
        out_path: Optional[str] = None,
        ax: Optional[Axes] = None,
        nsweeps: int = 1,
    ):
        """Scatter-plots a point-cloud on top of equirectangular image.

        Args:
            sample_token: Sample token.
            dot_size: Scatter plot dot size.
            downsample: Downsampling factor.
            pointsensor_channel: Pointcloud channel name, e.g. 'MVS'.
            camera_channel: Camera channel name, e.g. 'CAM_FRONT'.
            out_path: Optional path to save the rendered figure to disk.
            ax: Optional existing matplotlib axes object to draw on.
            nsweeps: Number of sweeps for lidar and radar.
        """
        sample_record = self.get("sample", sample_token)

        # Here we just grab the front camera and the point sensor.
        pointsensor_token = sample_record["data"][pointsensor_channel]
        camera_token = sample_record["data"][camera_channel]

        points, coloring, im = self.map_pointcloud_to_image(
            pointsensor_token,
            camera_token,
            nsweeps=nsweeps,
        )

        # Init axes.
        if ax is None:
            _, ax = plt.subplots(1, 1, figsize=(18, 16), dpi=160)

        ax.imshow(im)
        ax.scatter(
            points[0, 0::downsample],
            points[1, 0::downsample],
            c=coloring[0::downsample],
            s=dot_size,
        )
        ax.axis("off")

        if out_path is not None:
            plt.xlim(0, im.size[0])
            plt.ylim(im.size[1], 0)

            with pathmgr.open(out_path, "wb") as fid:
                plt.savefig(fid, bbox_inches="tight", pad_inches=0)

    def map_pointcloud_to_image(
        self,
        pointsensor_token: str,
        camera_token: str,
        nsweeps: int = 1,
        min_dist: float = 1.0,
    ) -> Tuple:
        """Given a point sensor (e.g. lidar / mvs) token and camera sample_data token,
        load point-cloud and map it to an image.

        Args:
            pointsensor_token: Lidar/mvs sample_data token.
            camera_token: Camera sample_data token.
            nsweeps: Number of sweeps for lidar and mvs.
            min_dist: Distance from the camera below which points are discarded.

        Returns:
            pointcloud
            coloring
            image
        """
        # Find all relevant records
        sd_cam_record = self.get("sample_data", camera_token)
        sd_pts_record = self.get("sample_data", pointsensor_token)
        cs_cam_record = self.get(
            "calibrated_sensor", sd_cam_record["calibrated_sensor_token"]
        )
        cs_pts_record = self.get(
            "calibrated_sensor", sd_pts_record["calibrated_sensor_token"]
        )
        sample_pts_record = self.get("sample", sd_pts_record["sample_token"])
        ego_pose_cam = self.get("ego_pose", sd_cam_record["ego_pose_token"])
        ego_pose_pts = self.get("ego_pose", sd_pts_record["ego_pose_token"])

        # Load the point cloud
        pc, times = LidarPointCloud.from_file_multisweep(
            self,
            sample_pts_record,
            sd_pts_record["channel"],
            sd_pts_record["channel"],
            nsweeps,
        )

        # Load the image
        with pathmgr.open(
            path.join(self.dataroot, sd_cam_record["filename"]), "rb"
        ) as fid:
            im = Image.open(fid)
            im.load()

        # Transform the point cloud to the camera frame
        # 1. From point sensor to world
        pc.rotate(Quaternion(cs_pts_record["rotation"]).rotation_matrix)
        pc.translate(np.array(cs_pts_record["translation"]))
        pc.rotate(Quaternion(ego_pose_pts["rotation"]).rotation_matrix)
        pc.translate(np.array(ego_pose_pts["translation"]))

        # 2. From world to camera
        pc.translate(-np.array(ego_pose_cam["translation"]))
        pc.rotate(Quaternion(ego_pose_cam["rotation"]).rotation_matrix.T)
        pc.translate(-np.array(cs_cam_record["translation"]))
        pc.rotate(Quaternion(cs_cam_record["rotation"]).rotation_matrix.T)

        # Project the points to the camera plane / sphere
        if sd_cam_record["channel"] == "CAM_EQUIRECTANGULAR":
            # Equirectangular mode
            depths = np.sqrt((pc.points[:3, :] ** 2).sum(axis=0))

            # Do projection
            points = view_points_eq(pc.points[:3, :], im.size[0], im.size[1])
        else:
            # Projective mode
            depths = pc.points[2, :]

            # Do projection
            points = view_points(
                pc.points[:3, :],
                np.array(cs_cam_record["camera_intrinsic"]),
                normalize=True,
            )

            # Filter out points that are behind the camera, outside the frame or too close
            mask = np.ones(depths.shape[0], dtype=bool)
            mask = np.logical_and(mask, depths > min_dist)
            mask = np.logical_and(mask, points[0, :] > 1)
            mask = np.logical_and(mask, points[0, :] < im.size[0] - 1)
            mask = np.logical_and(mask, points[1, :] > 1)
            mask = np.logical_and(mask, points[1, :] < im.size[1] - 1)

            points = points[:, mask]
            depths = depths[mask]

        return points, depths, im

    def render_sample_data(  # noqa C901
        self,
        sample_data_token: str,
        axes_limit: float = 40,
        ax: Optional[Axes] = None,
        nsweeps: int = 1,
        out_path: Optional[str] = None,
        use_flat_vehicle_coordinates: bool = True,
        show_3d_boxes: bool = False,
        show_all_visible_3d_boxes: bool = False,
        verbose: bool = False,
    ) -> None:
        """Render sample data onto axis.

        Args:
            sample_data_token: Sample_data token.
            axes_limit: Axes limit for lidar / mvs (measured in meters).
            ax: Axes onto which to render.
            nsweeps: Number of sweeps for lidar / mvs.
            out_path: Optional path to save the rendered figure to disk.
            use_flat_vehicle_coordinates: Instead of the current sensor's coordinate
                frame, use ego frame which is aligned to z-plane in the world.
            show_3d_boxes: When rendering images, the default is to show 2D boxes.
                If this is set to True, show 3D boxes instead.
            show_all_visible_3d_boxes: If true, when rendering 3D boxes we show all
                those that are potentially visible from this sample_data (i.e. those
                that have a 2D annotation in the corresponding 360 image). Otherwise,
                we only show those that have been annotated directly on this sample.
            verbose: Whether to display the image after it is rendered.
        """
        # Get sensor modality.
        sd_record = self.get("sample_data", sample_data_token)

        sensor_modality = sd_record["sensor_modality"]

        if sensor_modality == "lidar" or sensor_modality == "mvs":
            sample_rec = self.get("sample", sd_record["sample_token"])
            chan = sd_record["channel"]
            ref_chan = "MVS"
            ref_sd_token = sample_rec["data"][ref_chan]
            ref_sd_record = self.get("sample_data", ref_sd_token)

            # Get aggregated lidar point cloud in lidar frame.
            pc, times = LidarPointCloud.from_file_multisweep(
                self, sample_rec, chan, ref_chan, nsweeps=nsweeps
            )

            # By default we render the sample_data top down in the sensor frame.
            # This is slightly inaccurate when rendering the map as the sensor frame may not be perfectly upright.
            # Using use_flat_vehicle_coordinates we can render the map in the ego frame instead.
            if use_flat_vehicle_coordinates:
                # Retrieve transformation matrices for reference point cloud.
                cs_record = self.get(
                    "calibrated_sensor", ref_sd_record["calibrated_sensor_token"]
                )
                pose_record = self.get("ego_pose", ref_sd_record["ego_pose_token"])
                ref_to_ego = transform_matrix(
                    translation=cs_record["translation"],
                    rotation=Quaternion(cs_record["rotation"]),
                )

                # Compute rotation between 3D vehicle pose and "flat" vehicle pose (parallel to global z plane).
                ego_yaw = Quaternion(pose_record["rotation"]).yaw_pitch_roll[0]
                rotation_vehicle_flat_from_vehicle = np.dot(
                    Quaternion(
                        scalar=np.cos(ego_yaw / 2), vector=[0, 0, np.sin(ego_yaw / 2)]
                    ).rotation_matrix,
                    Quaternion(pose_record["rotation"]).inverse.rotation_matrix,
                )
                vehicle_flat_from_vehicle = np.eye(4)
                vehicle_flat_from_vehicle[:3, :3] = rotation_vehicle_flat_from_vehicle
                viewpoint = np.dot(vehicle_flat_from_vehicle, ref_to_ego)
            else:
                viewpoint = np.eye(4)

            # Init axes.
            if ax is None:
                _, ax = plt.subplots(1, 1, figsize=(9, 9))

            # Show point cloud.
            points = view_points(pc.points[:3, :], viewpoint, normalize=False)
            dists = np.sqrt(np.sum(pc.points[:2, :] ** 2, axis=0))
            colors = np.minimum(1, dists / axes_limit / np.sqrt(2))
            point_scale = 0.2

            ax.scatter(points[0, :], points[1, :], c=colors, s=point_scale)

            # Show ego vehicle.
            ax.plot(0, 0, "x", color="red")

            # Get boxes in lidar frame.
            _, boxes, _, _ = self.get_sample_data(
                ref_sd_token,
                use_flat_vehicle_coordinates=use_flat_vehicle_coordinates,
                get_all_visible_boxes=show_all_visible_3d_boxes,
            )

            # Show boxes.
            for box in boxes:
                c = np.array(self.get_color(box.name)) / 255.0
                box.render(ax, view=np.eye(4), colors=(c, c, c))

            # Limit visible range.
            ax.set_xlim(-axes_limit, axes_limit)
            ax.set_ylim(-axes_limit, axes_limit)
        elif sensor_modality == "camera" or sensor_modality == "depth":
            # Load boxes and image.
            data_path, boxes, boxes_2d, camera_intrinsic = self.get_sample_data(
                sample_data_token,
                get_all_visible_boxes=show_all_visible_3d_boxes,
            )
            with pathmgr.open(data_path, "rb") as fid:
                data = Image.open(fid)
                data.load()

            # depthmaps are stored in 16 bit pngs, so it needs to be converted in meters
            if sensor_modality == "depth":
                data = np.array(data, dtype=np.float32) / 256.0
                max_depth_color = 120.0
                depth_map_color = plot_deph_normalized_colormap(data, max_depth_color)
                data = Image.fromarray(depth_map_color.astype(np.uint8))

            # Init axes.
            if ax is None:
                _, ax = plt.subplots(1, 1, figsize=(9, 16))

            # Show image.
            ax.imshow(data)

            # Show boxes.
            if show_3d_boxes:
                for box in boxes:
                    c = np.array(self.get_color(box.name)) / 255.0
                    if sd_record["channel"] == "CAM_EQUIRECTANGULAR":
                        box.render_eq(ax, data.size, colors=(c, c, c))
                    else:
                        box.render(
                            ax, view=camera_intrinsic, normalize=True, colors=(c, c, c)
                        )
            else:
                for box in boxes_2d:
                    c = np.array(self.get_color(box.name)) / 255.0
                    if isinstance(box, Box2d):
                        box.render(ax, sd_record["width"], color=c, linewidth=1)
                    else:
                        box.render(ax, color=c)

            # Limit visible range.
            ax.set_xlim(0, data.size[0])
            ax.set_ylim(data.size[1], 0)

        else:
            raise ValueError("Error: Unknown sensor modality!")

        ax.axis("off")
        ax.set_title(sd_record["channel"])
        ax.set_aspect("equal")

        if out_path is not None:
            with pathmgr.open(out_path, "wb") as fid:
                plt.savefig(fid, bbox_inches="tight", pad_inches=0, dpi=200)

        if verbose:
            plt.show()

    def render_aerial_view(
        self,
        sample_data_token: str,
        axes_limit: float = 40,
        ax: Optional[Axes] = None,
        nsweeps: int = 1,
        out_path: Optional[str] = None,
        verbose: bool = False,
    ) -> None:
        """Render a view of a point cloud onto an aerial image.

        NOTE: this function requires the GDAL library.

        Args:
            sample_data_token: Sample_data token of the point cloud to render.
            axes_limit: Axes limit for lidar / mvs (measured in meters).
            ax: Axes onto which to render.
            nsweeps: Number of sweeps for lidar / mvs.
            out_path: Optional path to save the rendered figure to disk.
            verbose: Whether to display the image after it is rendered.
        """
        assert gdal is not None, "GDAL is required to use this function"

        # Get sensor modality.
        sd_record = self.get("sample_data", sample_data_token)
        sensor_modality = sd_record["sensor_modality"]

        assert (
            sensor_modality == "lidar" or sensor_modality == "mvs"
        ), "This function is only available for pointclouds"

        # Get the point cloud
        sample_rec = self.get("sample", sd_record["sample_token"])
        pc, _ = LidarPointCloud.from_file_multisweep(
            self, sample_rec, sd_record["channel"], "MVS", nsweeps=nsweeps
        )

        # Create coordinates converter
        converter = TopocentricConverter(
            self.geo["reference"]["lat"],
            self.geo["reference"]["lon"],
            self.geo["reference"]["alt"],
        )

        # Transform to global coordinates
        ref_sd_record = self.get("sample_data", sample_rec["data"]["MVS"])
        cs_record = self.get(
            "calibrated_sensor", ref_sd_record["calibrated_sensor_token"]
        )
        pose_record = self.get("ego_pose", ref_sd_record["ego_pose_token"])

        sens_to_ego = transform_matrix(
            rotation=Quaternion(cs_record["rotation"]),
            translation=np.array(cs_record["translation"]),
        )
        ego_to_world = transform_matrix(
            rotation=Quaternion(pose_record["rotation"]),
            translation=np.array(pose_record["translation"]),
        )
        sens_to_world = np.dot(ego_to_world, sens_to_ego)
        pc.transform(sens_to_world)

        # Get center and view area corners in geo-referenced coordinates
        y_c_lla, x_c_lla, _ = converter.to_lla(
            pose_record["translation"][0],
            pose_record["translation"][1],
            pose_record["translation"][2],
        )
        y0_lla, x0_lla, _ = converter.to_lla(
            pose_record["translation"][0] - axes_limit,
            pose_record["translation"][1] - axes_limit,
            pose_record["translation"][2],
        )
        y1_lla, x1_lla, _ = converter.to_lla(
            pose_record["translation"][0] + axes_limit,
            pose_record["translation"][1] + axes_limit,
            pose_record["translation"][2],
        )

        # Take a crop of the aerial data
        crop = gdal.Warp(
            "",
            path.join(self.dataroot, self.geo["aerial"]["filename"]),
            dstSRS="WGS84",
            format="VRT",
            outputBounds=(x0_lla, y0_lla, x1_lla, y1_lla),
        )
        img = np.stack(
            [crop.GetRasterBand(i + 1).ReadAsArray() for i in range(3)], axis=-1
        )
        img = Image.fromarray(img)

        # Init axes.
        if ax is None:
            _, ax = plt.subplots(1, 1, figsize=(9, 9))

        # Draw the aerial data
        ax.imshow(img)

        # Project the point cloud to image coordinates
        y_p_lla, x_p_lla, _ = converter.to_lla(
            pc.points[0, :], pc.points[1, :], pc.points[2, :]
        )

        gt = crop.GetGeoTransform()
        R = np.array([[gt[1], gt[2]], [gt[4], gt[5]]])
        t = np.array([[gt[0]], [gt[3]]])
        points = np.linalg.solve(R, np.vstack([x_p_lla, y_p_lla]) - t)

        # Draw the points
        dists = np.sqrt(
            np.sum(
                (
                    pc.points[:2, :]
                    - np.array(pose_record["translation"][:2]).reshape(2, 1)
                )
                ** 2,
                axis=0,
            )
        )
        colors = np.minimum(1, dists / axes_limit / np.sqrt(2))
        ax.scatter(points[0, :], points[1, :], c=colors, s=0.2)

        # Show ego vehicle
        ev_pos = np.linalg.solve(R, np.array([[x_c_lla], [y_c_lla]]) - t)
        ax.plot(ev_pos[0], ev_pos[1], "x", color="red")

        # Limit visible range
        ax.set_xlim(0, img.size[0])
        ax.set_ylim(img.size[1], 0)

        # Produce final output and optionally save
        ax.axis("off")
        ax.set_title(sd_record["channel"])
        ax.set_aspect("equal")

        if out_path is not None:
            with pathmgr.open(out_path, "wb") as fid:
                plt.savefig(fid, bbox_inches="tight", pad_inches=0, dpi=200)

        if verbose:
            plt.show()

    def get_panoptic_mask(
        self, sample_data_token: str
    ) -> Tuple[Dict[str, Any], np.ndarray]:
        """Get the panoptic mask for a given image

        Since panoptic masks are originally computed on the equirectangular images,
        this function performs a warp operation in case the requested image is a
        perspective one (e.g. "CAM_FRONT"), which might introduce small artificats
        in the returned mask.

        Args:
            sample_data_token: Sample_data token of the image.

        Returns:
            pano_record: The panoptic record dictionary.
            pano: The panoptic mask, possibly re-projected to the requested image.
        """
        # Get the sample data
        sd_record = self.get("sample_data", sample_data_token)
        assert (
            sd_record["sensor_modality"] == "camera"
        ), "Panoptic masks can only be computed for images"

        # Get the corresponding sample and panoptic records
        sample_record = self.get("sample", sd_record["sample_token"])
        pano_record = self.get("panoptic", sample_record["panoptic_token"])

        # Load the panoptic image
        pano_record = self.get("panoptic", sample_record["panoptic_token"])
        with pathmgr.open(
            path.join(self.dataroot, pano_record["filename"]), "rb"
        ) as fid:
            pano = Image.open(fid)
            pano = np.array(pano, copy=True)

        # Warp the mask to the requested view if needed
        if sd_record["channel"] != "CAM_EQUIRECTANGULAR":
            # Get the corresponding equirectangulare image sample data
            sd_eq_record = self.get(
                "sample_data", sample_record["data"]["CAM_EQUIRECTANGULAR"]
            )

            # Get the sensor records
            cs_eq_record = self.get(
                "calibrated_sensor", sd_eq_record["calibrated_sensor_token"]
            )
            cs_record = self.get(
                "calibrated_sensor", sd_record["calibrated_sensor_token"]
            )

            # 3D transformation CURRENT -> EQ
            t_ego_to_eq = transform_matrix(
                np.array(cs_eq_record["translation"]),
                Quaternion(cs_eq_record["rotation"]),
                inverse=True,
            )
            t_cur_to_ego = transform_matrix(
                np.array(cs_record["translation"]), Quaternion(cs_record["rotation"])
            )
            t_cur_to_eq = np.dot(t_ego_to_eq, t_cur_to_ego)

            # Warp the panoptic mask using scikit image
            dtype = pano.dtype
            pano = warp(
                pano,
                inverse_map_eq,
                map_args={
                    "transform": t_cur_to_eq,
                    "intrinsics": np.array(cs_record["camera_intrinsic"]),
                    "eq_size": (sd_eq_record["width"], sd_eq_record["height"]),
                },
                output_shape=(sd_record["height"], sd_record["width"]),
                order=0,
                preserve_range=True,
            ).astype(dtype)

        return pano_record, pano

    def render_panoptic(
        self, sample_data_token: str, out_path: Optional[str] = None
    ) -> Image.Image:
        """Overlay an image with the corresponding panoptic segmentation

        Since panoptic masks are originally computed on the equirectangular images,
        this function performs a warp operation in case the requested image is a
        perspective one (e.g. "CAM_FRONT"), which might introduce small artificats
        in the returned mask.

        Args:
            sample_data_token: Sample_data token of the image.
            out_path: Path to save the rendered figure to disk.

        Returns:
            The rendered image as a PIL Image
        """
        # Load the image
        with pathmgr.open(
            self.get_sample_data_path(sample_data_token),
            "rb",
        ) as fid:
            img = Image.open(fid)
            img.load()

        # Load the panoptic image
        pano_record, pano = self.get_panoptic_mask(sample_data_token)

        # Make a palette for this image
        palette = []
        for cat_token in pano_record["category_tokens"]:
            if cat_token is not None:
                cat_name = self.get("category", cat_token)["name"]
                palette.append(self.get_color(cat_name))
            else:
                palette.append((0, 0, 0))

        # Remap panoptic mask to colors
        palette = np.array(palette, dtype=np.uint8)
        pano_color = palette[pano, :]
        pano_color = Image.fromarray(pano_color)

        # Highlight the edges between instances
        is_stuff = np.array(
            [it is None for it in pano_record["instance_tokens"]], dtype=bool
        )
        pano[is_stuff[pano]] = 0

        contours = (
            find_boundaries(pano, mode="outer", background=0).astype(np.uint8) * 255
        )
        contours = dilation(contours)

        contours = np.expand_dims(contours, -1).repeat(4, -1)
        contours_img = Image.fromarray(contours, mode="RGBA")

        # Render and save
        out = Image.blend(img, pano_color, 0.5).convert(mode="RGBA")
        out = Image.alpha_composite(out, contours_img)
        out = out.convert(mode="RGB")

        if out_path is not None:
            with pathmgr.open(out_path, "wb") as fid:
                out.save(fid)
        return out