Baselines for panoptic challenges (nutonomy#633)

* Add scripts to fuse predictions for lidarseg and tracking / detection to get predictions for panoptic tracking / detection

* Update description in headers

* Tidy up resulting folder structure which will be created when running baseline.py

* Add docstrings for scripts

* Fix some typos in example usage of baseline.py

* Combine merge_lidarseg_and_tracking.py and merge_lidarseg_and_detection.py

* Rename merge_lidarseg_and_tracking.py to merge_for_panoptic.py

* Add comments, prevent divison by zero

* Fix list comprehensions in sort_confidence

* Change name from merge_for_panoptic.py to get_panoptic_from_seg_det_or_track.py

* No confidence thresholding of detection boxes

python-sdk/nuscenes/eval/panoptic/baselines.py

@@ -0,0 +1,164 @@
+"""
+Script to generate baselines for nuScenes-panoptic tasks.
+Code written by Motional and the Robot Learning Lab, University of Freiburg.
+"""
+import argparse
+import os
+import itertools
+from typing import List
+import zipfile
+
+from nuscenes.nuscenes import NuScenes
+from nuscenes.eval.panoptic.evaluate import NuScenesPanopticEval
+from nuscenes.eval.panoptic.get_panoptic_from_seg_det_or_track import generate_panoptic_labels
+
+
+def prepare_files(method_names: List[str], root_dir: str) -> None:
+    """
+    Prepare the files containing the predictions of the various method names.
+    :param method_names: A list of method names.
+    :param root_dir: The directory where the predictions of the various methods are stored at.
+    """
+    for method_name in method_names:
+        zip_path_to_predictions_by_method = os.path.join(root_dir, method_name + '.zip')
+        dir_path_to_predictions_by_method = os.path.join(root_dir, method_name)
+        assert os.path.exists(zip_path_to_predictions_by_method), 'Error: Zip file for method {} does not exist at {}.'\
+            .format(method_name, zip_path_to_predictions_by_method)
+        zip_ref = zipfile.ZipFile(zip_path_to_predictions_by_method, 'r')
+        zip_ref.extractall(dir_path_to_predictions_by_method)
+        zip_ref.close()
+
+
+def get_prediction_json_path(prediction_dir: str) -> str:
+    """
+    Get the name of the json file in a directory (abort if there is more than one).
+    :param prediction_dir: Path to the directory to check for the json file.
+    :return: Absolute path to the json file.
+    """
+
+    files_in_dir = os.listdir(prediction_dir)
+    files_in_dir = [f for f in files_in_dir if f.endswith('.json')]
+    assert len(files_in_dir) == 1, 'Error: The submission .zip file must contain exactly one .json file.'
+
+    prediction_json_path = os.path.join(prediction_dir, files_in_dir[0])
+    assert os.path.exists(prediction_json_path), \
+        'Error: JSON result file {} does not exist!'.format(prediction_json_path)
+
+    return prediction_json_path
+
+
+def main(out_dir: str,
+         lidarseg_preds_dir: str,
+         lidarseg_method_names: List[str],
+         det_or_track_preds_dir: str,
+         det_or_track_method_names: List[str],
+         task: str = 'tracking',
+         version: str = 'v1.0-test',
+         dataroot: str = '/data/sets/nuscenes') -> None:
+    """
+    Create baselines for a given panoptic task by merging the predictions of lidarseg and either tracking or detection
+    methods.
+    :param out_dir: Path to save any output to.
+    :param lidarseg_preds_dir: Path to the directory where the lidarseg predictions are stored.
+    :param lidarseg_method_names: List of lidarseg method names.
+    :param det_or_track_preds_dir: Path to the directory which contains the predictions from some methods to merge with
+        those of lidarseg to create panoptic predictions of a particular task.
+    :param det_or_track_method_names: List of tracking (or detection) method names to merge with lidarseg to create
+        panoptic predictions.
+    :param task: The task to create the panoptic predictions for and run evaluation on (either tracking or
+        segmentation).
+    :param version: Version of nuScenes to use (e.g. "v1.0", ...).
+    :param dataroot: Path to the tables and data for the specified version of nuScenes.
+    """
+    # Prepare the required files.
+    prepare_files(lidarseg_method_names, lidarseg_preds_dir)
+    prepare_files(det_or_track_method_names, det_or_track_preds_dir)
+
+    nusc = NuScenes(version=version, dataroot=dataroot)
+    eval_set = nusc.version.split('-')[-1]
+
+    # Get all possible pairwise permutations.
+    baselines = list(itertools.product(lidarseg_method_names, det_or_track_method_names))
+    print('There are {} baselines: {}'.format(len(baselines), baselines))
+
+    # Get the predictions for the panoptic task at hand.
+    for i, (lidarseg_method, det_or_track_method) in enumerate(baselines):
+        print('{:02d}/{:02d}: Getting predictions for panoptic {} from {} and {}.'
+              .format(i + 1, len(baselines), task, lidarseg_method, det_or_track_method))
+
+        dir_to_save_panoptic_preds_to = os.path.join(out_dir, task, 'panoptic_predictions',
+                                                     '{}_with_{}'.format(lidarseg_method, det_or_track_method))
+        os.makedirs(dir_to_save_panoptic_preds_to, exist_ok=True)
+
+        dir_of_lidarseg_method_preds = os.path.join(lidarseg_preds_dir, lidarseg_method)
+
+        json_of_preds_by_det_or_track_method = get_prediction_json_path(
+            os.path.join(det_or_track_preds_dir, det_or_track_method))
+
+        generate_panoptic_labels(nusc,
+                                 dir_of_lidarseg_method_preds,
+                                 json_of_preds_by_det_or_track_method,
+                                 eval_set=eval_set,
+                                 task=task,
+                                 out_dir=dir_to_save_panoptic_preds_to,
+                                 verbose=True)
+        print('Panoptic {} predictions saved at {}.'.format(task, dir_to_save_panoptic_preds_to))
+
+        print('{:02d}/{:02d}: Evaluating predictions for panoptic {} from {} and {}.'
+              .format(i + 1, len(baselines), task, lidarseg_method, det_or_track_method))
+        dir_to_save_evaluation_results_to = os.path.join(out_dir, task, 'panoptic_eval_results',
+                                                         '{}_with_{}'.format(lidarseg_method, det_or_track_method))
+        os.makedirs(dir_to_save_evaluation_results_to, exist_ok=True)
+        dir_of_panoptic_preds = dir_to_save_panoptic_preds_to
+        evaluator = NuScenesPanopticEval(nusc=nusc,
+                                         results_folder=dir_of_panoptic_preds,
+                                         eval_set=eval_set,
+                                         task=task,
+                                         min_inst_points=15,
+                                         out_dir=dir_to_save_evaluation_results_to,
+                                         verbose=True)
+        evaluator.evaluate()
+        print('Evaluation for panoptic {} using predictions merged from {} and {} saved at {}.'
+              .format(task, lidarseg_method, det_or_track_method, dir_to_save_evaluation_results_to))
+
+
+if __name__ == '__main__':
+    """
+    Example usage:
+        python baselines.py --out_dir ~/Desktop/logs/panoptic \
+                            --lidarseg_preds_dir ~/Desktop/logs/panoptic/submissions/lidarseg \
+                            --lidarseg_method_names 2D3DNet mit_han_lab \
+                            --det_or_track_preds_dir ~/Desktop/logs/panoptic/submissions/detection \
+                            --det_or_track_method_names crossfusion mmfusion polarstream \
+                            --task segmentation              
+    """
+    parser = argparse.ArgumentParser(description='Create baselines for a panoptic task (tracking or segmentation).')
+    parser.add_argument('--out_dir', type=str, help='Path to save any output to.')
+    parser.add_argument('--lidarseg_preds_dir', type=str,
+                        help='Path to the directory where the lidarseg predictions are stored.')
+    parser.add_argument('--lidarseg_method_names', nargs='+', help='List of lidarseg method names.')
+    parser.add_argument('--det_or_track_preds_dir', type=str,
+                        help='Path to the directory which contains the predictions from some methods to merge with '
+                             'those of lidarseg to create panoptic predictions of a particular task.')
+    parser.add_argument('--det_or_track_method_names', nargs='+',
+                        help='List of tracking (or detection) method names to merge with lidarseg to create panoptic '
+                             'predictions.')
+    parser.add_argument('--task', type=str, default='tracking',
+                        help='The task to create the panoptic predictions for and run evaluation on (either tracking '
+                             'or segmentation')
+    parser.add_argument('--version', type=str, default='v1.0-test',
+                        help='Version of nuScenes to use (e.g. "v1.0", ...).')
+    parser.add_argument('--dataroot', type=str, default='/data/sets/nuscenes',
+                        help='Path to the tables and data for the specified version of nuScenes.')
+
+    args = parser.parse_args()
+    print(args)
+
+    main(out_dir=args.out_dir,
+         lidarseg_preds_dir=args.lidarseg_preds_dir,
+         lidarseg_method_names=args.lidarseg_method_names,
+         det_or_track_preds_dir=args.det_or_track_preds_dir,
+         det_or_track_method_names=args.det_or_track_method_names,
+         task=args.task,
+         version=args.version,
+         dataroot=args.dataroot)

python-sdk/nuscenes/eval/panoptic/get_panoptic_from_seg_det_or_track.py

@@ -0,0 +1,188 @@
+"""
+Script to generate nuScenes-panoptic predictions from nuScene-lidarseg predictions and nuScenes-tracking or
+nuScenes-detection predictions.
+Code written by Motional and the Robot Learning Lab, University of Freiburg.
+"""
+import argparse
+import os
+from typing import List, Tuple, Union
+
+import numpy as np
+from tqdm import tqdm
+
+from nuscenes.eval.common.loaders import load_prediction, add_center_dist
+from nuscenes.eval.common.utils import boxes_to_sensor
+from nuscenes.eval.detection.data_classes import DetectionBox
+from nuscenes.eval.lidarseg.utils import get_samples_in_eval_set
+from nuscenes.eval.panoptic.utils import PanopticClassMapper
+from nuscenes.eval.tracking.data_classes import TrackingBox
+from nuscenes.nuscenes import NuScenes
+from nuscenes.utils.data_classes import LidarSegPointCloud
+from nuscenes.utils.geometry_utils import points_in_box
+
+
+OVERLAP_THRESHOLD = 0.5  # Amount by which an instance can overlap with other instances, before it is discarded.
+STUFF_START_COARSE_CLASS_ID = 11
+
+
+def generate_panoptic_labels(nusc: NuScenes,
+                             lidarseg_preds_folder: str,
+                             preds_json: str,
+                             eval_set: str,
+                             task: str = 'segmentation',
+                             out_dir: str = None,
+                             verbose: bool = False):
+    """
+    Generate NuScenes lidar panoptic predictions.
+    :param nusc: A instance of NuScenes.
+    :param lidarseg_preds_folder: Path to the directory where the lidarseg predictions are stored.
+    :param preds_json: Path of the json where the tracking / detection predictions are stored.
+    :param eval_set: Which dataset split to evaluate on, train, val or test.
+    :param task: The task to create the panoptic predictions for (either tracking or segmentation).
+    :param out_dir: Path to save any output to.
+    :param verbose: Whether to print any output.
+    """
+    assert task in ['tracking', 'segmentation'], \
+        'Error: Task can only be either `tracking` or `segmentation, not {}'.format(task)
+
+    sample_tokens = get_samples_in_eval_set(nusc, eval_set)
+    num_samples = len(sample_tokens)
+    mapper = PanopticClassMapper(nusc)
+    coarse2idx = mapper.get_coarse2idx()
+    if verbose:
+        print(f'There are {num_samples} samples.')
+
+    panoptic_subdir = os.path.join('panoptic', eval_set)
+    panoptic_dir = os.path.join(out_dir, panoptic_subdir)
+    os.makedirs(panoptic_dir, exist_ok=True)
+
+    box_type = TrackingBox if task == 'tracking' else DetectionBox
+
+    # Load the predictions.
+    pred_boxes_all, meta = load_prediction(preds_json, 1000, box_type, verbose=verbose)
+    pred_boxes_all = add_center_dist(nusc, pred_boxes_all)
+
+    inst_tok2id = {}  # Only used if task == 'tracking'.
+    for sample_token in tqdm(sample_tokens, disable=not verbose):
+        sample = nusc.get('sample', sample_token)
+
+        scene_token = sample['scene_token']  # Only used if task == 'tracking'.
+        if task == 'tracking' and scene_token not in inst_tok2id:
+            inst_tok2id[scene_token] = {}
+
+        # Get the sample data token of the point cloud.
+        sd_record = nusc.get('sample_data', sample['data']['LIDAR_TOP'])
+        cs_record = nusc.get('calibrated_sensor', sd_record['calibrated_sensor_token'])
+        pose_record = nusc.get('ego_pose', sd_record['ego_pose_token'])
+
+        # Load the predictions for the point cloud.
+        lidar_path = os.path.join(nusc.dataroot, sd_record['filename'])
+        lidarseg_pred_filename = os.path.join(lidarseg_preds_folder, 'lidarseg', nusc.version.split('-')[-1],
+                                              sd_record['token'] + '_lidarseg.bin')
+        lidar_seg = LidarSegPointCloud(lidar_path, lidarseg_pred_filename)
+
+        panop_labels = np.zeros(lidar_seg.labels.shape, dtype=np.uint16)
+        overlaps = np.zeros(lidar_seg.labels.shape, dtype=np.uint8)
+
+        pred_boxes = pred_boxes_all[sample_token]
+
+        # Tracking IDs will be None if pred_boxes is of DetectionBox type.
+        sorted_pred_boxes, pred_cls, tracking_ids = sort_confidence(pred_boxes)
+        sorted_pred_boxes = boxes_to_sensor(sorted_pred_boxes, pose_record, cs_record)
+
+        # Go through each box (a.k.a instance) and obtain the panoptic label for each.
+        for instance_id, (pred_box, cl, tracking_id) in enumerate(zip(sorted_pred_boxes, pred_cls, tracking_ids)):
+            cl_id = coarse2idx[cl]
+
+            if task == 'tracking':
+                if not inst_tok2id[scene_token]:
+                    inst_tok2id[scene_token][tracking_id] = 1
+                elif tracking_id not in inst_tok2id[scene_token]:
+                    inst_tok2id[scene_token][tracking_id] = len(inst_tok2id[scene_token]) + 1
+
+            msk = np.zeros(lidar_seg.labels.shape, dtype=np.uint8)
+            indices = np.where(points_in_box(pred_box, lidar_seg.points[:, :3].T))[0]
+
+            msk[indices] = 1
+            msk[np.logical_and(lidar_seg.labels != cl_id, msk == 1)] = 0
+            intersection = np.logical_and(overlaps, msk)
+
+            # If the current instance overlaps with previous instances by a certain threshold, then ignore it (note
+            # that the instances are processed in order of decreasing confidence).
+            if np.sum(intersection) / (np.float32(np.sum(msk)) + 1e-32) > OVERLAP_THRESHOLD:
+                continue
+            # Add non-overlapping part to output.
+            msk = msk - intersection
+            if task == 'tracking':
+                panop_labels[msk != 0] = (cl_id * 1000 + inst_tok2id[scene_token][tracking_id])
+            else:
+                panop_labels[msk != 0] = (cl_id * 1000 + instance_id + 1)
+            overlaps += msk
+
+        stuff_msk = np.logical_and(panop_labels == 0, lidar_seg.labels >= STUFF_START_COARSE_CLASS_ID)
+        panop_labels[stuff_msk] = lidar_seg.labels[stuff_msk] * 1000
+        panoptic_file = sd_record['token'] + '_panoptic.npz'
+        np.savez_compressed(os.path.join(panoptic_dir, panoptic_file), data=panop_labels.astype(np.uint16))
+
+
+def sort_confidence(boxes: List[Union[DetectionBox, TrackingBox]]) \
+        -> Tuple[List[Union[DetectionBox, TrackingBox]], List[str], List[Union[str, None]]]:
+    """
+    Sort a list of boxes by confidence.
+    :param boxes: A list of boxes.
+    :return: A list of boxes sorted by confidence and a list of classes and a list of tracking IDs (if available)
+        corresponding to each box.
+    """
+    scores = [box.tracking_score if isinstance(box, TrackingBox) else box.detection_score
+              for box in boxes]
+    inds = np.argsort(scores)[::-1]
+
+    sorted_bboxes = [boxes[ind] for ind in inds]
+
+    sorted_cls = [box.tracking_name if isinstance(box, TrackingBox) else box.detection_name
+                  for box in sorted_bboxes]
+
+    tracking_ids = [box.tracking_id if isinstance(box, TrackingBox) else None
+                    for box in sorted_bboxes]
+
+    return sorted_bboxes, sorted_cls, tracking_ids
+
+
+def main():
+    """
+    The main method for this script.
+    """
+    parser = argparse.ArgumentParser(
+        description='Generate panoptic from nuScenes-LiDAR segmentation and tracking results.')
+    parser.add_argument('--seg_path', type=str, help='The path to the segmentation results folder.')
+    parser.add_argument('--track_path', type=str, help='The path to the track json file.')
+    parser.add_argument('--eval_set', type=str, default='val',
+                        help='Which dataset split to evaluate on, train, val or test.')
+    parser.add_argument('--eval_set', type=str, default='val',
+                        help='Which dataset split to evaluate on, train, val or test.')
+    parser.add_argument('--task', type=str, default='segmentation',
+                        help='The task to create the panoptic predictions for (either tracking or segmentation).')
+    parser.add_argument('--dataroot', type=str, default='/data/sets/nuscenes', help='Default nuScenes data directory.')
+    parser.add_argument('--version', type=str, default='v1.0-trainval',
+                        help='Which version of the nuScenes dataset to evaluate on, e.g. v1.0-trainval.')
+    parser.add_argument('--verbose', type=bool, default=False, help='Whether to print to stdout.')
+    parser.add_argument('--out_dir', type=str, default=None, help='Folder to write the panoptic labels to.')
+    args = parser.parse_args()
+
+    out_dir = args.out_dir if args.out_dir is not None else f'nuScenes-panoptictrack-merged-prediction-{args.version}'
+
+    print(f'Start Generation... \nArguments: {args}')
+    nusc = NuScenes(version=args.version, dataroot=args.dataroot, verbose=args.verbose)
+
+    generate_panoptic_labels(nusc=nusc,
+                             lidarseg_preds_folder=args.seg_path,
+                             preds_json=args.track_path,
+                             eval_set=args.eval_set,
+                             task=args.task,
+                             out_dir=out_dir,
+                             verbose=args.verbose)
+    print(f'Generated results saved at {out_dir}. \nFinished.')
+
+
+if __name__ == '__main__':
+    main()

python-sdk/nuscenes/panoptic/generate_panoptic_labels.py

@@ -1,5 +1,5 @@
 """
-Script to generate nuScenes-panoptic ground truth data from nuScenes databse and nuScene-lidarseg.
+Script to generate nuScenes-panoptic ground truth data from nuScenes database and nuScene-lidarseg.
 Code written by Motional and the Robot Learning Lab, University of Freiburg.
 
 Example usage: