nutonomy / nuscenes-devkit

The devkit of the nuScenes dataset.
https://www.nuScenes.org
Other
2.24k stars 617 forks source link

Associate instance IDs from panoptic labels with bounding box #920

Closed thomas-enxuli closed 1 year ago

thomas-enxuli commented 1 year ago

Is there a way in this devkit to associate the 3D bounding box annotations with the instance ID labels from the panoptic ground truth? Thanks.

lubing-motional commented 1 year ago

HI @thomas-enxuli , we don't have direct function to do above mapping. But think it's possible. Note not all boxes have panoptic ground truth. Each panoptic labels from things should have only one box annotation based on generate_panoptic_labels.py. You may try and feedback with following code snippet for your prupose (I did not tested it).

import os
import numpy as np
from tqdm import tqdm

from nuscenes.nuscenes import NuScenes
from nuscenes.panoptic.panoptic_utils import STUFF_START_CLASS_ID
from nuscenes.utils.data_classes import LidarPointCloud
from nuscenes.utils.geometry_utils import points_in_box
from nuscenes.utils.data_io import load_bin_file

def get_ann_token_to_panoptic_label_mapping(nusc: NuScenes) -> None:
    """
    :param nusc: NuScenes instance with panoptic data.
    Get the mapping from box annotation to panoptic labels.
    Note not every box has valid panoptic labels, i.e. some cases there is no point inside the box with the
    same semantic category as the box type.
    """
    num_samples = len(nusc.sample)
    cat_name_to_idx = nusc.lidarseg_name2idx_mapping
    annotation_token_to_panoptic_map = {}  # {scene_token: {annotation_token: (panoptic_label, sample_token)}}
    for sample_idx in tqdm(range(num_samples)):
        curr_sample = nusc.sample[sample_idx]
        scene_token = curr_sample['scene_token']
        if scene_token not in annotation_token_to_panoptic_map:
            annotation_token_to_panoptic_map[scene_token] = {}
        lidar_token = curr_sample['data']['LIDAR_TOP']
        # Load points.
        point_path = os.path.join(nusc.dataroot, nusc.get('sample_data', lidar_token)['filename'])
        points = LidarPointCloud.from_file(point_path).points.T  # (N, 4), where N is the number of points.
        # Load panoptic labels.
        label_path = os.path.join(nusc.dataroot, nusc.get('panoptic', lidar_token)['filename'])
        panoptic_labels = load_bin_file(label_path, 'panoptic')  # (N,)
        # Get point semantic labels.
        semantic_labels = panoptic_labels // 1000  # Semantic categories.
        # Get only thing points.
        thing_mask = np.logical_and(semantic_labels > 0, semantic_labels < STUFF_START_CLASS_ID)
        thing_points = points[thing_mask]
        thing_panoptic_labels = panoptic_labels[thing_mask]
        thing_semantic_labels = semantic_labels[thing_mask]
        # Traverse all annotation boxes, and its panoptic label is obtained as the panoptic label of any inside thing points if there is any.
        for ann_token in curr_sample['anns']:
            ann = nusc.get('sample_annotation', ann_token)
            _, boxes, _ = nusc.get_sample_data(lidar_token, selected_anntokens=[ann_token])
            indices = np.where(points_in_box(boxes[0], thing_points[:, :3].T))[0]  # Each box only include maximally one GT thing category.
            for index in range(len(indices)):
                if thing_semantic_labels[indices[index]] == cat_name_to_idx[ann['category_name']]:
                    panoptic_label = thing_panoptic_labels[indices[index]]
                    annotation_token_to_panoptic_map[scene_token][ann_token] = (panoptic_label, curr_sample.token)
                    break
    return annotation_token_to_panoptic_map