init0905

projects/mmdet3d_plugin/datasets/evaluation/map/vector_eval.py

+import prettytable
+from typing import Dict, List, Optional
+from time import time
+from copy import deepcopy
+from multiprocessing import Pool
+from logging import Logger
+from functools import partial, cached_property
+
+import numpy as np
+from numpy.typing import NDArray
+from shapely.geometry import LineString
+
+import mmcv
+from mmcv import Config
+from mmdet.datasets import build_dataset, build_dataloader
+
+from .AP import instance_match, average_precision
+
+INTERP_NUM = 200 # number of points to interpolate during evaluation
+THRESHOLDS = [0.5, 1.0, 1.5] # AP thresholds
+N_WORKERS = 16 # num workers to parallel
+
+class VectorEvaluate(object):
+    """Evaluator for vectorized map.
+
+    Args:
+        dataset_cfg (Config): dataset cfg for gt
+        n_workers (int): num workers to parallel
+    """
+
+    def __init__(self, dataset_cfg: Config, n_workers: int=N_WORKERS) -> None:
+        self.dataset = build_dataset(dataset_cfg)
+        self.dataloader = build_dataloader(
+            self.dataset, samples_per_gpu=1, workers_per_gpu=n_workers, shuffle=False, dist=False)
+        classes = self.dataset.MAP_CLASSES
+        self.cat2id = {cls: i for i, cls in enumerate(classes)}
+        self.id2cat = {v: k for k, v in self.cat2id.items()}
+        self.n_workers = n_workers
+        self.thresholds = [0.5, 1.0, 1.5]
+        
+    @cached_property
+    def gts(self) -> Dict[str, Dict[int, List[NDArray]]]:
+        print('collecting gts...')
+        gts = {}
+        pbar = mmcv.ProgressBar(len(self.dataloader))
+        for data in self.dataloader:
+            token = deepcopy(data['img_metas'].data[0][0]['token'])
+            gt = deepcopy(data['vectors'].data[0][0])
+            gts[token] = gt
+            pbar.update()
+            del data # avoid dataloader memory crash
+        
+        return gts
+    
+    def interp_fixed_num(self, 
+                         vector: NDArray, 
+                         num_pts: int) -> NDArray:
+        ''' Interpolate a polyline.
+        
+        Args:
+            vector (array): line coordinates, shape (M, 2)
+            num_pts (int): 
+        
+        Returns:
+            sampled_points (array): interpolated coordinates
+        '''
+        line = LineString(vector)
+        distances = np.linspace(0, line.length, num_pts)
+        sampled_points = np.array([list(line.interpolate(distance).coords) 
+            for distance in distances]).squeeze()
+        
+        return sampled_points
+    
+    def interp_fixed_dist(self, 
+                          vector: NDArray,
+                          sample_dist: float) -> NDArray:
+        ''' Interpolate a line at fixed interval.
+        
+        Args:
+            vector (LineString): vector
+            sample_dist (float): sample interval
+        
+        Returns:
+            points (array): interpolated points, shape (N, 2)
+        '''
+        line = LineString(vector)
+        distances = list(np.arange(sample_dist, line.length, sample_dist))
+        # make sure to sample at least two points when sample_dist > line.length
+        distances = [0,] + distances + [line.length,] 
+        
+        sampled_points = np.array([list(line.interpolate(distance).coords)
+                                for distance in distances]).squeeze()
+        
+        return sampled_points
+
+    def _evaluate_single(self, 
+                         pred_vectors: List, 
+                         scores: List, 
+                         groundtruth: List, 
+                         thresholds: List, 
+                         metric: str='metric') -> Dict[int, NDArray]:
+        ''' Do single-frame matching for one class.
+        
+        Args:
+            pred_vectors (List): List[vector(ndarray) (different length)], 
+            scores (List): List[score(float)]
+            groundtruth (List): List of vectors
+            thresholds (List): List of thresholds
+        
+        Returns:
+            tp_fp_score_by_thr (Dict): matching results at different thresholds
+                e.g. {0.5: (M, 2), 1.0: (M, 2), 1.5: (M, 2)}
+        '''
+        pred_lines = []
+
+        # interpolate predictions
+        for vector in pred_vectors:
+            vector = np.array(vector)
+            vector_interp = self.interp_fixed_num(vector, INTERP_NUM)
+            pred_lines.append(vector_interp)
+        if pred_lines:
+            pred_lines = np.stack(pred_lines)
+        else:
+            pred_lines = np.zeros((0, INTERP_NUM, 2))
+
+        # interpolate groundtruth
+        gt_lines = []
+        for vector in groundtruth:
+            vector_interp = self.interp_fixed_num(vector, INTERP_NUM)
+            gt_lines.append(vector_interp)
+        if gt_lines:
+            gt_lines = np.stack(gt_lines)
+        else:
+            gt_lines = np.zeros((0, INTERP_NUM, 2))
+        
+        scores = np.array(scores)
+        tp_fp_list = instance_match(pred_lines, scores, gt_lines, thresholds, metric) # (M, 2)
+        tp_fp_score_by_thr = {}
+        for i, thr in enumerate(thresholds):
+            tp, fp = tp_fp_list[i]
+            tp_fp_score = np.hstack([tp[:, None], fp[:, None], scores[:, None]])
+            tp_fp_score_by_thr[thr] = tp_fp_score
+        
+        return tp_fp_score_by_thr # {0.5: (M, 2), 1.0: (M, 2), 1.5: (M, 2)}
+        
+    def evaluate(self, 
+                 result_path: str, 
+                 metric: str='chamfer', 
+                 logger: Optional[Logger]=None) -> Dict[str, float]:
+        ''' Do evaluation for a submission file and print evalution results to `logger` if specified.
+        The submission will be aligned by tokens before evaluation. We use multi-worker to speed up.
+        
+        Args:
+            result_path (str): path to submission file
+            metric (str): distance metric. Default: 'chamfer'
+            logger (Logger): logger to print evaluation result, Default: None
+        
+        Returns:
+            new_result_dict (Dict): evaluation results. AP by categories.
+        '''
+        results = mmcv.load(result_path)
+        results = results['results']
+        
+        # re-group samples and gt by label
+        samples_by_cls = {label: [] for label in self.id2cat.keys()}
+        num_gts = {label: 0 for label in self.id2cat.keys()}
+        num_preds = {label: 0 for label in self.id2cat.keys()}
+
+        # align by token
+        for token, gt in self.gts.items():
+            if token in results.keys():
+                pred = results[token]
+            else:
+                pred = {'vectors': [], 'scores': [], 'labels': []}
+            
+            # for every sample
+            vectors_by_cls = {label: [] for label in self.id2cat.keys()}
+            scores_by_cls = {label: [] for label in self.id2cat.keys()}
+
+            for i in range(len(pred['labels'])):
+                # i-th pred line in sample
+                label = pred['labels'][i]
+                vector = pred['vectors'][i]
+                score = pred['scores'][i]
+
+                vectors_by_cls[label].append(vector)
+                scores_by_cls[label].append(score)
+
+            for label in self.id2cat.keys():
+                new_sample = (vectors_by_cls[label], scores_by_cls[label], gt[label])
+                num_gts[label] += len(gt[label])
+                num_preds[label] += len(scores_by_cls[label])
+                samples_by_cls[label].append(new_sample)
+
+        result_dict = {}
+
+        print(f'\nevaluating {len(self.id2cat)} categories...')
+        start = time()
+        if self.n_workers > 0:
+            pool = Pool(self.n_workers)
+        
+        sum_mAP = 0
+        pbar = mmcv.ProgressBar(len(self.id2cat))
+        for label in self.id2cat.keys():
+            samples = samples_by_cls[label] # List[(pred_lines, scores, gts)]
+            result_dict[self.id2cat[label]] = {
+                'num_gts': num_gts[label],
+                'num_preds': num_preds[label]
+            }
+            sum_AP = 0
+
+            fn = partial(self._evaluate_single, thresholds=self.thresholds, metric=metric)
+            if self.n_workers > 0 and len(samples) > 81:
+                tpfp_score_list = pool.starmap(fn, samples)
+            else:
+                tpfp_score_list = []
+                for sample in samples:
+                    tpfp_score_list.append(fn(*sample))
+            
+            for thr in self.thresholds:
+                tp_fp_score = [i[thr] for i in tpfp_score_list]
+                tp_fp_score = np.vstack(tp_fp_score) # (num_dets, 3)
+                sort_inds = np.argsort(-tp_fp_score[:, -1])
+
+                tp = tp_fp_score[sort_inds, 0] # (num_dets,)
+                fp = tp_fp_score[sort_inds, 1] # (num_dets,)
+                tp = np.cumsum(tp, axis=0)
+                fp = np.cumsum(fp, axis=0)
+                eps = np.finfo(np.float32).eps
+                recalls = tp / np.maximum(num_gts[label], eps)
+                precisions = tp / np.maximum((tp + fp), eps)
+
+                AP = average_precision(recalls, precisions, 'area')
+                sum_AP += AP
+                result_dict[self.id2cat[label]].update({f'AP@{thr}': AP})
+
+            pbar.update()
+            
+            AP = sum_AP / len(self.thresholds)
+            sum_mAP += AP
+
+            result_dict[self.id2cat[label]].update({f'AP': AP})
+        
+        if self.n_workers > 0:
+            pool.close()
+        
+        mAP = sum_mAP / len(self.id2cat.keys())
+        result_dict.update({'mAP': mAP})
+        
+        print(f"finished in {time() - start:.2f}s")
+
+        # print results
+        table = prettytable.PrettyTable(['category', 'num_preds', 'num_gts'] + 
+                [f'AP@{thr}' for thr in self.thresholds] + ['AP'])
+        for label in self.id2cat.keys():
+            table.add_row([
+                self.id2cat[label], 
+                result_dict[self.id2cat[label]]['num_preds'],
+                result_dict[self.id2cat[label]]['num_gts'],
+                *[round(result_dict[self.id2cat[label]][f'AP@{thr}'], 4) for thr in self.thresholds],
+                round(result_dict[self.id2cat[label]]['AP'], 4),
+            ])
+        
+        from mmcv.utils import print_log
+        print_log('\n'+str(table), logger=logger)
+        mAP_normal = 0
+        for label in self.id2cat.keys():
+            for thr in self.thresholds:
+                mAP_normal += result_dict[self.id2cat[label]][f'AP@{thr}']
+        mAP_normal = mAP_normal / 9
+
+        print_log(f'mAP_normal = {mAP_normal:.4f}\n', logger=logger)
+        # print_log(f'mAP_hard = {mAP_easy:.4f}\n', logger=logger)
+
+        new_result_dict = {}
+        for name in self.cat2id:
+            new_result_dict[name] = result_dict[name]['AP']
+        new_result_dict['mAP_normal'] = mAP_normal
+        return new_result_dict
\ No newline at end of file

projects/mmdet3d_plugin/datasets/evaluation/motion/motion_eval_uniad.py

+# nuScenes dev-kit.
+# Code written by Holger Caesar & Oscar Beijbom, 2018.
+
+import argparse
+import json
+import os
+import random
+import time
+import tqdm
+from typing import Tuple, Dict, Any
+
+import numpy as np
+
+from nuscenes import NuScenes
+from nuscenes.eval.common.config import config_factory
+from nuscenes.eval.common.data_classes import EvalBoxes
+from nuscenes.eval.common.loaders import add_center_dist, filter_eval_boxes
+from nuscenes.eval.detection.algo import accumulate, calc_ap, calc_tp
+from nuscenes.eval.detection.constants import DETECTION_NAMES, ATTRIBUTE_NAMES, TP_METRICS
+from nuscenes.eval.detection.data_classes import DetectionConfig, DetectionMetrics, DetectionBox, \
+    DetectionMetricDataList, DetectionMetricData
+from nuscenes.eval.detection.render import summary_plot, class_pr_curve, class_tp_curve, dist_pr_curve, visualize_sample
+from nuscenes.prediction import PredictHelper, convert_local_coords_to_global
+from nuscenes.utils.splits import create_splits_scenes
+from nuscenes.eval.detection.utils import category_to_detection_name
+from nuscenes.eval.common.utils import quaternion_yaw, Quaternion
+from nuscenes.eval.common.utils import center_distance, scale_iou, yaw_diff, velocity_l2, attr_acc, cummean
+
+from .motion_utils import MotionBox, load_prediction, load_gt, accumulate
+MOTION_TP_METRICS = ['min_ade_err', 'min_fde_err', 'miss_rate_err']
+
+
+class MotionEval:
+    """
+    This is the official nuScenes detection evaluation code.
+    Results are written to the provided output_dir.
+
+    nuScenes uses the following detection metrics:
+    - Mean Average Precision (mAP): Uses center-distance as matching criterion; averaged over distance thresholds.
+    - True Positive (TP) metrics: Average of translation, velocity, scale, orientation and attribute errors.
+    - nuScenes Detection Score (NDS): The weighted sum of the above.
+
+    Here is an overview of the functions in this method:
+    - init: Loads GT annotations and predictions stored in JSON format and filters the boxes.
+    - run: Performs evaluation and dumps the metric data to disk.
+    - render: Renders various plots and dumps to disk.
+
+    We assume that:
+    - Every sample_token is given in the results, although there may be not predictions for that sample.
+
+    Please see https://www.nuscenes.org/object-detection for more details.
+    """
+    def __init__(self,
+                 nusc: NuScenes,
+                 config: DetectionConfig,
+                 result_path: str,
+                 eval_set: str,
+                 output_dir: str = None,
+                 verbose: bool = True,
+                 seconds: int = 12):
+        """
+        Initialize a DetectionEval object.
+        :param nusc: A NuScenes object.
+        :param config: A DetectionConfig object.
+        :param result_path: Path of the nuScenes JSON result file.
+        :param eval_set: The dataset split to evaluate on, e.g. train, val or test.
+        :param output_dir: Folder to save plots and results to.
+        :param verbose: Whether to print to stdout.
+        """
+        self.nusc = nusc
+        self.result_path = result_path
+        self.eval_set = eval_set
+        self.output_dir = output_dir
+        self.verbose = verbose
+        self.cfg = config
+
+        # Check result file exists.
+        # assert os.path.exists(result_path), 'Error: The result file does not exist!'
+
+        # Make dirs.
+        self.plot_dir = os.path.join(self.output_dir, 'plots')
+        if not os.path.isdir(self.output_dir):
+            os.makedirs(self.output_dir)
+        if not os.path.isdir(self.plot_dir):
+            os.makedirs(self.plot_dir)
+
+        # Load data.
+        if verbose:
+            print('Initializing nuScenes detection evaluation')
+        self.pred_boxes, self.meta = load_prediction(self.result_path, self.cfg.max_boxes_per_sample, MotionBox,
+                                                     verbose=verbose)
+        self.gt_boxes = load_gt(self.nusc, self.eval_set, MotionBox, verbose=verbose, seconds=seconds)
+
+        assert set(self.pred_boxes.sample_tokens) == set(self.gt_boxes.sample_tokens), \
+            "Samples in split doesn't match samples in predictions."
+
+        # Add center distances.
+        self.pred_boxes = add_center_dist(nusc, self.pred_boxes)
+        self.gt_boxes = add_center_dist(nusc, self.gt_boxes)
+
+        # Filter boxes (distance, points per box, etc.).
+        if verbose:
+            print('Filtering predictions')
+        self.pred_boxes = filter_eval_boxes(nusc, self.pred_boxes, self.cfg.class_range, verbose=verbose)
+        if verbose:
+            print('Filtering ground truth annotations')
+        self.gt_boxes = filter_eval_boxes(nusc, self.gt_boxes, self.cfg.class_range, verbose=verbose)
+
+        self.sample_tokens = self.gt_boxes.sample_tokens
+
+    def evaluate(self) -> Tuple[DetectionMetrics, DetectionMetricDataList]:
+        """
+        Performs the actual evaluation.
+        :return: A tuple of high-level and the raw metric data.
+        """
+        start_time = time.time()
+        self.cfg.class_names = ['car', 'pedestrian']
+        self.cfg.dist_ths = [2.0]
+
+        # -----------------------------------
+        # Step 1: Accumulate metric data for all classes and distance thresholds.
+        # -----------------------------------
+        if self.verbose:
+            print('Accumulating metric data...')
+        metric_data_list = DetectionMetricDataList()
+        metrics = {}
+        for class_name in self.cfg.class_names:
+            for dist_th in self.cfg.dist_ths:
+                md, EPA, EPA_ = accumulate(self.gt_boxes, self.pred_boxes, class_name, self.cfg.dist_fcn_callable, dist_th)
+                metric_data_list.set(class_name, dist_th, md)
+                metrics[f'{class_name}_EPA'] = EPA_
+
+        # -----------------------------------
+        # Step 2: Calculate metrics from the data.
+        # -----------------------------------
+        if self.verbose:
+            print('Calculating metrics...')
+        for class_name in self.cfg.class_names:
+            # Compute TP metrics.
+            for metric_name in MOTION_TP_METRICS:
+                metric_data = metric_data_list[(class_name, self.cfg.dist_th_tp)]
+                tp = calc_tp(metric_data, self.cfg.min_recall, metric_name)
+                metrics[f'{class_name}_{metric_name}']  = tp
+
+        return metrics, metric_data_list
+
+    def render(self, metrics: DetectionMetrics, md_list: DetectionMetricDataList) -> None:
+        """
+        Renders various PR and TP curves.
+        :param metrics: DetectionMetrics instance.
+        :param md_list: DetectionMetricDataList instance.
+        """
+        if self.verbose:
+            print('Rendering PR and TP curves')
+
+        def savepath(name):
+            return os.path.join(self.plot_dir, name + '.pdf')
+
+        summary_plot(md_list, metrics, min_precision=self.cfg.min_precision, min_recall=self.cfg.min_recall,
+                     dist_th_tp=self.cfg.dist_th_tp, savepath=savepath('summary'))
+
+        for detection_name in self.cfg.class_names:
+            class_pr_curve(md_list, metrics, detection_name, self.cfg.min_precision, self.cfg.min_recall,
+                           savepath=savepath(detection_name + '_pr'))
+
+            class_tp_curve(md_list, metrics, detection_name, self.cfg.min_recall, self.cfg.dist_th_tp,
+                           savepath=savepath(detection_name + '_tp'))
+
+        for dist_th in self.cfg.dist_ths:
+            dist_pr_curve(md_list, metrics, dist_th, self.cfg.min_precision, self.cfg.min_recall,
+                          savepath=savepath('dist_pr_' + str(dist_th)))
+
+    def main(self,
+             plot_examples: int = 0,
+             render_curves: bool = True) -> Dict[str, Any]:
+        """
+        Main function that loads the evaluation code, visualizes samples, runs the evaluation and renders stat plots.
+        :param plot_examples: How many example visualizations to write to disk.
+        :param render_curves: Whether to render PR and TP curves to disk.
+        :return: A dict that stores the high-level metrics and meta data.
+        """
+        if plot_examples > 0:
+            # Select a random but fixed subset to plot.
+            random.seed(42)
+            sample_tokens = list(self.sample_tokens)
+            random.shuffle(sample_tokens)
+            sample_tokens = sample_tokens[:plot_examples]
+
+            # Visualize samples.
+            example_dir = os.path.join(self.output_dir, 'examples')
+            if not os.path.isdir(example_dir):
+                os.mkdir(example_dir)
+            for sample_token in sample_tokens:
+                visualize_sample(self.nusc,
+                                 sample_token,
+                                 self.gt_boxes if self.eval_set != 'test' else EvalBoxes(),
+                                 # Don't render test GT.
+                                 self.pred_boxes,
+                                 eval_range=max(self.cfg.class_range.values()),
+                                 savepath=os.path.join(example_dir, '{}.png'.format(sample_token)))
+
+        # Run evaluation.
+        metrics, metric_data_list = self.evaluate()
+
+        return metrics
+
+class NuScenesEval(MotionEval):
+    """
+    Dummy class for backward-compatibility. Same as MotionEval.
+    """
+
+if __name__ == "__main__":
+
+    # Settings.
+    parser = argparse.ArgumentParser(description='Evaluate nuScenes detection results.',
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('result_path', type=str, help='The submission as a JSON file.')
+    parser.add_argument('--output_dir', type=str, default='~/nuscenes-metrics',
+                        help='Folder to store result metrics, graphs and example visualizations.')
+    parser.add_argument('--eval_set', type=str, default='val',
+                        help='Which dataset split to evaluate on, train, val or test.')
+    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('--config_path', type=str, default='',
+                        help='Path to the configuration file.'
+                             'If no path given, the CVPR 2019 configuration will be used.')
+    parser.add_argument('--plot_examples', type=int, default=10,
+                        help='How many example visualizations to write to disk.')
+    parser.add_argument('--render_curves', type=int, default=1,
+                        help='Whether to render PR and TP curves to disk.')
+    parser.add_argument('--verbose', type=int, default=1,
+                        help='Whether to print to stdout.')
+    args = parser.parse_args()
+
+    result_path_ = os.path.expanduser(args.result_path)
+    output_dir_ = os.path.expanduser(args.output_dir)
+    eval_set_ = args.eval_set
+    dataroot_ = args.dataroot
+    version_ = args.version
+    config_path = args.config_path
+    plot_examples_ = args.plot_examples
+    render_curves_ = bool(args.render_curves)
+    verbose_ = bool(args.verbose)
+
+    if config_path == '':
+        cfg_ = config_factory('detection_cvpr_2019')
+    else:
+        with open(config_path, 'r') as _f:
+            cfg_ = DetectionConfig.deserialize(json.load(_f))
+
+    nusc_ = NuScenes(version=version_, verbose=verbose_, dataroot=dataroot_)
+    nusc_eval = DetectionEval(nusc_, config=cfg_, result_path=result_path_, eval_set=eval_set_,
+                              output_dir=output_dir_, verbose=verbose_)
+    nusc_eval.main(plot_examples=plot_examples_, render_curves=render_curves_)

projects/mmdet3d_plugin/datasets/evaluation/motion/motion_utils.py

+# nuScenes dev-kit.
+# Code written by Holger Caesar & Oscar Beijbom, 2018.
+
+import argparse
+import json
+import os
+import random
+import time
+import tqdm
+from typing import Tuple, Dict, Any, Callable
+
+import numpy as np
+
+from nuscenes import NuScenes
+from nuscenes.eval.common.config import config_factory
+from nuscenes.eval.common.data_classes import EvalBoxes
+from nuscenes.eval.common.loaders import add_center_dist, filter_eval_boxes
+from nuscenes.eval.detection.algo import calc_ap, calc_tp
+from nuscenes.eval.detection.constants import DETECTION_NAMES, ATTRIBUTE_NAMES, TP_METRICS
+from nuscenes.eval.detection.data_classes import DetectionConfig, DetectionMetrics, DetectionBox, \
+    DetectionMetricDataList, DetectionMetricData
+from nuscenes.eval.detection.render import summary_plot, class_pr_curve, class_tp_curve, dist_pr_curve, visualize_sample
+from nuscenes.prediction import PredictHelper, convert_local_coords_to_global
+from nuscenes.utils.splits import create_splits_scenes
+from nuscenes.eval.detection.utils import category_to_detection_name
+from nuscenes.eval.common.utils import quaternion_yaw, Quaternion
+from nuscenes.eval.common.utils import center_distance, scale_iou, yaw_diff, velocity_l2, attr_acc, cummean
+
+
+motion_name_mapping = {
+    'car': 'car',
+    'truck': 'car',
+    'construction_vehicle': 'car',
+    'bus': 'car',
+    'trailer': 'car',
+    'motorcycle': 'car',
+    'bicycle': 'car',
+    'pedestrian': 'pedestrian',
+    'traffic_cone': 'barrier',
+    'barrier': 'barrier',
+}
+
+
+class MotionBox(DetectionBox):
+    """ Data class used during detection evaluation. Can be a prediction or ground truth."""
+
+    def __init__(self,
+                 sample_token: str = "",
+                 translation: Tuple[float, float, float] = (0, 0, 0),
+                 size: Tuple[float, float, float] = (0, 0, 0),
+                 rotation: Tuple[float, float, float, float] = (0, 0, 0, 0),
+                 velocity: Tuple[float, float] = (0, 0),
+                 ego_translation: [float, float, float] = (0, 0, 0),  # Translation to ego vehicle in meters.
+                 num_pts: int = -1,  # Nbr. LIDAR or RADAR inside the box. Only for gt boxes.
+                 detection_name: str = 'car',  # The class name used in the detection challenge.
+                 detection_score: float = -1.0,  # GT samples do not have a score.
+                 attribute_name: str = '',  # Box attribute. Each box can have at most 1 attribute.
+                 traj=None):  
+
+        super().__init__(sample_token, translation, size, rotation, velocity, ego_translation, num_pts)
+
+        assert detection_name is not None, 'Error: detection_name cannot be empty!'
+        assert detection_name in DETECTION_NAMES, 'Error: Unknown detection_name %s' % detection_name
+
+        assert attribute_name in ATTRIBUTE_NAMES or attribute_name == '', \
+            'Error: Unknown attribute_name %s' % attribute_name
+
+        assert type(detection_score) == float, 'Error: detection_score must be a float!'
+        assert not np.any(np.isnan(detection_score)), 'Error: detection_score may not be NaN!'
+
+        # Assign.
+        self.detection_name = detection_name
+        self.detection_score = detection_score
+        self.attribute_name = attribute_name
+        self.traj = traj
+
+    def __eq__(self, other):
+        return (self.sample_token == other.sample_token and
+                self.translation == other.translation and
+                self.size == other.size and
+                self.rotation == other.rotation and
+                self.velocity == other.velocity and
+                self.ego_translation == other.ego_translation and
+                self.num_pts == other.num_pts and
+                self.detection_name == other.detection_name and
+                self.detection_score == other.detection_score and
+                self.attribute_name == other.attribute_name and
+                np.all(self.traj == other.traj))
+
+    def serialize(self) -> dict:
+        """ Serialize instance into json-friendly format. """
+        return {
+            'sample_token': self.sample_token,
+            'translation': self.translation,
+            'size': self.size,
+            'rotation': self.rotation,
+            'velocity': self.velocity,
+            'ego_translation': self.ego_translation,
+            'num_pts': self.num_pts,
+            'detection_name': self.detection_name,
+            'detection_score': self.detection_score,
+            'attribute_name': self.attribute_name,
+            'traj': self.traj,
+        }
+
+    @classmethod
+    def deserialize(cls, content: dict):
+        """ Initialize from serialized content. """
+        return cls(sample_token=content['sample_token'],
+                   translation=tuple(content['translation']),
+                   size=tuple(content['size']),
+                   rotation=tuple(content['rotation']),
+                   velocity=tuple(content['velocity']),
+                   ego_translation=(0.0, 0.0, 0.0) if 'ego_translation' not in content
+                   else tuple(content['ego_translation']),
+                   num_pts=-1 if 'num_pts' not in content else int(content['num_pts']),
+                   detection_name=content['detection_name'],
+                   detection_score=-1.0 if 'detection_score' not in content else float(content['detection_score']),
+                   attribute_name=content['attribute_name'],
+                   traj=content['trajs'],)
+
+
+def load_prediction(result_path: str, max_boxes_per_sample: int, box_cls, verbose: bool = False) \
+        -> Tuple[EvalBoxes, Dict]:
+    """
+    Loads object predictions from file.
+    :param result_path: Path to the .json result file provided by the user.
+    :param max_boxes_per_sample: Maximim number of boxes allowed per sample.
+    :param box_cls: Type of box to load, e.g. DetectionBox or TrackingBox.
+    :param verbose: Whether to print messages to stdout.
+    :return: The deserialized results and meta data.
+    """
+
+    # Load from file and check that the format is correct.
+    # with open(result_path) as f:
+    #     data = json.load(f)
+    data = result_path
+    assert 'results' in data, 'Error: No field `results` in result file. Please note that the result format changed.' \
+                              'See https://www.nuscenes.org/object-detection for more information.'
+
+    # motion name mapping
+    for key in data['results'].keys():
+        for i in range(len(data['results'][key])):
+            cls_name = data['results'][key][i]['detection_name']
+            if cls_name in motion_name_mapping:
+                cls_name = motion_name_mapping[cls_name]
+            data['results'][key][i]['detection_name'] = cls_name
+    
+    # Deserialize results and get meta data.
+    all_results = EvalBoxes.deserialize(data['results'], box_cls)
+    meta = data['meta']
+    if verbose:
+        print("Loaded results from {}. Found detections for {} samples."
+              .format(result_path, len(all_results.sample_tokens)))
+
+    # Check that each sample has no more than x predicted boxes.
+    for sample_token in all_results.sample_tokens:
+        assert len(all_results.boxes[sample_token]) <= max_boxes_per_sample, \
+            "Error: Only <= %d boxes per sample allowed!" % max_boxes_per_sample
+
+    return all_results, meta
+
+
+def load_gt(nusc: NuScenes, eval_split: str, box_cls, verbose: bool = False, seconds: int = 12) -> EvalBoxes:
+    """
+    Loads ground truth boxes from DB.
+    :param nusc: A NuScenes instance.
+    :param eval_split: The evaluation split for which we load GT boxes.
+    :param box_cls: Type of box to load, e.g. DetectionBox or TrackingBox.
+    :param verbose: Whether to print messages to stdout.
+    :return: The GT boxes.
+    """
+    predict_helper = PredictHelper(nusc)
+    # Init.
+    if box_cls == MotionBox:
+        attribute_map = {a['token']: a['name'] for a in nusc.attribute}
+
+    if verbose:
+        print('Loading annotations for {} split from nuScenes version: {}'.format(eval_split, nusc.version))
+    # Read out all sample_tokens in DB.
+    sample_tokens_all = [s['token'] for s in nusc.sample]
+    assert len(sample_tokens_all) > 0, "Error: Database has no samples!"
+
+    # Only keep samples from this split.
+    splits = create_splits_scenes()
+
+    # Check compatibility of split with nusc_version.
+    version = nusc.version
+    if eval_split in {'train', 'val', 'train_detect', 'train_track'}:
+        assert version.endswith('trainval'), \
+            'Error: Requested split {} which is not compatible with NuScenes version {}'.format(eval_split, version)
+    elif eval_split in {'mini_train', 'mini_val'}:
+        assert version.endswith('mini'), \
+            'Error: Requested split {} which is not compatible with NuScenes version {}'.format(eval_split, version)
+    elif eval_split == 'test':
+        assert version.endswith('test'), \
+            'Error: Requested split {} which is not compatible with NuScenes version {}'.format(eval_split, version)
+    else:
+        raise ValueError('Error: Requested split {} which this function cannot map to the correct NuScenes version.'
+                         .format(eval_split))
+
+    if eval_split == 'test':
+        # Check that you aren't trying to cheat :).
+        assert len(nusc.sample_annotation) > 0, \
+            'Error: You are trying to evaluate on the test set but you do not have the annotations!'
+
+    sample_tokens = []
+    for sample_token in sample_tokens_all:
+        scene_token = nusc.get('sample', sample_token)['scene_token']
+        scene_record = nusc.get('scene', scene_token)
+        if scene_record['name'] in splits[eval_split]:
+            sample_tokens.append(sample_token)
+
+    all_annotations = EvalBoxes()
+
+    # Load annotations and filter predictions and annotations.
+    tracking_id_set = set()
+    for sample_token in tqdm.tqdm(sample_tokens, leave=verbose):
+
+        sample = nusc.get('sample', sample_token)
+        sample_annotation_tokens = sample['anns']
+
+        sample_boxes = []
+        for sample_annotation_token in sample_annotation_tokens:
+
+            sample_annotation = nusc.get('sample_annotation', sample_annotation_token)
+            if box_cls == MotionBox:
+                # Get label name in detection task and filter unused labels.
+                detection_name = category_to_detection_name(sample_annotation['category_name'])
+                # motion name mapping
+                if detection_name in motion_name_mapping:
+                    detection_name = motion_name_mapping[detection_name]
+
+                if detection_name is None:
+                    continue
+
+                # Get attribute_name.
+                attr_tokens = sample_annotation['attribute_tokens']
+                attr_count = len(attr_tokens)
+                if attr_count == 0:
+                    attribute_name = ''
+                elif attr_count == 1:
+                    attribute_name = attribute_map[attr_tokens[0]]
+                else:
+                    raise Exception('Error: GT annotations must not have more than one attribute!')
+
+                # get future trajs
+                instance_token = nusc.get('sample_annotation', sample_annotation['token'])['instance_token']
+                fut_traj_local = predict_helper.get_future_for_agent(
+                    instance_token, 
+                    sample_token, 
+                    seconds=seconds, 
+                    in_agent_frame=True
+                )
+                if fut_traj_local.shape[0] > 0:
+                    _, boxes, _ = nusc.get_sample_data(sample['data']['LIDAR_TOP'], selected_anntokens=[sample_annotation['token']])
+                    box = boxes[0]
+                    trans = box.center
+                    rot = Quaternion(matrix=box.rotation_matrix)
+                    fut_traj_scence_centric = convert_local_coords_to_global(fut_traj_local, trans, rot) 
+                else:
+                    fut_traj_scence_centric = np.zeros((0,))
+
+                sample_boxes.append(
+                    box_cls(
+                        sample_token=sample_token,
+                        translation=sample_annotation['translation'],
+                        size=sample_annotation['size'],
+                        rotation=sample_annotation['rotation'],
+                        velocity=nusc.box_velocity(sample_annotation['token'])[:2],
+                        num_pts=sample_annotation['num_lidar_pts'] + sample_annotation['num_radar_pts'],
+                        detection_name=detection_name,
+                        detection_score=-1.0,  # GT samples do not have a score.
+                        attribute_name=attribute_name,
+                        traj=fut_traj_scence_centric
+                    )
+                )
+            elif box_cls == TrackingBox:
+                # Use nuScenes token as tracking id.
+                tracking_id = sample_annotation['instance_token']
+                tracking_id_set.add(tracking_id)
+
+                # Get label name in detection task and filter unused labels.
+                # Import locally to avoid errors when motmetrics package is not installed.
+                from nuscenes.eval.tracking.utils import category_to_tracking_name
+                tracking_name = category_to_tracking_name(sample_annotation['category_name'])
+                if tracking_name is None:
+                    continue
+
+                sample_boxes.append(
+                    box_cls(
+                        sample_token=sample_token,
+                        translation=sample_annotation['translation'],
+                        size=sample_annotation['size'],
+                        rotation=sample_annotation['rotation'],
+                        velocity=nusc.box_velocity(sample_annotation['token'])[:2],
+                        num_pts=sample_annotation['num_lidar_pts'] + sample_annotation['num_radar_pts'],
+                        tracking_id=tracking_id,
+                        tracking_name=tracking_name,
+                        tracking_score=-1.0  # GT samples do not have a score.
+                    )
+                )
+            else:
+                raise NotImplementedError('Error: Invalid box_cls %s!' % box_cls)
+
+        all_annotations.add_boxes(sample_token, sample_boxes)
+
+    if verbose:
+        print("Loaded ground truth annotations for {} samples.".format(len(all_annotations.sample_tokens)))
+
+    return all_annotations
+
+
+def accumulate(gt_boxes: EvalBoxes,
+               pred_boxes: EvalBoxes,
+               class_name: str,
+               dist_fcn: Callable,
+               dist_th: float,
+               verbose: bool = False) -> DetectionMetricData:
+    """
+    Average Precision over predefined different recall thresholds for a single distance threshold.
+    The recall/conf thresholds and other raw metrics will be used in secondary metrics.
+    :param gt_boxes: Maps every sample_token to a list of its sample_annotations.
+    :param pred_boxes: Maps every sample_token to a list of its sample_results.
+    :param class_name: Class to compute AP on.
+    :param dist_fcn: Distance function used to match detections and ground truths.
+    :param dist_th: Distance threshold for a match.
+    :param verbose: If true, print debug messages.
+    :return: (average_prec, metrics). The average precision value and raw data for a number of metrics.
+    """
+    # ---------------------------------------------
+    # Organize input and initialize accumulators.
+    # ---------------------------------------------
+
+    # Count the positives.
+    npos = len([1 for gt_box in gt_boxes.all if gt_box.detection_name == class_name])
+    if verbose:
+        print("Found {} GT of class {} out of {} total across {} samples.".
+              format(npos, class_name, len(gt_boxes.all), len(gt_boxes.sample_tokens)))
+
+    # For missing classes in the GT, return a data structure corresponding to no predictions.
+    if npos == 0:
+        return DetectionMetricData.no_predictions(), 0
+
+    # Organize the predictions in a single list.
+    pred_boxes_list = [box for box in pred_boxes.all if box.detection_name == class_name]
+    pred_confs = [box.detection_score for box in pred_boxes_list]
+
+    if verbose:
+        print("Found {} PRED of class {} out of {} total across {} samples.".
+              format(len(pred_confs), class_name, len(pred_boxes.all), len(pred_boxes.sample_tokens)))
+
+    # Sort by confidence.
+    sortind = [i for (v, i) in sorted((v, i) for (i, v) in enumerate(pred_confs))][::-1]
+
+    # Do the actual matching.
+    tp = []  # Accumulator of true positives.
+    fp = []  # Accumulator of false positives.
+    conf = []  # Accumulator of confidences.
+    hit = 0 # Accumulator of matched and hit
+
+    # match_data holds the extra metrics we calculate for each match.
+    match_data = {'conf': [],
+                  'min_ade': [],
+                  'min_fde': [],
+                  'miss_rate': []}
+
+    # ---------------------------------------------
+    # Match and accumulate match data.
+    # ---------------------------------------------
+
+    taken = set()  # Initially no gt bounding box is matched.
+    for ind in sortind:
+        pred_box = pred_boxes_list[ind]
+        min_dist = np.inf
+        match_gt_idx = None
+
+        for gt_idx, gt_box in enumerate(gt_boxes[pred_box.sample_token]):
+
+            # Find closest match among ground truth boxes
+            if gt_box.detection_name == class_name and not (pred_box.sample_token, gt_idx) in taken:
+                this_distance = dist_fcn(gt_box, pred_box)
+                if this_distance < min_dist:
+                    min_dist = this_distance
+                    match_gt_idx = gt_idx
+
+        # If the closest match is close enough according to threshold we have a match!
+        is_match = min_dist < dist_th
+
+        if is_match:
+            taken.add((pred_box.sample_token, match_gt_idx))
+
+            #  Update tp, fp and confs.
+            tp.append(1)
+            fp.append(0)
+            conf.append(pred_box.detection_score)
+
+            # Since it is a match, update match data also.
+            gt_box_match = gt_boxes[pred_box.sample_token][match_gt_idx]
+
+            match_data['conf'].append(pred_box.detection_score)
+
+            minade, minfde, mr = prediction_metrics(gt_box_match, pred_box)
+            match_data['min_ade'].append(minade)
+            match_data['min_fde'].append(minfde)
+            match_data['miss_rate'].append(mr)
+
+            if minfde < 2.0:
+                hit += 1
+
+        else:
+            # No match. Mark this as a false positive.
+            tp.append(0)
+            fp.append(1)
+            conf.append(pred_box.detection_score)
+
+    # Check if we have any matches. If not, just return a "no predictions" array.
+    if len(match_data['min_ade']) == 0:
+        return MotionMetricData.no_predictions()
+
+    # Accumulate.
+    N_tp = np.sum(tp)
+    N_fp = np.sum(fp)
+    tp = np.cumsum(tp).astype(float)
+    fp = np.cumsum(fp).astype(float)
+    conf = np.array(conf)
+
+    # Calculate precision and recall.
+    prec = tp / (fp + tp)
+    rec = tp / float(npos)
+
+    rec_interp = np.linspace(0, 1, DetectionMetricData.nelem)  # 101 steps, from 0% to 100% recall.
+    prec = np.interp(rec_interp, rec, prec, right=0)
+    conf = np.interp(rec_interp, rec, conf, right=0)
+    rec = rec_interp
+
+    # ---------------------------------------------
+    # Re-sample the match-data to match, prec, recall and conf.
+    # ---------------------------------------------
+
+    for key in match_data.keys():
+        if key == "conf":
+            continue  # Confidence is used as reference to align with fp and tp. So skip in this step.
+
+        else:
+            # For each match_data, we first calculate the accumulated mean.
+            tmp = cummean(np.array(match_data[key]))
+
+            # Then interpolate based on the confidences. (Note reversing since np.interp needs increasing arrays)
+            match_data[key] = np.interp(conf[::-1], match_data['conf'][::-1], tmp[::-1])[::-1]
+
+    EPA = (hit - 0.5 * N_fp) / npos
+
+    ## match based on traj
+    traj_matched = 0
+    taken = set()  # Initially no gt bounding box is matched.
+    for ind in sortind:
+        pred_box = pred_boxes_list[ind]
+        min_dist = np.inf
+        match_gt_idx = None
+
+        for gt_idx, gt_box in enumerate(gt_boxes[pred_box.sample_token]):
+
+            # Find closest match among ground truth boxes
+            if gt_box.detection_name == class_name and not (pred_box.sample_token, gt_idx) in taken:
+                this_distance = dist_fcn(gt_box, pred_box)
+                if this_distance < min_dist:
+                    min_dist = this_distance
+                    match_gt_idx = gt_idx
+                    fde_distance = traj_fde(gt_box, pred_box, final_step=12)
+
+        # If the closest match is close enough according to threshold we have a match!
+        is_match = min_dist < dist_th and fde_distance < 2.0
+        if is_match:
+            taken.add((pred_box.sample_token, match_gt_idx))
+            traj_matched += 1
+    EPA_ = (traj_matched - 0.5 * N_fp) / npos  ## same as UniAD
+
+    # ---------------------------------------------
+    # Done. Instantiate MetricData and return
+    # ---------------------------------------------
+    return MotionMetricData(recall=rec,
+                               precision=prec,
+                               confidence=conf,
+                               min_ade_err=match_data['min_ade'],
+                               min_fde_err=match_data['min_fde'],
+                               miss_rate_err=match_data['miss_rate']), EPA, EPA_
+
+
+def prediction_metrics(gt_box_match, pred_box, miss_thresh=2):
+    gt_traj = np.array(gt_box_match.traj)
+    pred_traj = np.array(pred_box.traj)
+
+    valid_step = gt_traj.shape[0]
+    if valid_step <= 0:
+        return 0, 0, 0
+
+    pred_traj_valid = pred_traj[:, :valid_step, :]
+    dist = np.linalg.norm(pred_traj_valid - gt_traj[np.newaxis], axis=2)
+
+    minade = dist.mean(axis=1).min()
+    minfde = dist[:, -1].min()
+    mr = dist.max(axis=1).min() > miss_thresh
+
+    return minade, minfde, mr
+
+def traj_fde(gt_box, pred_box, final_step):
+    if gt_box.traj.shape[0] <= 0:
+        return np.inf
+    final_step = min(gt_box.traj.shape[0], final_step)
+    gt_final = gt_box.traj[None, final_step-1]
+    pred_final = np.array(pred_box.traj)[:,final_step-1,:]
+    err = gt_final - pred_final
+    err = np.sqrt(np.sum(np.square(gt_final - pred_final), axis=-1))
+    return np.min(err)
+
+
+class MotionMetricDataList(DetectionMetricDataList):
+    """ This stores a set of MetricData in a dict indexed by (name, match-distance). """
+    @classmethod
+    def deserialize(cls, content: dict):
+        mdl = cls()
+        for key, md in content.items():
+            name, distance = key.split(':')
+            mdl.set(name, float(distance), MotionMetricData.deserialize(md))
+        return mdl
+
+class MotionMetricData(DetectionMetricData):
+    """ This class holds accumulated and interpolated data required to calculate the detection metrics. """
+
+    nelem = 101
+
+    def __init__(self,
+                 recall: np.array,
+                 precision: np.array,
+                 confidence: np.array,
+                 min_ade_err: np.array,
+                 min_fde_err: np.array,
+                 miss_rate_err: np.array):
+
+        # Assert lengths.
+        assert len(recall) == self.nelem
+        assert len(precision) == self.nelem
+        assert len(confidence) == self.nelem
+        assert len(min_ade_err) == self.nelem
+        assert len(min_fde_err) == self.nelem
+        assert len(miss_rate_err) == self.nelem
+
+        # Assert ordering.
+        assert all(confidence == sorted(confidence, reverse=True))  # Confidences should be descending.
+        assert all(recall == sorted(recall))  # Recalls should be ascending.
+
+        # Set attributes explicitly to help IDEs figure out what is going on.
+        self.recall = recall
+        self.precision = precision
+        self.confidence = confidence
+        self.min_ade_err = min_ade_err
+        self.min_fde_err = min_fde_err
+        self.miss_rate_err = miss_rate_err
+
+    def __eq__(self, other):
+        eq = True
+        for key in self.serialize().keys():
+            eq = eq and np.array_equal(getattr(self, key), getattr(other, key))
+        return eq
+
+    @property
+    def max_recall_ind(self):
+        """ Returns index of max recall achieved. """
+
+        # Last instance of confidence > 0 is index of max achieved recall.
+        non_zero = np.nonzero(self.confidence)[0]
+        if len(non_zero) == 0:  # If there are no matches, all the confidence values will be zero.
+            max_recall_ind = 0
+        else:
+            max_recall_ind = non_zero[-1]
+
+        return max_recall_ind
+
+    @property
+    def max_recall(self):
+        """ Returns max recall achieved. """
+
+        return self.recall[self.max_recall_ind]
+
+    def serialize(self):
+        """ Serialize instance into json-friendly format. """
+        return {
+            'recall': self.recall.tolist(),
+            'precision': self.precision.tolist(),
+            'confidence': self.confidence.tolist(),
+            'min_ade_err': self.min_ade_err.tolist(),
+            'min_fde_err': self.min_fde_err.tolist(),
+            'miss_rate_err': self.miss_rate_err.tolist(),
+        }
+
+    @classmethod
+    def deserialize(cls, content: dict):
+        """ Initialize from serialized content. """
+        return cls(recall=np.array(content['recall']),
+                   precision=np.array(content['precision']),
+                   confidence=np.array(content['confidence']),
+                   min_ade_err=np.array(content['min_ade_err']),
+                   min_fde_err=np.array(content['min_fde_err']),
+                   miss_rate_err=np.array(content['miss_rate_err']))
+
+    @classmethod
+    def no_predictions(cls):
+        """ Returns a md instance corresponding to having no predictions. """
+        return cls(recall=np.linspace(0, 1, cls.nelem),
+                   precision=np.zeros(cls.nelem),
+                   confidence=np.zeros(cls.nelem),
+                   min_ade_err=np.ones(cls.nelem),
+                   min_fde_err=np.ones(cls.nelem),
+                   miss_rate_err=np.ones(cls.nelem))
+
+    @classmethod
+    def random_md(cls):
+        """ Returns an md instance corresponding to a random results. """
+        return cls(recall=np.linspace(0, 1, cls.nelem),
+                   precision=np.random.random(cls.nelem),
+                   confidence=np.linspace(0, 1, cls.nelem)[::-1],
+                   min_ade_err=np.random.random(cls.nelem),
+                   min_fde_err=np.random.random(cls.nelem),
+                   miss_rate_err=np.random.random(cls.nelem))
+

projects/mmdet3d_plugin/datasets/evaluation/planning/planning_eval.py

+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import numpy as np
+from shapely.geometry import Polygon
+
+from mmcv.utils import print_log
+from mmdet.datasets import build_dataset, build_dataloader
+
+from projects.mmdet3d_plugin.datasets.utils import box3d_to_corners
+
+
+def check_collision(ego_box, boxes):
+    '''
+        ego_box: tensor with shape [7], [x, y, z, w, l, h, yaw]
+        boxes: tensor with shape [N, 7]
+    '''
+    if  boxes.shape[0] == 0:
+        return False
+
+    # follow uniad, add a 0.5m offset
+    ego_box[0] += 0.5 * torch.cos(ego_box[6])
+    ego_box[1] += 0.5 * torch.sin(ego_box[6])
+    ego_corners_box = box3d_to_corners(ego_box.unsqueeze(0))[0, [0, 3, 7, 4], :2]
+    corners_box = box3d_to_corners(boxes)[:, [0, 3, 7, 4], :2]
+    ego_poly = Polygon([(point[0], point[1]) for point in ego_corners_box])
+    for i in range(len(corners_box)):
+        box_poly =  Polygon([(point[0], point[1]) for point in corners_box[i]])
+        collision = ego_poly.intersects(box_poly)
+        if collision:
+            return True
+
+    return False
+
+def get_yaw(traj):
+    start = traj[0]
+    end = traj[-1]
+    dist = torch.linalg.norm(end - start, dim=-1)
+    if dist < 0.5:
+        return traj.new_ones(traj.shape[0]) * np.pi / 2
+
+    zeros = traj.new_zeros((1, 2))
+    traj_cat = torch.cat([zeros, traj], dim=0)
+    yaw = traj.new_zeros(traj.shape[0]+1)
+    yaw[..., 1:-1] = torch.atan2(
+        traj_cat[..., 2:, 1] - traj_cat[..., :-2, 1],
+        traj_cat[..., 2:, 0] - traj_cat[..., :-2, 0],
+    )
+    yaw[..., -1] = torch.atan2(
+        traj_cat[..., -1, 1] - traj_cat[..., -2, 1],
+        traj_cat[..., -1, 0] - traj_cat[..., -2, 0],
+    )
+    return yaw[1:]
+
+class PlanningMetric():
+    def __init__(
+        self,
+        n_future=6,
+        compute_on_step: bool = False,
+    ):
+        self.W = 1.85
+        self.H = 4.084
+
+        self.n_future = n_future
+        self.reset()
+
+    def reset(self):
+        self.obj_col = torch.zeros(self.n_future)
+        self.obj_box_col = torch.zeros(self.n_future)
+        self.L2 = torch.zeros(self.n_future)
+        self.total = torch.tensor(0)
+
+    def evaluate_single_coll(self, traj, fut_boxes):
+        n_future = traj.shape[0]
+        yaw = get_yaw(traj)
+        ego_box = traj.new_zeros((n_future, 7))
+        ego_box[:, :2] = traj
+        ego_box[:, 3:6] = ego_box.new_tensor([self.H, self.W, 1.56])
+        ego_box[:, 6] = yaw
+        collision = torch.zeros(n_future, dtype=torch.bool)
+
+        for t in range(n_future):
+            ego_box_t = ego_box[t].clone()
+            boxes = fut_boxes[t][0].clone()
+            collision[t] = check_collision(ego_box_t, boxes)
+        return collision
+
+    def evaluate_coll(self, trajs, gt_trajs, fut_boxes):
+        B, n_future, _ = trajs.shape
+        trajs = trajs * torch.tensor([-1, 1], device=trajs.device)
+        gt_trajs = gt_trajs * torch.tensor([-1, 1], device=gt_trajs.device)
+
+        obj_coll_sum = torch.zeros(n_future, device=trajs.device)
+        obj_box_coll_sum = torch.zeros(n_future, device=trajs.device)
+
+        assert B == 1, 'only supprt bs=1'
+        for i in range(B):
+            gt_box_coll = self.evaluate_single_coll(gt_trajs[i], fut_boxes)
+            box_coll = self.evaluate_single_coll(trajs[i], fut_boxes)
+            box_coll = torch.logical_and(box_coll, torch.logical_not(gt_box_coll))
+            
+            obj_coll_sum += gt_box_coll.long()
+            obj_box_coll_sum += box_coll.long()
+
+        return obj_coll_sum, obj_box_coll_sum
+
+    def compute_L2(self, trajs, gt_trajs, gt_trajs_mask):
+        '''
+        trajs: torch.Tensor (B, n_future, 3)
+        gt_trajs: torch.Tensor (B, n_future, 3)
+        '''
+        return torch.sqrt((((trajs[:, :, :2] - gt_trajs[:, :, :2]) ** 2) * gt_trajs_mask).sum(dim=-1)) 
+
+    def update(self, trajs, gt_trajs, gt_trajs_mask, fut_boxes):
+        assert trajs.shape == gt_trajs.shape
+        trajs[..., 0] = - trajs[..., 0]
+        gt_trajs[..., 0] = - gt_trajs[..., 0]
+        L2 = self.compute_L2(trajs, gt_trajs, gt_trajs_mask)
+        obj_coll_sum, obj_box_coll_sum = self.evaluate_coll(trajs[:,:,:2], gt_trajs[:,:,:2], fut_boxes)
+
+        self.obj_col += obj_coll_sum
+        self.obj_box_col += obj_box_coll_sum
+        self.L2 += L2.sum(dim=0)
+        self.total +=len(trajs)
+
+    def compute(self):
+        return {
+            'obj_col': self.obj_col / self.total,
+            'obj_box_col': self.obj_box_col / self.total,
+            'L2' : self.L2 / self.total
+        }
+
+
+def planning_eval(results, eval_config, logger):
+    dataset = build_dataset(eval_config)
+    dataloader = build_dataloader(
+            dataset, samples_per_gpu=1, workers_per_gpu=1, shuffle=False, dist=False)
+    planning_metrics = PlanningMetric()
+    for i, data in enumerate(tqdm(dataloader)):
+        sdc_planning = data['gt_ego_fut_trajs'].cumsum(dim=-2).unsqueeze(1)
+        sdc_planning_mask = data['gt_ego_fut_masks'].unsqueeze(-1).repeat(1, 1, 2).unsqueeze(1)
+        command = data['gt_ego_fut_cmd'].argmax(dim=-1).item()
+        fut_boxes = data['fut_boxes']
+        if not sdc_planning_mask.all(): ## for incomplete gt, we do not count this sample
+            continue
+        res = results[i]
+        pred_sdc_traj = res['img_bbox']['final_planning'].unsqueeze(0)
+        planning_metrics.update(pred_sdc_traj[:, :6, :2], sdc_planning[0,:, :6, :2], sdc_planning_mask[0,:, :6, :2], fut_boxes)
+       
+    planning_results = planning_metrics.compute()
+    planning_metrics.reset()
+    from prettytable import PrettyTable
+    planning_tab = PrettyTable()
+    metric_dict = {}
+
+    planning_tab.field_names = [
+    "metrics", "0.5s", "1.0s", "1.5s", "2.0s", "2.5s", "3.0s", "avg"]
+    for key in planning_results.keys():
+        value = planning_results[key].tolist()
+        new_values = []
+        for i in range(len(value)):
+            new_values.append(np.array(value[:i+1]).mean())
+        value = new_values
+        avg = [value[1], value[3], value[5]]
+        avg = sum(avg) / len(avg)
+        value.append(avg)
+        metric_dict[key] = avg
+        row_value = []
+        row_value.append(key)
+        for i in range(len(value)):
+            if 'col' in key:
+                row_value.append('%.3f' % float(value[i]*100) + '%')
+            else:
+                row_value.append('%.4f' % float(value[i]))
+        planning_tab.add_row(row_value)
+
+    print_log('\n'+str(planning_tab), logger=logger)
+    return metric_dict

projects/mmdet3d_plugin/datasets/map_utils/nuscmap_extractor.py

+from shapely.geometry import LineString, box, Polygon
+from shapely import ops, strtree
+
+import numpy as np
+from nuscenes.map_expansion.map_api import NuScenesMap, NuScenesMapExplorer
+from nuscenes.eval.common.utils import quaternion_yaw
+from pyquaternion import Quaternion
+from .utils import split_collections, get_drivable_area_contour, \
+        get_ped_crossing_contour
+from numpy.typing import NDArray
+from typing import Dict, List, Tuple, Union
+
+class NuscMapExtractor(object):
+    """NuScenes map ground-truth extractor.
+
+    Args:
+        data_root (str): path to nuScenes dataset
+        roi_size (tuple or list): bev range
+    """
+    def __init__(self, data_root: str, roi_size: Union[List, Tuple]) -> None:
+        self.roi_size = roi_size
+        self.MAPS = ['boston-seaport', 'singapore-hollandvillage',
+                     'singapore-onenorth', 'singapore-queenstown']
+        
+        self.nusc_maps = {}
+        self.map_explorer = {}
+        for loc in self.MAPS:
+            self.nusc_maps[loc] = NuScenesMap(
+                dataroot=data_root, map_name=loc)
+            self.map_explorer[loc] = NuScenesMapExplorer(self.nusc_maps[loc])
+        
+        # local patch in nuScenes format
+        self.local_patch = box(-roi_size[0] / 2, -roi_size[1] / 2, 
+                roi_size[0] / 2, roi_size[1] / 2)
+    
+    def _union_ped(self, ped_geoms: List[Polygon]) -> List[Polygon]:
+        ''' merge close ped crossings.
+        
+        Args:
+            ped_geoms (list): list of Polygon
+        
+        Returns:
+            union_ped_geoms (Dict): merged ped crossings 
+        '''
+
+        def get_rec_direction(geom):
+            rect = geom.minimum_rotated_rectangle
+            rect_v_p = np.array(rect.exterior.coords)[:3]
+            rect_v = rect_v_p[1:]-rect_v_p[:-1]
+            v_len = np.linalg.norm(rect_v, axis=-1)
+            longest_v_i = v_len.argmax()
+
+            return rect_v[longest_v_i], v_len[longest_v_i]
+
+        tree = strtree.STRtree(ped_geoms)
+        index_by_id = dict((id(pt), i) for i, pt in enumerate(ped_geoms))
+
+        final_pgeom = []
+        remain_idx = [i for i in range(len(ped_geoms))]
+        for i, pgeom in enumerate(ped_geoms):
+
+            if i not in remain_idx:
+                continue
+            # update
+            remain_idx.pop(remain_idx.index(i))
+            pgeom_v, pgeom_v_norm = get_rec_direction(pgeom)
+            final_pgeom.append(pgeom)
+
+            for o in tree.query(pgeom):
+                o_idx = index_by_id[id(o)]
+                if o_idx not in remain_idx:
+                    continue
+
+                o_v, o_v_norm = get_rec_direction(o)
+                cos = pgeom_v.dot(o_v)/(pgeom_v_norm*o_v_norm)
+                if 1 - np.abs(cos) < 0.01:  # theta < 8 degrees.
+                    final_pgeom[-1] =\
+                        final_pgeom[-1].union(o)
+                    # update
+                    remain_idx.pop(remain_idx.index(o_idx))
+
+        results = []
+        for p in final_pgeom:
+            results.extend(split_collections(p))
+        return results
+        
+    def get_map_geom(self, 
+                     location: str, 
+                     translation: Union[List, NDArray],
+                     rotation: Union[List, NDArray]) -> Dict[str, List[Union[LineString, Polygon]]]:
+        ''' Extract geometries given `location` and self pose, self may be lidar or ego.
+        
+        Args:
+            location (str): city name
+            translation (array): self2global translation, shape (3,)
+            rotation (array): self2global quaternion, shape (4, )
+            
+        Returns:
+            geometries (Dict): extracted geometries by category.
+        '''
+
+        # (center_x, center_y, len_y, len_x) in nuscenes format
+        patch_box = (translation[0], translation[1], 
+                self.roi_size[1], self.roi_size[0])
+        rotation = Quaternion(rotation)
+        yaw = quaternion_yaw(rotation) / np.pi * 180
+
+        # get dividers
+        lane_dividers = self.map_explorer[location]._get_layer_line(
+                    patch_box, yaw, 'lane_divider')
+        
+        road_dividers = self.map_explorer[location]._get_layer_line(
+                    patch_box, yaw, 'road_divider')
+        
+        all_dividers = []
+        for line in lane_dividers + road_dividers:
+            all_dividers += split_collections(line)
+
+        # get ped crossings
+        ped_crossings = []
+        ped = self.map_explorer[location]._get_layer_polygon(
+                    patch_box, yaw, 'ped_crossing')
+        
+        for p in ped:
+            ped_crossings += split_collections(p)
+        # some ped crossings are split into several small parts
+        # we need to merge them
+        ped_crossings = self._union_ped(ped_crossings)
+        
+        ped_crossing_lines = []
+        for p in ped_crossings:
+            # extract exteriors to get a closed polyline
+            line = get_ped_crossing_contour(p, self.local_patch)
+            if line is not None:
+                ped_crossing_lines.append(line)
+
+        # get boundaries
+        # we take the union of road segments and lanes as drivable areas
+        # we don't take drivable area layer in nuScenes since its definition may be ambiguous
+        road_segments = self.map_explorer[location]._get_layer_polygon(
+                    patch_box, yaw, 'road_segment')
+        lanes = self.map_explorer[location]._get_layer_polygon(
+                    patch_box, yaw, 'lane')
+        union_roads = ops.unary_union(road_segments)
+        union_lanes = ops.unary_union(lanes)
+        drivable_areas = ops.unary_union([union_roads, union_lanes])
+        
+        drivable_areas = split_collections(drivable_areas)
+        
+        # boundaries are defined as the contour of drivable areas
+        boundaries = get_drivable_area_contour(drivable_areas, self.roi_size)
+
+        return dict(
+            divider=all_dividers, # List[LineString]
+            ped_crossing=ped_crossing_lines, # List[LineString]
+            boundary=boundaries, # List[LineString]
+            drivable_area=drivable_areas, # List[Polygon],
+        )
+

projects/mmdet3d_plugin/datasets/map_utils/utils.py

+from shapely.geometry import LineString, box, Polygon, LinearRing
+from shapely.geometry.base import BaseGeometry
+from shapely import ops
+import numpy as np
+from scipy.spatial import distance
+from typing import List, Optional, Tuple
+from numpy.typing import NDArray
+
+def split_collections(geom: BaseGeometry) -> List[Optional[BaseGeometry]]:
+    ''' Split Multi-geoms to list and check is valid or is empty.
+        
+    Args:
+        geom (BaseGeometry): geoms to be split or validate.
+    
+    Returns:
+        geometries (List): list of geometries.
+    '''
+    assert geom.geom_type in ['MultiLineString', 'LineString', 'MultiPolygon', 
+        'Polygon', 'GeometryCollection'], f"got geom type {geom.geom_type}"
+    if 'Multi' in geom.geom_type:
+        outs = []
+        for g in geom.geoms:
+            if g.is_valid and not g.is_empty:
+                outs.append(g)
+        return outs
+    else:
+        if geom.is_valid and not geom.is_empty:
+            return [geom,]
+        else:
+            return []
+
+def get_drivable_area_contour(drivable_areas: List[Polygon], 
+                              roi_size: Tuple) -> List[LineString]:
+    ''' Extract drivable area contours to get list of boundaries.
+
+    Args:
+        drivable_areas (list): list of drivable areas.
+        roi_size (tuple): bev range size
+    
+    Returns:
+        boundaries (List): list of boundaries.
+    '''
+    max_x = roi_size[0] / 2
+    max_y = roi_size[1] / 2
+
+    # a bit smaller than roi to avoid unexpected boundaries on edges
+    local_patch = box(-max_x + 0.2, -max_y + 0.2, max_x - 0.2, max_y - 0.2)
+    
+    exteriors = []
+    interiors = []
+    
+    for poly in drivable_areas:
+        exteriors.append(poly.exterior)
+        for inter in poly.interiors:
+            interiors.append(inter)
+    
+    results = []
+    for ext in exteriors:
+        # NOTE: we make sure all exteriors are clock-wise
+        # such that each boundary's right-hand-side is drivable area
+        # and left-hand-side is walk way
+        
+        if ext.is_ccw:
+            ext = LinearRing(list(ext.coords)[::-1])
+        lines = ext.intersection(local_patch)
+        if lines.geom_type == 'MultiLineString':
+            lines = ops.linemerge(lines)
+        assert lines.geom_type in ['MultiLineString', 'LineString']
+        
+        results.extend(split_collections(lines))
+
+    for inter in interiors:
+        # NOTE: we make sure all interiors are counter-clock-wise
+        if not inter.is_ccw:
+            inter = LinearRing(list(inter.coords)[::-1])
+        lines = inter.intersection(local_patch)
+        if lines.geom_type == 'MultiLineString':
+            lines = ops.linemerge(lines)
+        assert lines.geom_type in ['MultiLineString', 'LineString']
+        
+        results.extend(split_collections(lines))
+
+    return results
+
+def get_ped_crossing_contour(polygon: Polygon, 
+                             local_patch: box) -> Optional[LineString]:
+    ''' Extract ped crossing contours to get a closed polyline.
+    Different from `get_drivable_area_contour`, this function ensures a closed polyline.
+
+    Args:
+        polygon (Polygon): ped crossing polygon to be extracted.
+        local_patch (tuple): local patch params
+    
+    Returns:
+        line (LineString): a closed line
+    '''
+
+    ext = polygon.exterior
+    if not ext.is_ccw:
+        ext = LinearRing(list(ext.coords)[::-1])
+    lines = ext.intersection(local_patch)
+    if lines.type != 'LineString':
+        # remove points in intersection results
+        lines = [l for l in lines.geoms if l.geom_type != 'Point']
+        lines = ops.linemerge(lines)
+        
+        # same instance but not connected.
+        if lines.type != 'LineString':
+            ls = []
+            for l in lines.geoms:
+                ls.append(np.array(l.coords))
+            
+            lines = np.concatenate(ls, axis=0)
+            lines = LineString(lines)
+    if not lines.is_empty:
+        return lines
+    
+    return None
+

projects/mmdet3d_plugin/datasets/nuscenes_3d_dataset.py

+import random
+import math
+import os
+from os import path as osp
+import cv2
+import tempfile
+import copy
+import prettytable
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import pyquaternion
+from shapely.geometry import LineString
+from nuscenes.utils.data_classes import Box as NuScenesBox
+from nuscenes.eval.detection.config import config_factory as det_configs
+from nuscenes.eval.common.config import config_factory as track_configs
+
+import mmcv
+from mmcv.utils import print_log
+from mmdet.datasets import DATASETS
+from mmdet.datasets.pipelines import Compose
+from .utils import (
+    draw_lidar_bbox3d_on_img,
+    draw_lidar_bbox3d_on_bev,
+)
+
+
+@DATASETS.register_module()
+class NuScenes3DDataset(Dataset):
+    DefaultAttribute = {
+        "car": "vehicle.parked",
+        "pedestrian": "pedestrian.moving",
+        "trailer": "vehicle.parked",
+        "truck": "vehicle.parked",
+        "bus": "vehicle.moving",
+        "motorcycle": "cycle.without_rider",
+        "construction_vehicle": "vehicle.parked",
+        "bicycle": "cycle.without_rider",
+        "barrier": "",
+        "traffic_cone": "",
+    }
+    ErrNameMapping = {
+        "trans_err": "mATE",
+        "scale_err": "mASE",
+        "orient_err": "mAOE",
+        "vel_err": "mAVE",
+        "attr_err": "mAAE",
+    }
+    CLASSES = (
+        "car",
+        "truck",
+        "trailer",
+        "bus",
+        "construction_vehicle",
+        "bicycle",
+        "motorcycle",
+        "pedestrian",
+        "traffic_cone",
+        "barrier",
+    )
+    MAP_CLASSES = (
+        'ped_crossing',
+        'divider',
+        'boundary',
+    )
+    ID_COLOR_MAP = [
+        (59, 59, 238),
+        (0, 255, 0),
+        (0, 0, 255),
+        (255, 255, 0),
+        (0, 255, 255),
+        (255, 0, 255),
+        (255, 255, 255),
+        (0, 127, 255),
+        (71, 130, 255),
+        (127, 127, 0),
+    ]
+
+    def __init__(
+        self,
+        ann_file,
+        pipeline=None,
+        data_root=None,
+        classes=None,
+        map_classes=None,
+        load_interval=1,
+        with_velocity=True,
+        modality=None,
+        test_mode=False,
+        det3d_eval_version="detection_cvpr_2019",
+        track3d_eval_version="tracking_nips_2019",
+        version="v1.0-trainval",
+        use_valid_flag=False,
+        vis_score_threshold=0.25,
+        data_aug_conf=None,
+        sequences_split_num=1,
+        with_seq_flag=False,
+        keep_consistent_seq_aug=True,
+        work_dir=None,
+        eval_config=None,
+    ):
+        self.version = version
+        self.load_interval = load_interval
+        self.use_valid_flag = use_valid_flag
+        super().__init__()
+        self.data_root = data_root
+        self.ann_file = ann_file
+        self.test_mode = test_mode
+        self.modality = modality
+        self.box_mode_3d = 0
+
+        if classes is not None:
+            self.CLASSES = classes
+        if map_classes is not None: 
+            self.MAP_CLASSES = map_classes
+        self.cat2id = {name: i for i, name in enumerate(self.CLASSES)}
+        self.data_infos = self.load_annotations(self.ann_file)
+
+        if pipeline is not None:
+            self.pipeline = Compose(pipeline)
+
+        self.with_velocity = with_velocity
+        self.det3d_eval_version = det3d_eval_version
+        self.det3d_eval_configs = det_configs(self.det3d_eval_version)
+        self.det3d_eval_configs.class_names = list(self.det3d_eval_configs.class_range.keys())
+        self.track3d_eval_version = track3d_eval_version
+        self.track3d_eval_configs = track_configs(self.track3d_eval_version)
+        self.track3d_eval_configs.class_names = list(self.track3d_eval_configs.class_range.keys())
+        if self.modality is None:
+            self.modality = dict(
+                use_camera=False,
+                use_lidar=True,
+                use_radar=False,
+                use_map=False,
+                use_external=False,
+            )
+        self.vis_score_threshold = vis_score_threshold
+
+        self.data_aug_conf = data_aug_conf
+        self.sequences_split_num = sequences_split_num
+        self.keep_consistent_seq_aug = keep_consistent_seq_aug
+        if with_seq_flag:
+            self._set_sequence_group_flag()
+        
+        self.work_dir = work_dir
+        self.eval_config = eval_config
+
+    def __len__(self):
+        return len(self.data_infos)
+
+    def _set_sequence_group_flag(self):
+        """
+        Set each sequence to be a different group
+        """
+        if self.sequences_split_num == -1:
+            self.flag = np.arange(len(self.data_infos))
+            return
+        
+        res = []
+
+        curr_sequence = 0
+        for idx in range(len(self.data_infos)):
+            if idx != 0 and len(self.data_infos[idx]["sweeps"]) == 0:
+                # Not first frame and # of sweeps is 0 -> new sequence
+                curr_sequence += 1
+            res.append(curr_sequence)
+
+        self.flag = np.array(res, dtype=np.int64)
+
+        if self.sequences_split_num != 1:
+            if self.sequences_split_num == "all":
+                self.flag = np.array(
+                    range(len(self.data_infos)), dtype=np.int64
+                )
+            else:
+                bin_counts = np.bincount(self.flag)
+                new_flags = []
+                curr_new_flag = 0
+                for curr_flag in range(len(bin_counts)):
+                    curr_sequence_length = np.array(
+                        list(
+                            range(
+                                0,
+                                bin_counts[curr_flag],
+                                math.ceil(
+                                    bin_counts[curr_flag]
+                                    / self.sequences_split_num
+                                ),
+                            )
+                        )
+                        + [bin_counts[curr_flag]]
+                    )
+
+                    for sub_seq_idx in (
+                        curr_sequence_length[1:] - curr_sequence_length[:-1]
+                    ):
+                        for _ in range(sub_seq_idx):
+                            new_flags.append(curr_new_flag)
+                        curr_new_flag += 1
+
+                assert len(new_flags) == len(self.flag)
+                assert (
+                    len(np.bincount(new_flags))
+                    == len(np.bincount(self.flag)) * self.sequences_split_num
+                )
+                self.flag = np.array(new_flags, dtype=np.int64)
+
+    def get_augmentation(self):
+        if self.data_aug_conf is None:
+            return None
+        H, W = self.data_aug_conf["H"], self.data_aug_conf["W"]
+        fH, fW = self.data_aug_conf["final_dim"]
+        if not self.test_mode:
+            resize = np.random.uniform(*self.data_aug_conf["resize_lim"])
+            resize_dims = (int(W * resize), int(H * resize))
+            newW, newH = resize_dims
+            crop_h = (
+                int(
+                    (1 - np.random.uniform(*self.data_aug_conf["bot_pct_lim"]))
+                    * newH
+                )
+                - fH
+            )
+            crop_w = int(np.random.uniform(0, max(0, newW - fW)))
+            crop = (crop_w, crop_h, crop_w + fW, crop_h + fH)
+            flip = False
+            if self.data_aug_conf["rand_flip"] and np.random.choice([0, 1]):
+                flip = True
+            rotate = np.random.uniform(*self.data_aug_conf["rot_lim"])
+            rotate_3d = np.random.uniform(*self.data_aug_conf["rot3d_range"])
+        else:
+            resize = max(fH / H, fW / W)
+            resize_dims = (int(W * resize), int(H * resize))
+            newW, newH = resize_dims
+            crop_h = (
+                int((1 - np.mean(self.data_aug_conf["bot_pct_lim"])) * newH)
+                - fH
+            )
+            crop_w = int(max(0, newW - fW) / 2)
+            crop = (crop_w, crop_h, crop_w + fW, crop_h + fH)
+            flip = False
+            rotate = 0
+            rotate_3d = 0
+        aug_config = {
+            "resize": resize,
+            "resize_dims": resize_dims,
+            "crop": crop,
+            "flip": flip,
+            "rotate": rotate,
+            "rotate_3d": rotate_3d,
+        }
+        return aug_config
+
+    def __getitem__(self, idx):
+        if isinstance(idx, dict):
+            aug_config = idx["aug_config"]
+            idx = idx["idx"]
+        else:
+            aug_config = self.get_augmentation()
+        data = self.get_data_info(idx)
+        data["aug_config"] = aug_config
+        data = self.pipeline(data)
+        return data
+
+    def get_cat_ids(self, idx):
+        info = self.data_infos[idx]
+        if self.use_valid_flag:
+            mask = info["valid_flag"]
+            gt_names = set(info["gt_names"][mask])
+        else:
+            gt_names = set(info["gt_names"])
+
+        cat_ids = []
+        for name in gt_names:
+            if name in self.CLASSES:
+                cat_ids.append(self.cat2id[name])
+        return cat_ids
+
+    def load_annotations(self, ann_file):
+        data = mmcv.load(ann_file, file_format="pkl")
+        data_infos = list(sorted(data["infos"], key=lambda e: e["timestamp"]))
+        data_infos = data_infos[:: self.load_interval]
+        self.metadata = data["metadata"]
+        self.version = self.metadata["version"]
+        print(self.metadata)
+        return data_infos
+    
+    def anno2geom(self, annos):
+        map_geoms = {}
+        for label, anno_list in annos.items():
+            map_geoms[label] = []
+            for anno in anno_list:
+                geom = LineString(anno)
+                map_geoms[label].append(geom)
+        return map_geoms
+    
+    def get_data_info(self, index):
+        info = self.data_infos[index]
+        input_dict = dict(
+            token=info["token"],
+            map_location=info["map_location"],
+            pts_filename=info["lidar_path"],
+            sweeps=info["sweeps"],
+            timestamp=info["timestamp"] / 1e6,
+            lidar2ego_translation=info["lidar2ego_translation"],
+            lidar2ego_rotation=info["lidar2ego_rotation"],
+            ego2global_translation=info["ego2global_translation"],
+            ego2global_rotation=info["ego2global_rotation"],
+            ego_status=info['ego_status'].astype(np.float32),
+            map_infos=info["map_annos"],
+        )
+        lidar2ego = np.eye(4)
+        lidar2ego[:3, :3] = pyquaternion.Quaternion(
+            info["lidar2ego_rotation"]
+        ).rotation_matrix
+        lidar2ego[:3, 3] = np.array(info["lidar2ego_translation"])
+        ego2global = np.eye(4)
+        ego2global[:3, :3] = pyquaternion.Quaternion(
+            info["ego2global_rotation"]
+        ).rotation_matrix
+        ego2global[:3, 3] = np.array(info["ego2global_translation"])
+        input_dict["lidar2global"] = ego2global @ lidar2ego
+
+        map_geoms = self.anno2geom(info["map_annos"])
+        input_dict["map_geoms"] = map_geoms
+
+        if self.modality["use_camera"]:
+            image_paths = []
+            lidar2img_rts = []
+            lidar2cam_rts = []
+            cam_intrinsic = []
+            for cam_type, cam_info in info["cams"].items():
+                image_paths.append(cam_info["data_path"])
+                # obtain lidar to image transformation matrix
+                lidar2cam_r = np.linalg.inv(cam_info["sensor2lidar_rotation"])
+                lidar2cam_t = (
+                    cam_info["sensor2lidar_translation"] @ lidar2cam_r.T
+                )
+                lidar2cam_rt = np.eye(4)
+                lidar2cam_rt[:3, :3] = lidar2cam_r.T
+                lidar2cam_rt[3, :3] = -lidar2cam_t
+                intrinsic = copy.deepcopy(cam_info["cam_intrinsic"])
+                cam_intrinsic.append(intrinsic)
+                viewpad = np.eye(4)
+                viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic
+                lidar2img_rt = viewpad @ lidar2cam_rt.T
+                lidar2img_rts.append(lidar2img_rt)
+                lidar2cam_rts.append(lidar2cam_rt)
+
+            input_dict.update(
+                dict(
+                    img_filename=image_paths,
+                    lidar2img=lidar2img_rts,
+                    lidar2cam=lidar2cam_rts,
+                    cam_intrinsic=cam_intrinsic,
+                )
+            )
+
+        annos = self.get_ann_info(index)
+        input_dict.update(annos)
+        return input_dict
+
+    def get_ann_info(self, index):
+        info = self.data_infos[index]
+        if self.use_valid_flag:
+            mask = info["valid_flag"]
+        else:
+            mask = info["num_lidar_pts"] > 0
+        gt_bboxes_3d = info["gt_boxes"][mask]
+        gt_names_3d = info["gt_names"][mask]
+        gt_labels_3d = []
+        for cat in gt_names_3d:
+            if cat in self.CLASSES:
+                gt_labels_3d.append(self.CLASSES.index(cat))
+            else:
+                gt_labels_3d.append(-1)
+        gt_labels_3d = np.array(gt_labels_3d)
+
+        if self.with_velocity:
+            gt_velocity = info["gt_velocity"][mask]
+            nan_mask = np.isnan(gt_velocity[:, 0])
+            gt_velocity[nan_mask] = [0.0, 0.0]
+            gt_bboxes_3d = np.concatenate([gt_bboxes_3d, gt_velocity], axis=-1)
+
+        anns_results = dict(
+            gt_bboxes_3d=gt_bboxes_3d,
+            gt_labels_3d=gt_labels_3d,
+            gt_names=gt_names_3d,
+        )
+        if "instance_inds" in info:
+            instance_inds = np.array(info["instance_inds"], dtype=np.int)[mask]
+            anns_results["instance_inds"] = instance_inds
+            
+        if 'gt_agent_fut_trajs' in info:
+            anns_results['gt_agent_fut_trajs'] = info['gt_agent_fut_trajs'][mask]
+            anns_results['gt_agent_fut_masks'] = info['gt_agent_fut_masks'][mask]
+
+        if 'gt_ego_fut_trajs' in info:
+            anns_results['gt_ego_fut_trajs'] = info['gt_ego_fut_trajs']
+            anns_results['gt_ego_fut_masks'] = info['gt_ego_fut_masks']
+            anns_results['gt_ego_fut_cmd'] = info['gt_ego_fut_cmd']
+        
+            ## get future box for planning eval
+            fut_ts = int(info['gt_ego_fut_masks'].sum())
+            fut_boxes = []
+            cur_scene_token = info["scene_token"]
+            cur_T_global = get_T_global(info)
+            for i in range(1, fut_ts + 1):
+                fut_info = self.data_infos[index + i]
+                fut_scene_token = fut_info["scene_token"]
+                if cur_scene_token != fut_scene_token:
+                    break
+                if self.use_valid_flag:
+                    mask = fut_info["valid_flag"]
+                else:
+                    mask = fut_info["num_lidar_pts"] > 0
+
+                fut_gt_bboxes_3d = fut_info["gt_boxes"][mask]
+                
+                fut_T_global = get_T_global(fut_info)
+                T_fut2cur = np.linalg.inv(cur_T_global) @ fut_T_global
+
+                center = fut_gt_bboxes_3d[:, :3] @ T_fut2cur[:3, :3].T + T_fut2cur[:3, 3]
+                yaw = np.stack([np.cos(fut_gt_bboxes_3d[:, 6]), np.sin(fut_gt_bboxes_3d[:, 6])], axis=-1)
+                yaw = yaw @ T_fut2cur[:2, :2].T
+                yaw = np.arctan2(yaw[..., 1], yaw[..., 0])
+
+                fut_gt_bboxes_3d[:, :3] = center
+                fut_gt_bboxes_3d[:, 6] = yaw
+                fut_boxes.append(fut_gt_bboxes_3d)
+
+            anns_results['fut_boxes'] = fut_boxes
+        
+        return anns_results
+
+    def _format_bbox(self, results, jsonfile_prefix=None, tracking=False):
+        nusc_annos = {}
+        mapped_class_names = self.CLASSES
+
+        print("Start to convert detection format...")
+        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
+            annos = []
+            boxes = output_to_nusc_box(
+                det, threshold=self.tracking_threshold if tracking else None
+            )
+            sample_token = self.data_infos[sample_id]["token"]
+            boxes = lidar_nusc_box_to_global(
+                self.data_infos[sample_id],
+                boxes,
+                mapped_class_names,
+                self.det3d_eval_configs,
+                self.det3d_eval_version,
+            )
+            for i, box in enumerate(boxes):
+                name = mapped_class_names[box.label]
+                if tracking and name in [
+                    "barrier",
+                    "traffic_cone",
+                    "construction_vehicle",
+                ]:
+                    continue
+                if np.sqrt(box.velocity[0] ** 2 + box.velocity[1] ** 2) > 0.2:
+                    if name in [
+                        "car",
+                        "construction_vehicle",
+                        "bus",
+                        "truck",
+                        "trailer",
+                    ]:
+                        attr = "vehicle.moving"
+                    elif name in ["bicycle", "motorcycle"]:
+                        attr = "cycle.with_rider"
+                    else:
+                        attr = NuScenes3DDataset.DefaultAttribute[name]
+                else:
+                    if name in ["pedestrian"]:
+                        attr = "pedestrian.standing"
+                    elif name in ["bus"]:
+                        attr = "vehicle.stopped"
+                    else:
+                        attr = NuScenes3DDataset.DefaultAttribute[name]
+
+                nusc_anno = dict(
+                    sample_token=sample_token,
+                    translation=box.center.tolist(),
+                    size=box.wlh.tolist(),
+                    rotation=box.orientation.elements.tolist(),
+                    velocity=box.velocity[:2].tolist(),
+                )
+                if not tracking:
+                    nusc_anno.update(
+                        dict(
+                            detection_name=name,
+                            detection_score=box.score,
+                            attribute_name=attr,
+                        )
+                    )
+                else:
+                    nusc_anno.update(
+                        dict(
+                            tracking_name=name,
+                            tracking_score=box.score,
+                            tracking_id=str(box.token),
+                        )
+                    )
+
+                annos.append(nusc_anno)
+            nusc_annos[sample_token] = annos
+        nusc_submissions = {
+            "meta": self.modality,
+            "results": nusc_annos,
+        }
+
+        mmcv.mkdir_or_exist(jsonfile_prefix)
+        res_path = osp.join(jsonfile_prefix, "results_nusc.json")
+        print("Results writes to", res_path)
+        mmcv.dump(nusc_submissions, res_path)
+        return res_path
+
+    def _evaluate_single(
+        self, result_path, logger=None, result_name="img_bbox", tracking=False
+    ):
+        from nuscenes import NuScenes
+
+        output_dir = osp.join(*osp.split(result_path)[:-1])
+        nusc = NuScenes(
+            version=self.version, dataroot=self.data_root, verbose=False
+        )
+        eval_set_map = {
+            "v1.0-mini": "mini_val",
+            "v1.0-trainval": "val",
+        }
+        if not tracking:
+            from nuscenes.eval.detection.evaluate import NuScenesEval
+
+            nusc_eval = NuScenesEval(
+                nusc,
+                config=self.det3d_eval_configs,
+                result_path=result_path,
+                eval_set=eval_set_map[self.version],
+                output_dir=output_dir,
+                verbose=True,
+            )
+            nusc_eval.main(render_curves=False)
+
+            # record metrics
+            metrics = mmcv.load(osp.join(output_dir, "metrics_summary.json"))
+            detail = dict()
+            metric_prefix = f"{result_name}_NuScenes"
+            for name in self.CLASSES:
+                for k, v in metrics["label_aps"][name].items():
+                    val = float("{:.4f}".format(v))
+                    detail[
+                        "{}/{}_AP_dist_{}".format(metric_prefix, name, k)
+                    ] = val
+                for k, v in metrics["label_tp_errors"][name].items():
+                    val = float("{:.4f}".format(v))
+                    detail["{}/{}_{}".format(metric_prefix, name, k)] = val
+                for k, v in metrics["tp_errors"].items():
+                    val = float("{:.4f}".format(v))
+                    detail[
+                        "{}/{}".format(metric_prefix, self.ErrNameMapping[k])
+                    ] = val
+
+            detail["{}/NDS".format(metric_prefix)] = metrics["nd_score"]
+            detail["{}/mAP".format(metric_prefix)] = metrics["mean_ap"]
+        else:
+            from nuscenes.eval.tracking.evaluate import TrackingEval
+
+            nusc_eval = TrackingEval(
+                config=self.track3d_eval_configs,
+                result_path=result_path,
+                eval_set=eval_set_map[self.version],
+                output_dir=output_dir,
+                verbose=True,
+                nusc_version=self.version,
+                nusc_dataroot=self.data_root,
+            )
+            metrics = nusc_eval.main()
+
+            # record metrics
+            metrics = mmcv.load(osp.join(output_dir, "metrics_summary.json"))
+            print(metrics)
+            detail = dict()
+            metric_prefix = f"{result_name}_NuScenes"
+            keys = [
+                "amota",
+                "amotp",
+                "recall",
+                "motar",
+                "gt",
+                "mota",
+                "motp",
+                "mt",
+                "ml",
+                "faf",
+                "tp",
+                "fp",
+                "fn",
+                "ids",
+                "frag",
+                "tid",
+                "lgd",
+            ]
+            for key in keys:
+                detail["{}/{}".format(metric_prefix, key)] = metrics[key]
+
+        return detail
+
+    def format_results(self, results, jsonfile_prefix=None, tracking=False):
+        assert isinstance(results, list), "results must be a list"
+
+        if jsonfile_prefix is None:
+            tmp_dir = tempfile.TemporaryDirectory()
+            jsonfile_prefix = osp.join(tmp_dir.name, "results")
+        else:
+            tmp_dir = None
+
+        if not ("pts_bbox" in results[0] or "img_bbox" in results[0]):
+            result_files = self._format_bbox(
+                results, jsonfile_prefix, tracking=tracking
+            )
+        else:
+            result_files = dict()
+            for name in results[0]:
+                print(f"\nFormating bboxes of {name}")
+                results_ = [out[name] for out in results]
+                tmp_file_ = jsonfile_prefix
+                result_files.update(
+                    {
+                        name: self._format_bbox(
+                            results_, tmp_file_, tracking=tracking
+                        )
+                    }
+                )
+        return result_files, tmp_dir
+
+    def format_map_results(self, results, prefix=None):
+        submissions = {'results': {},}
+        
+        for j, pred in enumerate(results):
+            '''
+            For each case, the result should be formatted as Dict{'vectors': [], 'scores': [], 'labels': []}
+            'vectors': List of vector, each vector is a array([[x1, y1], [x2, y2] ...]),
+                contain all vectors predicted in this sample.
+            'scores: List of score(float), 
+                contain scores of all instances in this sample.
+            'labels': List of label(int), 
+                contain labels of all instances in this sample.
+            '''
+            if pred is None: # empty prediction
+                continue
+            pred = pred['img_bbox']
+
+            single_case = {'vectors': [], 'scores': [], 'labels': []}
+            token = self.data_infos[j]['token']
+            for i in range(len(pred['scores'])):
+                score = pred['scores'][i]
+                label = pred['labels'][i]
+                vector = pred['vectors'][i]
+
+                # A line should have >=2 points
+                if len(vector) < 2:
+                    continue
+                
+                single_case['vectors'].append(vector)
+                single_case['scores'].append(score)
+                single_case['labels'].append(label)
+            
+            submissions['results'][token] = single_case
+        
+        out_path = osp.join(prefix, 'submission_vector.json')
+        print(f'saving submissions results to {out_path}')
+        os.makedirs(os.path.dirname(out_path), exist_ok=True)
+        mmcv.dump(submissions, out_path)
+        return out_path
+
+    def format_motion_results(self, results, jsonfile_prefix=None, tracking=False, thresh=None):
+        nusc_annos = {}
+        mapped_class_names = self.CLASSES
+
+        print("Start to convert detection format...")
+        for sample_id, det in enumerate(mmcv.track_iter_progress(results)):
+            annos = []
+            boxes = output_to_nusc_box(
+                det['img_bbox'], threshold=None
+            )
+            sample_token = self.data_infos[sample_id]["token"]
+            boxes = lidar_nusc_box_to_global(
+                self.data_infos[sample_id],
+                boxes,
+                mapped_class_names,
+                self.det3d_eval_configs,
+                self.det3d_eval_version,
+                filter_with_cls_range=False,
+            )
+            for i, box in enumerate(boxes):
+                if thresh is not None and box.score < thresh:
+                    continue
+                name = mapped_class_names[box.label]
+                if tracking and name in [
+                    "barrier",
+                    "traffic_cone",
+                    "construction_vehicle",
+                ]:
+                    continue
+                if np.sqrt(box.velocity[0] ** 2 + box.velocity[1] ** 2) > 0.2:
+                    if name in [
+                        "car",
+                        "construction_vehicle",
+                        "bus",
+                        "truck",
+                        "trailer",
+                    ]:
+                        attr = "vehicle.moving"
+                    elif name in ["bicycle", "motorcycle"]:
+                        attr = "cycle.with_rider"
+                    else:
+                        attr = NuScenes3DDataset.DefaultAttribute[name]
+                else:
+                    if name in ["pedestrian"]:
+                        attr = "pedestrian.standing"
+                    elif name in ["bus"]:
+                        attr = "vehicle.stopped"
+                    else:
+                        attr = NuScenes3DDataset.DefaultAttribute[name]
+
+                nusc_anno = dict(
+                    sample_token=sample_token,
+                    translation=box.center.tolist(),
+                    size=box.wlh.tolist(),
+                    rotation=box.orientation.elements.tolist(),
+                    velocity=box.velocity[:2].tolist(),
+                )
+                if not tracking:
+                    nusc_anno.update(
+                        dict(
+                            detection_name=name,
+                            detection_score=box.score,
+                            attribute_name=attr,
+                        )
+                    )
+                else:
+                    nusc_anno.update(
+                        dict(
+                            tracking_name=name,
+                            tracking_score=box.score,
+                            tracking_id=str(box.token),
+                        )
+                    )
+                nusc_anno.update(
+                    dict(
+                        trajs=det['img_bbox']['trajs_3d'][i].numpy(),
+                    )
+                )
+                annos.append(nusc_anno)
+            nusc_annos[sample_token] = annos
+        nusc_submissions = {
+            "meta": self.modality,
+            "results": nusc_annos,
+        }
+
+        return nusc_submissions 
+
+    def _evaluate_single_motion(self,
+                         results,
+                         result_path,
+                         logger=None,
+                         metric='bbox',
+                         result_name='pts_bbox'):
+        """Evaluation for a single model in nuScenes protocol.
+
+        Args:
+            result_path (str): Path of the result file.
+            logger (logging.Logger | str | None): Logger used for printing
+                related information during evaluation. Default: None.
+            metric (str): Metric name used for evaluation. Default: 'bbox'.
+            result_name (str): Result name in the metric prefix.
+                Default: 'pts_bbox'.
+
+        Returns:
+            dict: Dictionary of evaluation details.
+        """
+        from nuscenes import NuScenes
+        from .evaluation.motion.motion_eval_uniad import NuScenesEval as NuScenesEvalMotion
+
+        output_dir = result_path
+        nusc = NuScenes(
+            version=self.version, dataroot=self.data_root, verbose=False)
+        eval_set_map = {
+            'v1.0-mini': 'mini_val',
+            'v1.0-trainval': 'val',
+        }
+        nusc_eval = NuScenesEvalMotion(
+            nusc,
+            config=copy.deepcopy(self.det3d_eval_configs),
+            result_path=results,
+            eval_set=eval_set_map[self.version],
+            output_dir=output_dir,
+            verbose=False,
+            seconds=6)
+        metrics = nusc_eval.main(render_curves=False)
+        
+        MOTION_METRICS = ['EPA', 'min_ade_err', 'min_fde_err', 'miss_rate_err']
+        class_names = ['car', 'pedestrian']
+
+        table = prettytable.PrettyTable()
+        table.field_names = ["class names"] + MOTION_METRICS
+        for class_name in class_names:
+            row_data = [class_name]
+            for m in MOTION_METRICS:
+                row_data.append('%.4f' % metrics[f'{class_name}_{m}'])
+            table.add_row(row_data)
+        print_log('\n'+str(table), logger=logger)
+        return metrics
+
+    def evaluate(
+        self,
+        results,
+        eval_mode,
+        metric=None,
+        logger=None,
+        jsonfile_prefix=None,
+        result_names=["img_bbox"],
+        show=False,
+        out_dir=None,
+        pipeline=None,
+    ):
+        res_path = "results.pkl" if "trainval" in self.version else "results_mini.pkl"
+        res_path = osp.join(self.work_dir, res_path)
+        print('All Results write to', res_path)
+        mmcv.dump(results, res_path)
+
+        results_dict = dict()
+        if eval_mode['with_det']:
+            self.tracking = eval_mode["with_tracking"]
+            self.tracking_threshold = eval_mode["tracking_threshold"]
+            for metric in ["detection", "tracking"]:
+                tracking = metric == "tracking"
+                if tracking and not self.tracking:
+                    continue
+                result_files, tmp_dir = self.format_results(
+                    results, jsonfile_prefix=self.work_dir, tracking=tracking
+                )
+
+                if isinstance(result_files, dict):
+                    for name in result_names:
+                        ret_dict = self._evaluate_single(
+                            result_files[name], tracking=tracking
+                        )
+                    results_dict.update(ret_dict)
+                elif isinstance(result_files, str):
+                    ret_dict = self._evaluate_single(
+                        result_files, tracking=tracking
+                    )
+                    results_dict.update(ret_dict)
+
+                if tmp_dir is not None:
+                    tmp_dir.cleanup()
+
+        if eval_mode['with_map']:
+            from .evaluation.map.vector_eval import VectorEvaluate
+            self.map_evaluator = VectorEvaluate(self.eval_config)
+            result_path = self.format_map_results(results, prefix=self.work_dir)
+            map_results_dict = self.map_evaluator.evaluate(result_path, logger=logger)
+            results_dict.update(map_results_dict)
+
+        if eval_mode['with_motion']:
+            thresh = eval_mode["motion_threshhold"]
+            result_files = self.format_motion_results(results, jsonfile_prefix=self.work_dir, thresh=thresh)
+            motion_results_dict = self._evaluate_single_motion(result_files, self.work_dir, logger=logger)
+            results_dict.update(motion_results_dict)
+        
+        if eval_mode['with_planning']:
+            from .evaluation.planning.planning_eval import planning_eval
+            planning_results_dict = planning_eval(results, self.eval_config, logger=logger)
+            results_dict.update(planning_results_dict)
+
+        if show or out_dir:
+            self.show(results, save_dir=out_dir, show=show, pipeline=pipeline)
+        
+        # print main metrics for recording
+        metric_str = '\n'
+        if "img_bbox_NuScenes/NDS" in results_dict:
+            metric_str += f'mAP: {results_dict.get("img_bbox_NuScenes/mAP"):.4f}\n'
+            metric_str += f'mATE: {results_dict.get("img_bbox_NuScenes/mATE"):.4f}\n'
+            metric_str += f'mASE: {results_dict.get("img_bbox_NuScenes/mASE"):.4f}\n'
+            metric_str += f'mAOE: {results_dict.get("img_bbox_NuScenes/mAOE"):.4f}\n' 
+            metric_str += f'mAVE: {results_dict.get("img_bbox_NuScenes/mAVE"):.4f}\n' 
+            metric_str += f'mAAE: {results_dict.get("img_bbox_NuScenes/mAAE"):.4f}\n' 
+            metric_str += f'NDS: {results_dict.get("img_bbox_NuScenes/NDS"):.4f}\n\n'
+        
+        if "img_bbox_NuScenes/amota" in results_dict:
+            metric_str += f'AMOTA: {results_dict["img_bbox_NuScenes/amota"]:.4f}\n' 
+            metric_str += f'AMOTP: {results_dict["img_bbox_NuScenes/amotp"]:.4f}\n' 
+            metric_str += f'RECALL: {results_dict["img_bbox_NuScenes/recall"]:.4f}\n' 
+            metric_str += f'MOTAR: {results_dict["img_bbox_NuScenes/motar"]:.4f}\n' 
+            metric_str += f'MOTA: {results_dict["img_bbox_NuScenes/mota"]:.4f}\n' 
+            metric_str += f'MOTP: {results_dict["img_bbox_NuScenes/motp"]:.4f}\n' 
+            metric_str += f'IDS: {results_dict["img_bbox_NuScenes/ids"]}\n\n' 
+        
+        if "mAP_normal" in results_dict:
+            metric_str += f'ped_crossing= {results_dict["ped_crossing"]:.4f}\n' 
+            metric_str += f'divider= {results_dict["divider"]:.4f}\n' 
+            metric_str += f'boundary= {results_dict["boundary"]:.4f}\n' 
+            metric_str += f'mAP_normal= {results_dict["mAP_normal"]:.4f}\n\n' 
+
+        if "car_EPA" in results_dict:
+            metric_str += f'Car / Ped\n' 
+            metric_str += f'epa= {results_dict["car_EPA"]:.4f} / {results_dict["pedestrian_EPA"]:.4f}\n'
+            metric_str += f'ade= {results_dict["car_min_ade_err"]:.4f} / {results_dict["pedestrian_min_ade_err"]:.4f}\n'
+            metric_str += f'fde= {results_dict["car_min_fde_err"]:.4f} / {results_dict["pedestrian_min_fde_err"]:.4f}\n'
+            metric_str += f'mr= {results_dict["car_miss_rate_err"]:.4f} / {results_dict["pedestrian_miss_rate_err"]:.4f}\n\n' 
+
+        if "L2" in results_dict:
+            metric_str += f'obj_box_col: {(results_dict["obj_box_col"]*100):.3f}%\n'
+            metric_str += f'L2: {results_dict["L2"]:.4f}\n\n'
+        
+        print_log(metric_str, logger=logger)
+        return results_dict
+
+    def show(self, results, save_dir=None, show=False, pipeline=None):
+        save_dir = "./" if save_dir is None else save_dir
+        save_dir = os.path.join(save_dir, "visual")
+        print_log(os.path.abspath(save_dir))
+        pipeline = Compose(pipeline)
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+
+        fourcc = cv2.VideoWriter_fourcc(*"MJPG")
+        videoWriter = None
+
+        for i, result in enumerate(results):
+            if "img_bbox" in result.keys():
+                result = result["img_bbox"]
+            data_info = pipeline(self.get_data_info(i))
+            imgs = []
+
+            raw_imgs = data_info["img"]
+            lidar2img = data_info["img_metas"].data["lidar2img"]
+            pred_bboxes_3d = result["boxes_3d"][
+                result["scores_3d"] > self.vis_score_threshold
+            ]
+            if "instance_ids" in result and self.tracking:
+                color = []
+                for id in result["instance_ids"].cpu().numpy().tolist():
+                    color.append(
+                        self.ID_COLOR_MAP[int(id % len(self.ID_COLOR_MAP))]
+                    )
+            elif "labels_3d" in result:
+                color = []
+                for id in result["labels_3d"].cpu().numpy().tolist():
+                    color.append(self.ID_COLOR_MAP[id])
+            else:
+                color = (255, 0, 0)
+
+            # ===== draw boxes_3d to images =====
+            for j, img_origin in enumerate(raw_imgs):
+                img = img_origin.copy()
+                if len(pred_bboxes_3d) != 0:
+                    img = draw_lidar_bbox3d_on_img(
+                        pred_bboxes_3d,
+                        img,
+                        lidar2img[j],
+                        img_metas=None,
+                        color=color,
+                        thickness=3,
+                    )
+                imgs.append(img)
+
+            # ===== draw boxes_3d to BEV =====
+            bev = draw_lidar_bbox3d_on_bev(
+                pred_bboxes_3d,
+                bev_size=img.shape[0] * 2,
+                color=color,
+            )
+
+            # ===== put text and concat =====
+            for j, name in enumerate(
+                [
+                    "front",
+                    "front right",
+                    "front left",
+                    "rear",
+                    "rear left",
+                    "rear right",
+                ]
+            ):
+                imgs[j] = cv2.rectangle(
+                    imgs[j],
+                    (0, 0),
+                    (440, 80),
+                    color=(255, 255, 255),
+                    thickness=-1,
+                )
+                w, h = cv2.getTextSize(name, cv2.FONT_HERSHEY_SIMPLEX, 2, 2)[0]
+                text_x = int(220 - w / 2)
+                text_y = int(40 + h / 2)
+
+                imgs[j] = cv2.putText(
+                    imgs[j],
+                    name,
+                    (text_x, text_y),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    2,
+                    (0, 0, 0),
+                    2,
+                    cv2.LINE_AA,
+                )
+            image = np.concatenate(
+                [
+                    np.concatenate([imgs[2], imgs[0], imgs[1]], axis=1),
+                    np.concatenate([imgs[5], imgs[3], imgs[4]], axis=1),
+                ],
+                axis=0,
+            )
+            image = np.concatenate([image, bev], axis=1)
+
+            # ===== save video =====
+            if videoWriter is None:
+                videoWriter = cv2.VideoWriter(
+                    os.path.join(save_dir, "video.avi"),
+                    fourcc,
+                    7,
+                    image.shape[:2][::-1],
+                )
+            cv2.imwrite(os.path.join(save_dir, f"{i}.jpg"), image)
+            videoWriter.write(image)
+        videoWriter.release()
+
+
+def output_to_nusc_box(detection, threshold=None):
+    box3d = detection["boxes_3d"]
+    scores = detection["scores_3d"].numpy()
+    labels = detection["labels_3d"].numpy()
+    if "instance_ids" in detection:
+        ids = detection["instance_ids"]  # .numpy()
+    if threshold is not None:
+        if "cls_scores" in detection:
+            mask = detection["cls_scores"].numpy() >= threshold
+        else:
+            mask = scores >= threshold
+        box3d = box3d[mask]
+        scores = scores[mask]
+        labels = labels[mask]
+        ids = ids[mask]
+
+    if hasattr(box3d, "gravity_center"):
+        box_gravity_center = box3d.gravity_center.numpy()
+        box_dims = box3d.dims.numpy()
+        nus_box_dims = box_dims[:, [1, 0, 2]]
+        box_yaw = box3d.yaw.numpy()
+    else:
+        box3d = box3d.numpy()
+        box_gravity_center = box3d[..., :3].copy()
+        box_dims = box3d[..., 3:6].copy()
+        nus_box_dims = box_dims[..., [1, 0, 2]]
+        box_yaw = box3d[..., 6].copy()
+
+    # TODO: check whether this is necessary
+    # with dir_offset & dir_limit in the head
+    # box_yaw = -box_yaw - np.pi / 2
+
+    box_list = []
+    for i in range(len(box3d)):
+        quat = pyquaternion.Quaternion(axis=[0, 0, 1], radians=box_yaw[i])
+        if hasattr(box3d, "gravity_center"):
+            velocity = (*box3d.tensor[i, 7:9], 0.0)
+        else:
+            velocity = (*box3d[i, 7:9], 0.0)
+        box = NuScenesBox(
+            box_gravity_center[i],
+            nus_box_dims[i],
+            quat,
+            label=labels[i],
+            score=scores[i],
+            velocity=velocity,
+        )
+        if "instance_ids" in detection:
+            box.token = ids[i]
+        box_list.append(box)
+    return box_list
+
+
+def lidar_nusc_box_to_global(
+    info,
+    boxes,
+    classes,
+    eval_configs,
+    eval_version="detection_cvpr_2019",
+    filter_with_cls_range=True,
+):
+    box_list = []
+    for i, box in enumerate(boxes):
+        # Move box to ego vehicle coord system
+        box.rotate(pyquaternion.Quaternion(info["lidar2ego_rotation"]))
+        box.translate(np.array(info["lidar2ego_translation"]))
+        # filter det in ego.
+        if filter_with_cls_range:
+            cls_range_map = eval_configs.class_range
+            radius = np.linalg.norm(box.center[:2], 2)
+            det_range = cls_range_map[classes[box.label]]
+            if radius > det_range:
+                continue
+        # Move box to global coord system
+        box.rotate(pyquaternion.Quaternion(info["ego2global_rotation"]))
+        box.translate(np.array(info["ego2global_translation"]))
+        box_list.append(box)
+    return box_list
+
+
+def get_T_global(info):
+    lidar2ego = np.eye(4)
+    lidar2ego[:3, :3] = pyquaternion.Quaternion(
+        info["lidar2ego_rotation"]
+    ).rotation_matrix
+    lidar2ego[:3, 3] = np.array(info["lidar2ego_translation"])
+    ego2global = np.eye(4)
+    ego2global[:3, :3] = pyquaternion.Quaternion(
+        info["ego2global_rotation"]
+    ).rotation_matrix
+    ego2global[:3, 3] = np.array(info["ego2global_translation"])
+    return ego2global @ lidar2ego
\ No newline at end of file

projects/mmdet3d_plugin/datasets/pipelines/__init__.py

+from .transform import (
+    InstanceNameFilter,
+    CircleObjectRangeFilter,
+    NormalizeMultiviewImage,
+    NuScenesSparse4DAdaptor,
+    MultiScaleDepthMapGenerator,
+)
+from .augment import (
+    ResizeCropFlipImage,
+    BBoxRotation,
+    PhotoMetricDistortionMultiViewImage,
+)
+from .loading import LoadMultiViewImageFromFiles, LoadPointsFromFile
+from .vectorize import VectorizeMap
+
+__all__ = [
+    "InstanceNameFilter",
+    "ResizeCropFlipImage",
+    "BBoxRotation",
+    "CircleObjectRangeFilter",
+    "MultiScaleDepthMapGenerator",
+    "NormalizeMultiviewImage",
+    "PhotoMetricDistortionMultiViewImage",
+    "NuScenesSparse4DAdaptor",
+    "LoadMultiViewImageFromFiles",
+    "LoadPointsFromFile",
+    "VectorizeMap",
+]

projects/mmdet3d_plugin/datasets/pipelines/augment.py

+import torch
+
+import numpy as np
+from numpy import random
+import mmcv
+from mmdet.datasets.builder import PIPELINES
+from PIL import Image
+
+
+@PIPELINES.register_module()
+class ResizeCropFlipImage(object):
+    def __call__(self, results):
+        aug_config = results.get("aug_config")
+        if aug_config is None:
+            return results
+        imgs = results["img"]
+        N = len(imgs)
+        new_imgs = []
+        for i in range(N):
+            img, mat = self._img_transform(
+                np.uint8(imgs[i]), aug_config,
+            )
+            new_imgs.append(np.array(img).astype(np.float32))
+            results["lidar2img"][i] = mat @ results["lidar2img"][i]
+            if "cam_intrinsic" in results:
+                results["cam_intrinsic"][i][:3, :3] *= aug_config["resize"]
+                # results["cam_intrinsic"][i][:3, :3] = (
+                #     mat[:3, :3] @ results["cam_intrinsic"][i][:3, :3]
+                # )
+
+        results["img"] = new_imgs
+        results["img_shape"] = [x.shape[:2] for x in new_imgs]
+        return results
+
+    def _img_transform(self, img, aug_configs):
+        H, W = img.shape[:2]
+        resize = aug_configs.get("resize", 1)
+        resize_dims = (int(W * resize), int(H * resize))
+        crop = aug_configs.get("crop", [0, 0, *resize_dims])
+        flip = aug_configs.get("flip", False)
+        rotate = aug_configs.get("rotate", 0)
+
+        origin_dtype = img.dtype
+        if origin_dtype != np.uint8:
+            min_value = img.min()
+            max_vaule = img.max()
+            scale = 255 / (max_vaule - min_value)
+            img = (img - min_value) * scale
+            img = np.uint8(img)
+        img = Image.fromarray(img)
+        img = img.resize(resize_dims).crop(crop)
+        if flip:
+            img = img.transpose(method=Image.FLIP_LEFT_RIGHT)
+        img = img.rotate(rotate)
+        img = np.array(img).astype(np.float32)
+        if origin_dtype != np.uint8:
+            img = img.astype(np.float32)
+            img = img / scale + min_value
+
+        transform_matrix = np.eye(3)
+        transform_matrix[:2, :2] *= resize
+        transform_matrix[:2, 2] -= np.array(crop[:2])
+        if flip:
+            flip_matrix = np.array(
+                [[-1, 0, crop[2] - crop[0]], [0, 1, 0], [0, 0, 1]]
+            )
+            transform_matrix = flip_matrix @ transform_matrix
+        rotate = rotate / 180 * np.pi
+        rot_matrix = np.array(
+            [
+                [np.cos(rotate), np.sin(rotate), 0],
+                [-np.sin(rotate), np.cos(rotate), 0],
+                [0, 0, 1],
+            ]
+        )
+        rot_center = np.array([crop[2] - crop[0], crop[3] - crop[1]]) / 2
+        rot_matrix[:2, 2] = -rot_matrix[:2, :2] @ rot_center + rot_center
+        transform_matrix = rot_matrix @ transform_matrix
+        extend_matrix = np.eye(4)
+        extend_matrix[:3, :3] = transform_matrix
+        return img, extend_matrix
+
+
+@PIPELINES.register_module()
+class BBoxRotation(object):
+    def __call__(self, results):
+        angle = results["aug_config"]["rotate_3d"]
+        rot_cos = np.cos(angle)
+        rot_sin = np.sin(angle)
+
+        rot_mat = np.array(
+            [
+                [rot_cos, -rot_sin, 0, 0],
+                [rot_sin, rot_cos, 0, 0],
+                [0, 0, 1, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+        rot_mat_inv = np.linalg.inv(rot_mat)
+
+        num_view = len(results["lidar2img"])
+        for view in range(num_view):
+            results["lidar2img"][view] = (
+                results["lidar2img"][view] @ rot_mat_inv
+            )
+        if "lidar2global" in results:
+            results["lidar2global"] = results["lidar2global"] @ rot_mat_inv
+        if "gt_bboxes_3d" in results:
+            results["gt_bboxes_3d"] = self.box_rotate(
+                results["gt_bboxes_3d"], angle
+            )
+        return results
+
+    @staticmethod
+    def box_rotate(bbox_3d, angle):
+        rot_cos = np.cos(angle)
+        rot_sin = np.sin(angle)
+        rot_mat_T = np.array(
+            [[rot_cos, rot_sin, 0], [-rot_sin, rot_cos, 0], [0, 0, 1]]
+        )
+        bbox_3d[:, :3] = bbox_3d[:, :3] @ rot_mat_T
+        bbox_3d[:, 6] += angle
+        if bbox_3d.shape[-1] > 7:
+            vel_dims = bbox_3d[:, 7:].shape[-1]
+            bbox_3d[:, 7:] = bbox_3d[:, 7:] @ rot_mat_T[:vel_dims, :vel_dims]
+        return bbox_3d
+
+
+@PIPELINES.register_module()
+class PhotoMetricDistortionMultiViewImage:
+    """Apply photometric distortion to image sequentially, every transformation
+    is applied with a probability of 0.5. The position of random contrast is in
+    second or second to last.
+    1. random brightness
+    2. random contrast (mode 0)
+    3. convert color from BGR to HSV
+    4. random saturation
+    5. random hue
+    6. convert color from HSV to BGR
+    7. random contrast (mode 1)
+    8. randomly swap channels
+    Args:
+        brightness_delta (int): delta of brightness.
+        contrast_range (tuple): range of contrast.
+        saturation_range (tuple): range of saturation.
+        hue_delta (int): delta of hue.
+    """
+
+    def __init__(
+        self,
+        brightness_delta=32,
+        contrast_range=(0.5, 1.5),
+        saturation_range=(0.5, 1.5),
+        hue_delta=18,
+    ):
+        self.brightness_delta = brightness_delta
+        self.contrast_lower, self.contrast_upper = contrast_range
+        self.saturation_lower, self.saturation_upper = saturation_range
+        self.hue_delta = hue_delta
+
+    def __call__(self, results):
+        """Call function to perform photometric distortion on images.
+        Args:
+            results (dict): Result dict from loading pipeline.
+        Returns:
+            dict: Result dict with images distorted.
+        """
+        imgs = results["img"]
+        new_imgs = []
+        for img in imgs:
+            assert img.dtype == np.float32, (
+                "PhotoMetricDistortion needs the input image of dtype np.float32,"
+                ' please set "to_float32=True" in "LoadImageFromFile" pipeline'
+            )
+            # random brightness
+            if random.randint(2):
+                delta = random.uniform(
+                    -self.brightness_delta, self.brightness_delta
+                )
+                img += delta
+
+            # mode == 0 --> do random contrast first
+            # mode == 1 --> do random contrast last
+            mode = random.randint(2)
+            if mode == 1:
+                if random.randint(2):
+                    alpha = random.uniform(
+                        self.contrast_lower, self.contrast_upper
+                    )
+                    img *= alpha
+
+            # convert color from BGR to HSV
+            img = mmcv.bgr2hsv(img)
+
+            # random saturation
+            if random.randint(2):
+                img[..., 1] *= random.uniform(
+                    self.saturation_lower, self.saturation_upper
+                )
+
+            # random hue
+            if random.randint(2):
+                img[..., 0] += random.uniform(-self.hue_delta, self.hue_delta)
+                img[..., 0][img[..., 0] > 360] -= 360
+                img[..., 0][img[..., 0] < 0] += 360
+
+            # convert color from HSV to BGR
+            img = mmcv.hsv2bgr(img)
+
+            # random contrast
+            if mode == 0:
+                if random.randint(2):
+                    alpha = random.uniform(
+                        self.contrast_lower, self.contrast_upper
+                    )
+                    img *= alpha
+
+            # randomly swap channels
+            if random.randint(2):
+                img = img[..., random.permutation(3)]
+            new_imgs.append(img)
+        results["img"] = new_imgs
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f"(\nbrightness_delta={self.brightness_delta},\n"
+        repr_str += "contrast_range="
+        repr_str += f"{(self.contrast_lower, self.contrast_upper)},\n"
+        repr_str += "saturation_range="
+        repr_str += f"{(self.saturation_lower, self.saturation_upper)},\n"
+        repr_str += f"hue_delta={self.hue_delta})"
+        return repr_str

projects/mmdet3d_plugin/datasets/pipelines/loading.py

+import numpy as np
+import mmcv
+from mmdet.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class LoadMultiViewImageFromFiles(object):
+    """Load multi channel images from a list of separate channel files.
+
+    Expects results['img_filename'] to be a list of filenames.
+
+    Args:
+        to_float32 (bool, optional): Whether to convert the img to float32.
+            Defaults to False.
+        color_type (str, optional): Color type of the file.
+            Defaults to 'unchanged'.
+    """
+
+    def __init__(self, to_float32=False, color_type="unchanged"):
+        self.to_float32 = to_float32
+        self.color_type = color_type
+
+    def __call__(self, results):
+        """Call function to load multi-view image from files.
+
+        Args:
+            results (dict): Result dict containing multi-view image filenames.
+
+        Returns:
+            dict: The result dict containing the multi-view image data.
+                Added keys and values are described below.
+
+                - filename (str): Multi-view image filenames.
+                - img (np.ndarray): Multi-view image arrays.
+                - img_shape (tuple[int]): Shape of multi-view image arrays.
+                - ori_shape (tuple[int]): Shape of original image arrays.
+                - pad_shape (tuple[int]): Shape of padded image arrays.
+                - scale_factor (float): Scale factor.
+                - img_norm_cfg (dict): Normalization configuration of images.
+        """
+        filename = results["img_filename"]
+        # img is of shape (h, w, c, num_views)
+        img = np.stack(
+            [mmcv.imread(name, self.color_type) for name in filename], axis=-1
+        )
+        if self.to_float32:
+            img = img.astype(np.float32)
+        results["filename"] = filename
+        # unravel to list, see `DefaultFormatBundle` in formatting.py
+        # which will transpose each image separately and then stack into array
+        results["img"] = [img[..., i] for i in range(img.shape[-1])]
+        results["img_shape"] = img.shape
+        results["ori_shape"] = img.shape
+        # Set initial values for default meta_keys
+        results["pad_shape"] = img.shape
+        results["scale_factor"] = 1.0
+        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
+        results["img_norm_cfg"] = dict(
+            mean=np.zeros(num_channels, dtype=np.float32),
+            std=np.ones(num_channels, dtype=np.float32),
+            to_rgb=False,
+        )
+        return results
+
+    def __repr__(self):
+        """str: Return a string that describes the module."""
+        repr_str = self.__class__.__name__
+        repr_str += f"(to_float32={self.to_float32}, "
+        repr_str += f"color_type='{self.color_type}')"
+        return repr_str
+
+
+@PIPELINES.register_module()
+class LoadPointsFromFile(object):
+    """Load Points From File.
+
+    Load points from file.
+
+    Args:
+        coord_type (str): The type of coordinates of points cloud.
+            Available options includes:
+            - 'LIDAR': Points in LiDAR coordinates.
+            - 'DEPTH': Points in depth coordinates, usually for indoor dataset.
+            - 'CAMERA': Points in camera coordinates.
+        load_dim (int, optional): The dimension of the loaded points.
+            Defaults to 6.
+        use_dim (list[int], optional): Which dimensions of the points to use.
+            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
+            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
+        shift_height (bool, optional): Whether to use shifted height.
+            Defaults to False.
+        use_color (bool, optional): Whether to use color features.
+            Defaults to False.
+        file_client_args (dict, optional): Config dict of file clients,
+            refer to
+            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
+            for more details. Defaults to dict(backend='disk').
+    """
+
+    def __init__(
+        self,
+        coord_type,
+        load_dim=6,
+        use_dim=[0, 1, 2],
+        shift_height=False,
+        use_color=False,
+        file_client_args=dict(backend="disk"),
+    ):
+        self.shift_height = shift_height
+        self.use_color = use_color
+        if isinstance(use_dim, int):
+            use_dim = list(range(use_dim))
+        assert (
+            max(use_dim) < load_dim
+        ), f"Expect all used dimensions < {load_dim}, got {use_dim}"
+        assert coord_type in ["CAMERA", "LIDAR", "DEPTH"]
+
+        self.coord_type = coord_type
+        self.load_dim = load_dim
+        self.use_dim = use_dim
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
+
+    def _load_points(self, pts_filename):
+        """Private function to load point clouds data.
+
+        Args:
+            pts_filename (str): Filename of point clouds data.
+
+        Returns:
+            np.ndarray: An array containing point clouds data.
+        """
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+        try:
+            pts_bytes = self.file_client.get(pts_filename)
+            points = np.frombuffer(pts_bytes, dtype=np.float32)
+        except ConnectionError:
+            mmcv.check_file_exist(pts_filename)
+            if pts_filename.endswith(".npy"):
+                points = np.load(pts_filename)
+            else:
+                points = np.fromfile(pts_filename, dtype=np.float32)
+
+        return points
+
+    def __call__(self, results):
+        """Call function to load points data from file.
+
+        Args:
+            results (dict): Result dict containing point clouds data.
+
+        Returns:
+            dict: The result dict containing the point clouds data.
+                Added key and value are described below.
+
+                - points (:obj:`BasePoints`): Point clouds data.
+        """
+        pts_filename = results["pts_filename"]
+        points = self._load_points(pts_filename)
+        points = points.reshape(-1, self.load_dim)
+        points = points[:, self.use_dim]
+        attribute_dims = None
+
+        if self.shift_height:
+            floor_height = np.percentile(points[:, 2], 0.99)
+            height = points[:, 2] - floor_height
+            points = np.concatenate(
+                [points[:, :3], np.expand_dims(height, 1), points[:, 3:]], 1
+            )
+            attribute_dims = dict(height=3)
+
+        if self.use_color:
+            assert len(self.use_dim) >= 6
+            if attribute_dims is None:
+                attribute_dims = dict()
+            attribute_dims.update(
+                dict(
+                    color=[
+                        points.shape[1] - 3,
+                        points.shape[1] - 2,
+                        points.shape[1] - 1,
+                    ]
+                )
+            )
+
+        results["points"] = points
+        return results

projects/mmdet3d_plugin/datasets/pipelines/transform.py

+import numpy as np
+import mmcv
+from mmcv.parallel import DataContainer as DC
+from mmdet.datasets.builder import PIPELINES
+from mmdet.datasets.pipelines import to_tensor
+
+
+@PIPELINES.register_module()
+class MultiScaleDepthMapGenerator(object):
+    def __init__(self, downsample=1, max_depth=60):
+        if not isinstance(downsample, (list, tuple)):
+            downsample = [downsample]
+        self.downsample = downsample
+        self.max_depth = max_depth
+
+    def __call__(self, input_dict):
+        points = input_dict["points"][..., :3, None]
+        gt_depth = []
+        for i, lidar2img in enumerate(input_dict["lidar2img"]):
+            H, W = input_dict["img_shape"][i][:2]
+
+            pts_2d = (
+                np.squeeze(lidar2img[:3, :3] @ points, axis=-1)
+                + lidar2img[:3, 3]
+            )
+            pts_2d[:, :2] /= pts_2d[:, 2:3]
+            U = np.round(pts_2d[:, 0]).astype(np.int32)
+            V = np.round(pts_2d[:, 1]).astype(np.int32)
+            depths = pts_2d[:, 2]
+            mask = np.logical_and.reduce(
+                [
+                    V >= 0,
+                    V < H,
+                    U >= 0,
+                    U < W,
+                    depths >= 0.1,
+                    # depths <= self.max_depth,
+                ]
+            )
+            V, U, depths = V[mask], U[mask], depths[mask]
+            sort_idx = np.argsort(depths)[::-1]
+            V, U, depths = V[sort_idx], U[sort_idx], depths[sort_idx]
+            depths = np.clip(depths, 0.1, self.max_depth)
+            for j, downsample in enumerate(self.downsample):
+                if len(gt_depth) < j + 1:
+                    gt_depth.append([])
+                h, w = (int(H / downsample), int(W / downsample))
+                u = np.floor(U / downsample).astype(np.int32)
+                v = np.floor(V / downsample).astype(np.int32)
+                depth_map = np.ones([h, w], dtype=np.float32) * -1
+                depth_map[v, u] = depths
+                gt_depth[j].append(depth_map)
+
+        input_dict["gt_depth"] = [np.stack(x) for x in gt_depth]
+        return input_dict
+
+
+@PIPELINES.register_module()
+class NuScenesSparse4DAdaptor(object):
+    def __init(self):
+        pass
+
+    def __call__(self, input_dict):
+        input_dict["projection_mat"] = np.float32(
+            np.stack(input_dict["lidar2img"])
+        )
+        input_dict["image_wh"] = np.ascontiguousarray(
+            np.array(input_dict["img_shape"], dtype=np.float32)[:, :2][:, ::-1]
+        )
+        input_dict["T_global_inv"] = np.linalg.inv(input_dict["lidar2global"])
+        input_dict["T_global"] = input_dict["lidar2global"]
+        if "cam_intrinsic" in input_dict:
+            input_dict["cam_intrinsic"] = np.float32(
+                np.stack(input_dict["cam_intrinsic"])
+            )
+            input_dict["focal"] = input_dict["cam_intrinsic"][..., 0, 0]
+        if "instance_inds" in input_dict:
+            input_dict["instance_id"] = input_dict["instance_inds"]
+
+        if "gt_bboxes_3d" in input_dict:
+            input_dict["gt_bboxes_3d"][:, 6] = self.limit_period(
+                input_dict["gt_bboxes_3d"][:, 6], offset=0.5, period=2 * np.pi
+            )
+            input_dict["gt_bboxes_3d"] = DC(
+                to_tensor(input_dict["gt_bboxes_3d"]).float()
+            )
+        if "gt_labels_3d" in input_dict:
+            input_dict["gt_labels_3d"] = DC(
+                to_tensor(input_dict["gt_labels_3d"]).long()
+            )
+
+        imgs = [img.transpose(2, 0, 1) for img in input_dict["img"]]
+        imgs = np.ascontiguousarray(np.stack(imgs, axis=0))
+        input_dict["img"] = DC(to_tensor(imgs), stack=True)
+
+        for key in [
+            'gt_map_labels', 
+            'gt_map_pts',
+            'gt_agent_fut_trajs',
+            'gt_agent_fut_masks',
+        ]:
+            if key not in input_dict:
+                continue
+            input_dict[key] = DC(to_tensor(input_dict[key]), stack=False, cpu_only=False) 
+
+        for key in [
+            'gt_ego_fut_trajs',
+            'gt_ego_fut_masks',
+            'gt_ego_fut_cmd',
+            'ego_status',
+        ]:
+            if key not in input_dict:
+                continue
+            input_dict[key] = DC(to_tensor(input_dict[key]), stack=True, cpu_only=False, pad_dims=None)
+        
+        return input_dict
+
+    def limit_period(
+        self, val: np.ndarray, offset: float = 0.5, period: float = np.pi
+    ) -> np.ndarray:
+        limited_val = val - np.floor(val / period + offset) * period
+        return limited_val
+
+
+@PIPELINES.register_module()
+class InstanceNameFilter(object):
+    """Filter GT objects by their names.
+
+    Args:
+        classes (list[str]): List of class names to be kept for training.
+    """
+
+    def __init__(self, classes):
+        self.classes = classes
+        self.labels = list(range(len(self.classes)))
+
+    def __call__(self, input_dict):
+        """Call function to filter objects by their names.
+
+        Args:
+            input_dict (dict): Result dict from loading pipeline.
+
+        Returns:
+            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d' \
+                keys are updated in the result dict.
+        """
+        gt_labels_3d = input_dict["gt_labels_3d"]
+        gt_bboxes_mask = np.array(
+            [n in self.labels for n in gt_labels_3d], dtype=np.bool_
+        )
+        input_dict["gt_bboxes_3d"] = input_dict["gt_bboxes_3d"][gt_bboxes_mask]
+        input_dict["gt_labels_3d"] = input_dict["gt_labels_3d"][gt_bboxes_mask]
+        if "instance_inds" in input_dict:
+            input_dict["instance_inds"] = input_dict["instance_inds"][gt_bboxes_mask]
+        if "gt_agent_fut_trajs" in input_dict:
+            input_dict["gt_agent_fut_trajs"] = input_dict["gt_agent_fut_trajs"][gt_bboxes_mask]
+            input_dict["gt_agent_fut_masks"] = input_dict["gt_agent_fut_masks"][gt_bboxes_mask]
+        return input_dict
+
+    def __repr__(self):
+        """str: Return a string that describes the module."""
+        repr_str = self.__class__.__name__
+        repr_str += f"(classes={self.classes})"
+        return repr_str
+
+
+@PIPELINES.register_module()
+class CircleObjectRangeFilter(object):
+    def __init__(
+        self, class_dist_thred=[52.5] * 5 + [31.5] + [42] * 3 + [31.5]
+    ):
+        self.class_dist_thred = class_dist_thred
+
+    def __call__(self, input_dict):
+        gt_bboxes_3d = input_dict["gt_bboxes_3d"]
+        gt_labels_3d = input_dict["gt_labels_3d"]
+        dist = np.sqrt(
+            np.sum(gt_bboxes_3d[:, :2] ** 2, axis=-1)
+        )
+        mask = np.array([False] * len(dist))
+        for label_idx, dist_thred in enumerate(self.class_dist_thred):
+            mask = np.logical_or(
+                mask,
+                np.logical_and(gt_labels_3d == label_idx, dist <= dist_thred),
+            )
+
+        gt_bboxes_3d = gt_bboxes_3d[mask]
+        gt_labels_3d = gt_labels_3d[mask]
+
+        input_dict["gt_bboxes_3d"] = gt_bboxes_3d
+        input_dict["gt_labels_3d"] = gt_labels_3d
+        if "instance_inds" in input_dict:
+            input_dict["instance_inds"] = input_dict["instance_inds"][mask]
+        if "gt_agent_fut_trajs" in input_dict:
+            input_dict["gt_agent_fut_trajs"] = input_dict["gt_agent_fut_trajs"][mask]
+            input_dict["gt_agent_fut_masks"] = input_dict["gt_agent_fut_masks"][mask]
+        return input_dict
+
+    def __repr__(self):
+        """str: Return a string that describes the module."""
+        repr_str = self.__class__.__name__
+        repr_str += f"(class_dist_thred={self.class_dist_thred})"
+        return repr_str
+
+
+@PIPELINES.register_module()
+class NormalizeMultiviewImage(object):
+    """Normalize the image.
+    Added key is "img_norm_cfg".
+    Args:
+        mean (sequence): Mean values of 3 channels.
+        std (sequence): Std values of 3 channels.
+        to_rgb (bool): Whether to convert the image from BGR to RGB,
+            default is true.
+    """
+
+    def __init__(self, mean, std, to_rgb=True):
+        self.mean = np.array(mean, dtype=np.float32)
+        self.std = np.array(std, dtype=np.float32)
+        self.to_rgb = to_rgb
+
+    def __call__(self, results):
+        """Call function to normalize images.
+        Args:
+            results (dict): Result dict from loading pipeline.
+        Returns:
+            dict: Normalized results, 'img_norm_cfg' key is added into
+                result dict.
+        """
+        results["img"] = [
+            mmcv.imnormalize(img, self.mean, self.std, self.to_rgb)
+            for img in results["img"]
+        ]
+        results["img_norm_cfg"] = dict(
+            mean=self.mean, std=self.std, to_rgb=self.to_rgb
+        )
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f"(mean={self.mean}, std={self.std}, to_rgb={self.to_rgb})"
+        return repr_str

projects/mmdet3d_plugin/datasets/pipelines/vectorize.py

+from typing import List, Tuple, Union, Dict
+
+import numpy as np
+from shapely.geometry import LineString
+from numpy.typing import NDArray
+
+from mmcv.parallel import DataContainer as DC
+from mmdet.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module(force=True)
+class VectorizeMap(object):
+    """Generate vectoized map and put into `semantic_mask` key.
+    Concretely, shapely geometry objects are converted into sample points (ndarray).
+    We use args `sample_num`, `sample_dist`, `simplify` to specify sampling method.
+
+    Args:
+        roi_size (tuple or list): bev range .
+        normalize (bool): whether to normalize points to range (0, 1).
+        coords_dim (int): dimension of point coordinates.
+        simplify (bool): whether to use simpily function. If true, `sample_num` \
+            and `sample_dist` will be ignored.
+        sample_num (int): number of points to interpolate from a polyline. Set to -1 to ignore.
+        sample_dist (float): interpolate distance. Set to -1 to ignore.
+    """
+
+    def __init__(self, 
+                 roi_size: Union[Tuple, List], 
+                 normalize: bool,
+                 coords_dim: int=2,
+                 simplify: bool=False, 
+                 sample_num: int=-1, 
+                 sample_dist: float=-1, 
+                 permute: bool=False
+        ):
+        self.coords_dim = coords_dim
+        self.sample_num = sample_num
+        self.sample_dist = sample_dist
+        self.roi_size = np.array(roi_size)
+        self.normalize = normalize
+        self.simplify = simplify
+        self.permute = permute
+
+        if sample_dist > 0:
+            assert sample_num < 0 and not simplify
+            self.sample_fn = self.interp_fixed_dist
+        elif sample_num > 0:
+            assert sample_dist < 0 and not simplify
+            self.sample_fn = self.interp_fixed_num
+        else:
+            assert simplify
+
+    def interp_fixed_num(self, line: LineString) -> NDArray:
+        ''' Interpolate a line to fixed number of points.
+        
+        Args:
+            line (LineString): line
+        
+        Returns:
+            points (array): interpolated points, shape (N, 2)
+        '''
+
+        distances = np.linspace(0, line.length, self.sample_num)
+        sampled_points = np.array([list(line.interpolate(distance).coords) 
+            for distance in distances]).squeeze()
+
+        return sampled_points
+
+    def interp_fixed_dist(self, line: LineString) -> NDArray:
+        ''' Interpolate a line at fixed interval.
+        
+        Args:
+            line (LineString): line
+        
+        Returns:
+            points (array): interpolated points, shape (N, 2)
+        '''
+
+        distances = list(np.arange(self.sample_dist, line.length, self.sample_dist))
+        # make sure to sample at least two points when sample_dist > line.length
+        distances = [0,] + distances + [line.length,] 
+        
+        sampled_points = np.array([list(line.interpolate(distance).coords)
+                                for distance in distances]).squeeze()
+        
+        return sampled_points
+    
+    def get_vectorized_lines(self, map_geoms: Dict) -> Dict:
+        ''' Vectorize map elements. Iterate over the input dict and apply the 
+        specified sample funcion.
+        
+        Args:
+            line (LineString): line
+        
+        Returns:
+            vectors (array): dict of vectorized map elements.
+        '''
+
+        vectors = {}
+        for label, geom_list in map_geoms.items():
+            vectors[label] = []
+            for geom in geom_list:
+                if geom.geom_type == 'LineString':
+                    if self.simplify:
+                        line = geom.simplify(0.2, preserve_topology=True)
+                        line = np.array(line.coords)
+                    else:
+                        line = self.sample_fn(geom)
+                    line = line[:, :self.coords_dim]
+
+                    if self.normalize:
+                        line = self.normalize_line(line)
+                    if self.permute:
+                        line = self.permute_line(line)
+                    vectors[label].append(line)
+
+                elif geom.geom_type == 'Polygon':
+                    # polygon objects will not be vectorized
+                    continue
+                
+                else:
+                    raise ValueError('map geoms must be either LineString or Polygon!')
+        return vectors
+    
+    def normalize_line(self, line: NDArray) -> NDArray:
+        ''' Convert points to range (0, 1).
+        
+        Args:
+            line (LineString): line
+        
+        Returns:
+            normalized (array): normalized points.
+        '''
+
+        origin = -np.array([self.roi_size[0]/2, self.roi_size[1]/2])
+
+        line[:, :2] = line[:, :2] - origin
+
+        # transform from range [0, 1] to (0, 1)
+        eps = 1e-5
+        line[:, :2] = line[:, :2] / (self.roi_size + eps)
+
+        return line
+    
+    def permute_line(self, line: np.ndarray, padding=1e5):
+        '''
+        (num_pts, 2) -> (num_permute, num_pts, 2)
+        where num_permute = 2 * (num_pts - 1)
+        '''
+        is_closed = np.allclose(line[0], line[-1], atol=1e-3)
+        num_points = len(line)
+        permute_num = num_points - 1
+        permute_lines_list = []
+        if is_closed:
+            pts_to_permute = line[:-1, :] # throw away replicate start end pts
+            for shift_i in range(permute_num):
+                permute_lines_list.append(np.roll(pts_to_permute, shift_i, axis=0))
+            flip_pts_to_permute = np.flip(pts_to_permute, axis=0)
+            for shift_i in range(permute_num):
+                permute_lines_list.append(np.roll(flip_pts_to_permute, shift_i, axis=0))
+        else:
+            permute_lines_list.append(line)
+            permute_lines_list.append(np.flip(line, axis=0))
+
+        permute_lines_array = np.stack(permute_lines_list, axis=0)
+
+        if is_closed:
+            tmp = np.zeros((permute_num * 2, num_points, self.coords_dim))
+            tmp[:, :-1, :] = permute_lines_array
+            tmp[:, -1, :] = permute_lines_array[:, 0, :] # add replicate start end pts
+            permute_lines_array = tmp
+
+        else:
+            # padding
+            padding = np.full([permute_num * 2 - 2, num_points, self.coords_dim], padding)
+            permute_lines_array = np.concatenate((permute_lines_array, padding), axis=0)
+        
+        return permute_lines_array
+    
+    def __call__(self, input_dict):
+        if "map_geoms" not in input_dict:
+            return input_dict
+        map_geoms = input_dict['map_geoms']
+        vectors = self.get_vectorized_lines(map_geoms)
+
+        if self.permute:
+            gt_map_labels, gt_map_pts = [], []
+            for label, vecs in vectors.items():
+                for vec in vecs:
+                    gt_map_labels.append(label)
+                    gt_map_pts.append(vec)
+            input_dict['gt_map_labels'] = np.array(gt_map_labels, dtype=np.int64)
+            input_dict['gt_map_pts'] = np.array(gt_map_pts, dtype=np.float32).reshape(-1, 2 * (self.sample_num - 1), self.sample_num, self.coords_dim)
+        else:
+            input_dict['vectors'] = DC(vectors, stack=False, cpu_only=True)
+        
+        return input_dict
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += f'(simplify={self.simplify}, '
+        repr_str += f'sample_num={self.sample_num}), '
+        repr_str += f'sample_dist={self.sample_dist}), ' 
+        repr_str += f'roi_size={self.roi_size})'
+        repr_str += f'normalize={self.normalize})'
+        repr_str += f'coords_dim={self.coords_dim})'
+
+        return repr_str
\ No newline at end of file

projects/mmdet3d_plugin/datasets/samplers/__init__.py

+from .group_sampler import DistributedGroupSampler
+from .distributed_sampler import DistributedSampler
+from .sampler import SAMPLER, build_sampler
+from .group_in_batch_sampler import (
+    GroupInBatchSampler,
+)

projects/mmdet3d_plugin/datasets/samplers/distributed_sampler.py

+import math
+
+import torch
+from torch.utils.data import DistributedSampler as _DistributedSampler
+from .sampler import SAMPLER
+
+import pdb
+import sys
+
+
+class ForkedPdb(pdb.Pdb):
+    def interaction(self, *args, **kwargs):
+        _stdin = sys.stdin
+        try:
+            sys.stdin = open("/dev/stdin")
+            pdb.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+
+
+def set_trace():
+    ForkedPdb().set_trace(sys._getframe().f_back)
+
+
+@SAMPLER.register_module()
+class DistributedSampler(_DistributedSampler):
+    def __init__(
+        self, dataset=None, num_replicas=None, rank=None, shuffle=True, seed=0
+    ):
+        super().__init__(
+            dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle
+        )
+        # for the compatibility from PyTorch 1.3+
+        self.seed = seed if seed is not None else 0
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        assert not self.shuffle
+        if "data_infos" in dir(self.dataset):
+            timestamps = [
+                x["timestamp"] / 1e6 for x in self.dataset.data_infos
+            ]
+            vehicle_idx = [
+                x["lidar_path"].split("/")[-1][:4]
+                if "lidar_path" in x
+                else None
+                for x in self.dataset.data_infos
+            ]
+        else:
+            timestamps = [
+                x["timestamp"] / 1e6
+                for x in self.dataset.datasets[0].data_infos
+            ] * len(self.dataset.datasets)
+            vehicle_idx = [
+                x["lidar_path"].split("/")[-1][:4]
+                if "lidar_path" in x
+                else None
+                for x in self.dataset.datasets[0].data_infos
+            ] * len(self.dataset.datasets)
+
+        sequence_splits = []
+        for i in range(len(timestamps)):
+            if i == 0 or (
+                abs(timestamps[i] - timestamps[i - 1]) > 4
+                or vehicle_idx[i] != vehicle_idx[i - 1]
+            ):
+                sequence_splits.append([i])
+            else:
+                sequence_splits[-1].append(i)
+
+        indices = []
+        perfix_sum = 0
+        split_length = len(self.dataset) // self.num_replicas
+        for i in range(len(sequence_splits)):
+            if perfix_sum >= (self.rank + 1) * split_length:
+                break
+            elif perfix_sum >= self.rank * split_length:
+                indices.extend(sequence_splits[i])
+            perfix_sum += len(sequence_splits[i])
+
+        self.num_samples = len(indices)
+        return iter(indices)

projects/mmdet3d_plugin/datasets/samplers/group_in_batch_sampler.py

+# https://github.com/Divadi/SOLOFusion/blob/main/mmdet3d/datasets/samplers/infinite_group_each_sample_in_batch_sampler.py
+import itertools
+import copy
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from mmcv.runner import get_dist_info
+from torch.utils.data.sampler import Sampler
+
+
+# https://github.com/open-mmlab/mmdetection/blob/3b72b12fe9b14de906d1363982b9fba05e7d47c1/mmdet/core/utils/dist_utils.py#L157
+def sync_random_seed(seed=None, device="cuda"):
+    """Make sure different ranks share the same seed.
+    All workers must call this function, otherwise it will deadlock.
+    This method is generally used in `DistributedSampler`,
+    because the seed should be identical across all processes
+    in the distributed group.
+    In distributed sampling, different ranks should sample non-overlapped
+    data in the dataset. Therefore, this function is used to make sure that
+    each rank shuffles the data indices in the same order based
+    on the same seed. Then different ranks could use different indices
+    to select non-overlapped data from the same data list.
+    Args:
+        seed (int, Optional): The seed. Default to None.
+        device (str): The device where the seed will be put on.
+            Default to 'cuda'.
+    Returns:
+        int: Seed to be used.
+    """
+    if seed is None:
+        seed = np.random.randint(2**31)
+    assert isinstance(seed, int)
+
+    rank, world_size = get_dist_info()
+
+    if world_size == 1:
+        return seed
+
+    if rank == 0:
+        random_num = torch.tensor(seed, dtype=torch.int32, device=device)
+    else:
+        random_num = torch.tensor(0, dtype=torch.int32, device=device)
+    dist.broadcast(random_num, src=0)
+    return random_num.item()
+
+
+class GroupInBatchSampler(Sampler):
+    """
+    Pardon this horrendous name. Basically, we want every sample to be from its own group.
+    If batch size is 4 and # of GPUs is 8, each sample of these 32 should be operating on
+    its own group.
+
+    Shuffling is only done for group order, not done within groups.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        batch_size=1,
+        world_size=None,
+        rank=None,
+        seed=0,
+        skip_prob=0.,
+        sequence_flip_prob=0.,
+    ):
+        _rank, _world_size = get_dist_info()
+        if world_size is None:
+            world_size = _world_size
+        if rank is None:
+            rank = _rank
+
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.world_size = world_size
+        self.rank = rank
+        self.seed = sync_random_seed(seed)
+
+        self.size = len(self.dataset)
+
+        assert hasattr(self.dataset, "flag")
+        self.flag = self.dataset.flag
+        self.group_sizes = np.bincount(self.flag)
+        self.groups_num = len(self.group_sizes)
+        self.global_batch_size = batch_size * world_size
+        assert self.groups_num >= self.global_batch_size
+
+        # Now, for efficiency, make a dict group_idx: List[dataset sample_idxs]
+        self.group_idx_to_sample_idxs = {
+            group_idx: np.where(self.flag == group_idx)[0].tolist()
+            for group_idx in range(self.groups_num)
+        }
+
+        # Get a generator per sample idx. Considering samples over all
+        # GPUs, each sample position has its own generator
+        self.group_indices_per_global_sample_idx = [
+            self._group_indices_per_global_sample_idx(
+                self.rank * self.batch_size + local_sample_idx
+            )
+            for local_sample_idx in range(self.batch_size)
+        ]
+
+        # Keep track of a buffer of dataset sample idxs for each local sample idx
+        self.buffer_per_local_sample = [[] for _ in range(self.batch_size)]
+        self.aug_per_local_sample = [None for _ in range(self.batch_size)]
+        self.skip_prob = skip_prob
+        self.sequence_flip_prob = sequence_flip_prob
+
+    def _infinite_group_indices(self):
+        g = torch.Generator()
+        g.manual_seed(self.seed)
+        while True:
+            yield from torch.randperm(self.groups_num, generator=g).tolist()
+
+    def _group_indices_per_global_sample_idx(self, global_sample_idx):
+        yield from itertools.islice(
+            self._infinite_group_indices(),
+            global_sample_idx,
+            None,
+            self.global_batch_size,
+        )
+
+    def __iter__(self):
+        while True:
+            curr_batch = []
+            for local_sample_idx in range(self.batch_size):
+                skip = (
+                    np.random.uniform() < self.skip_prob
+                    and len(self.buffer_per_local_sample[local_sample_idx]) > 1
+                )
+                if len(self.buffer_per_local_sample[local_sample_idx]) == 0:
+                    # Finished current group, refill with next group
+                    # skip = False
+                    new_group_idx = next(
+                        self.group_indices_per_global_sample_idx[
+                            local_sample_idx
+                        ]
+                    )
+                    self.buffer_per_local_sample[
+                        local_sample_idx
+                    ] = copy.deepcopy(
+                        self.group_idx_to_sample_idxs[new_group_idx]
+                    )
+                    if np.random.uniform() < self.sequence_flip_prob:
+                        self.buffer_per_local_sample[
+                            local_sample_idx
+                        ] = self.buffer_per_local_sample[local_sample_idx][
+                            ::-1
+                        ]
+                    if self.dataset.keep_consistent_seq_aug:
+                        self.aug_per_local_sample[
+                            local_sample_idx
+                        ] = self.dataset.get_augmentation()
+
+                if not self.dataset.keep_consistent_seq_aug:
+                    self.aug_per_local_sample[
+                        local_sample_idx
+                    ] = self.dataset.get_augmentation()
+
+                if skip:
+                    self.buffer_per_local_sample[local_sample_idx].pop(0)
+                curr_batch.append(
+                    dict(
+                        idx=self.buffer_per_local_sample[local_sample_idx].pop(
+                            0
+                        ),
+                        aug_config=self.aug_per_local_sample[local_sample_idx],
+                    )
+                )
+
+            yield curr_batch
+
+    def __len__(self):
+        """Length of base dataset."""
+        return self.size
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

projects/mmdet3d_plugin/datasets/samplers/group_sampler.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import math
+
+import numpy as np
+import torch
+from mmcv.runner import get_dist_info
+from torch.utils.data import Sampler
+from .sampler import SAMPLER
+import random
+from IPython import embed
+
+
+@SAMPLER.register_module()
+class DistributedGroupSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+        seed (int, optional): random seed used to shuffle the sampler if
+            ``shuffle=True``. This number should be identical across all
+            processes in the distributed group. Default: 0.
+    """
+
+    def __init__(
+        self, dataset, samples_per_gpu=1, num_replicas=None, rank=None, seed=0
+    ):
+        _rank, _num_replicas = get_dist_info()
+        if num_replicas is None:
+            num_replicas = _num_replicas
+        if rank is None:
+            rank = _rank
+        self.dataset = dataset
+        self.samples_per_gpu = samples_per_gpu
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.seed = seed if seed is not None else 0
+
+        assert hasattr(self.dataset, "flag")
+        self.flag = self.dataset.flag
+        self.group_sizes = np.bincount(self.flag)
+
+        self.num_samples = 0
+        for i, j in enumerate(self.group_sizes):
+            self.num_samples += (
+                int(
+                    math.ceil(
+                        self.group_sizes[i]
+                        * 1.0
+                        / self.samples_per_gpu
+                        / self.num_replicas
+                    )
+                )
+                * self.samples_per_gpu
+            )
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch + self.seed)
+
+        indices = []
+        for i, size in enumerate(self.group_sizes):
+            if size > 0:
+                indice = np.where(self.flag == i)[0]
+                assert len(indice) == size
+                # add .numpy() to avoid bug when selecting indice in parrots.
+                # TODO: check whether torch.randperm() can be replaced by
+                # numpy.random.permutation().
+                indice = indice[
+                    list(torch.randperm(int(size), generator=g).numpy())
+                ].tolist()
+                extra = int(
+                    math.ceil(
+                        size * 1.0 / self.samples_per_gpu / self.num_replicas
+                    )
+                ) * self.samples_per_gpu * self.num_replicas - len(indice)
+                # pad indice
+                tmp = indice.copy()
+                for _ in range(extra // size):
+                    indice.extend(tmp)
+                indice.extend(tmp[: extra % size])
+                indices.extend(indice)
+
+        assert len(indices) == self.total_size
+
+        indices = [
+            indices[j]
+            for i in list(
+                torch.randperm(
+                    len(indices) // self.samples_per_gpu, generator=g
+                )
+            )
+            for j in range(
+                i * self.samples_per_gpu, (i + 1) * self.samples_per_gpu
+            )
+        ]
+
+        # subsample
+        offset = self.num_samples * self.rank
+        indices = indices[offset : offset + self.num_samples]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

projects/mmdet3d_plugin/datasets/samplers/sampler.py

+from mmcv.utils.registry import Registry, build_from_cfg
+
+SAMPLER = Registry("sampler")
+
+
+def build_sampler(cfg, default_args):
+    return build_from_cfg(cfg, SAMPLER, default_args)

projects/mmdet3d_plugin/datasets/utils.py

+import copy
+
+import cv2
+import numpy as np
+import torch
+
+from projects.mmdet3d_plugin.core.box3d import *
+
+
+def box3d_to_corners(box3d):
+    if isinstance(box3d, torch.Tensor):
+        box3d = box3d.detach().cpu().numpy()
+    corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1)
+    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
+    # use relative origin [0.5, 0.5, 0]
+    corners_norm = corners_norm - np.array([0.5, 0.5, 0.5])
+    corners = box3d[:, None, [W, L, H]] * corners_norm.reshape([1, 8, 3])
+
+    # rotate around z axis
+    rot_cos = np.cos(box3d[:, YAW])
+    rot_sin = np.sin(box3d[:, YAW])
+    rot_mat = np.tile(np.eye(3)[None], (box3d.shape[0], 1, 1))
+    rot_mat[:, 0, 0] = rot_cos
+    rot_mat[:, 0, 1] = -rot_sin
+    rot_mat[:, 1, 0] = rot_sin
+    rot_mat[:, 1, 1] = rot_cos
+    corners = (rot_mat[:, None] @ corners[..., None]).squeeze(axis=-1)
+    corners += box3d[:, None, :3]
+    return corners
+
+
+def plot_rect3d_on_img(
+    img, num_rects, rect_corners, color=(0, 255, 0), thickness=1
+):
+    """Plot the boundary lines of 3D rectangular on 2D images.
+
+    Args:
+        img (numpy.array): The numpy array of image.
+        num_rects (int): Number of 3D rectangulars.
+        rect_corners (numpy.array): Coordinates of the corners of 3D
+            rectangulars. Should be in the shape of [num_rect, 8, 2].
+        color (tuple[int], optional): The color to draw bboxes.
+            Default: (0, 255, 0).
+        thickness (int, optional): The thickness of bboxes. Default: 1.
+    """
+    line_indices = (
+        (0, 1),
+        (0, 3),
+        (0, 4),
+        (1, 2),
+        (1, 5),
+        (3, 2),
+        (3, 7),
+        (4, 5),
+        (4, 7),
+        (2, 6),
+        (5, 6),
+        (6, 7),
+    )
+    h, w = img.shape[:2]
+    for i in range(num_rects):
+        corners = np.clip(rect_corners[i], -1e4, 1e5).astype(np.int32)
+        for start, end in line_indices:
+            if (
+                (corners[start, 1] >= h or corners[start, 1] < 0)
+                or (corners[start, 0] >= w or corners[start, 0] < 0)
+            ) and (
+                (corners[end, 1] >= h or corners[end, 1] < 0)
+                or (corners[end, 0] >= w or corners[end, 0] < 0)
+            ):
+                continue
+            if isinstance(color[0], int):
+                cv2.line(
+                    img,
+                    (corners[start, 0], corners[start, 1]),
+                    (corners[end, 0], corners[end, 1]),
+                    color,
+                    thickness,
+                    cv2.LINE_AA,
+                )
+            else:
+                cv2.line(
+                    img,
+                    (corners[start, 0], corners[start, 1]),
+                    (corners[end, 0], corners[end, 1]),
+                    color[i],
+                    thickness,
+                    cv2.LINE_AA,
+                )
+
+    return img.astype(np.uint8)
+
+
+def draw_lidar_bbox3d_on_img(
+    bboxes3d, raw_img, lidar2img_rt, img_metas=None, color=(0, 255, 0), thickness=1
+):
+    """Project the 3D bbox on 2D plane and draw on input image.
+
+    Args:
+        bboxes3d (:obj:`LiDARInstance3DBoxes`):
+            3d bbox in lidar coordinate system to visualize.
+        raw_img (numpy.array): The numpy array of image.
+        lidar2img_rt (numpy.array, shape=[4, 4]): The projection matrix
+            according to the camera intrinsic parameters.
+        img_metas (dict): Useless here.
+        color (tuple[int], optional): The color to draw bboxes.
+            Default: (0, 255, 0).
+        thickness (int, optional): The thickness of bboxes. Default: 1.
+    """
+    img = raw_img.copy()
+    # corners_3d = bboxes3d.corners
+    corners_3d = box3d_to_corners(bboxes3d)
+    num_bbox = corners_3d.shape[0]
+    pts_4d = np.concatenate(
+        [corners_3d.reshape(-1, 3), np.ones((num_bbox * 8, 1))], axis=-1
+    )
+    lidar2img_rt = copy.deepcopy(lidar2img_rt).reshape(4, 4)
+    if isinstance(lidar2img_rt, torch.Tensor):
+        lidar2img_rt = lidar2img_rt.cpu().numpy()
+    pts_2d = pts_4d @ lidar2img_rt.T
+
+    pts_2d[:, 2] = np.clip(pts_2d[:, 2], a_min=1e-5, a_max=1e5)
+    pts_2d[:, 0] /= pts_2d[:, 2]
+    pts_2d[:, 1] /= pts_2d[:, 2]
+    imgfov_pts_2d = pts_2d[..., :2].reshape(num_bbox, 8, 2)
+
+    return plot_rect3d_on_img(img, num_bbox, imgfov_pts_2d, color, thickness)
+
+
+def draw_points_on_img(points, img, lidar2img_rt, color=(0, 255, 0), circle=4):
+    img = img.copy()
+    N = points.shape[0]
+    points = points.cpu().numpy()
+    lidar2img_rt = copy.deepcopy(lidar2img_rt).reshape(4, 4)
+    if isinstance(lidar2img_rt, torch.Tensor):
+        lidar2img_rt = lidar2img_rt.cpu().numpy()
+    pts_2d = (
+        np.sum(points[:, :, None] * lidar2img_rt[:3, :3], axis=-1)
+        + lidar2img_rt[:3, 3]
+    )
+    pts_2d[..., 2] = np.clip(pts_2d[..., 2], a_min=1e-5, a_max=1e5)
+    pts_2d = pts_2d[..., :2] / pts_2d[..., 2:3]
+    pts_2d = np.clip(pts_2d, -1e4, 1e4).astype(np.int32)
+
+    for i in range(N):
+        for point in pts_2d[i]:
+            if isinstance(color[0], int):
+                color_tmp = color
+            else:
+                color_tmp = color[i]
+            cv2.circle(img, point.tolist(), circle, color_tmp, thickness=-1)
+    return img.astype(np.uint8)
+
+
+def draw_lidar_bbox3d_on_bev(
+    bboxes_3d, bev_size, bev_range=115, color=(255, 0, 0), thickness=3):
+    if isinstance(bev_size, (list, tuple)):
+        bev_h, bev_w = bev_size
+    else:
+        bev_h, bev_w = bev_size, bev_size
+    bev = np.zeros([bev_h, bev_w, 3])
+
+    marking_color = (127, 127, 127)
+    bev_resolution = bev_range / bev_h
+    for cir in range(int(bev_range / 2 / 10)):
+        cv2.circle(
+            bev,
+            (int(bev_h / 2), int(bev_w / 2)),
+            int((cir + 1) * 10 / bev_resolution),
+            marking_color,
+            thickness=thickness,
+        )
+    cv2.line(
+        bev,
+        (0, int(bev_h / 2)),
+        (bev_w, int(bev_h / 2)),
+        marking_color,
+    )
+    cv2.line(
+        bev,
+        (int(bev_w / 2), 0),
+        (int(bev_w / 2), bev_h),
+        marking_color,
+    )
+    if len(bboxes_3d) != 0:
+        bev_corners = box3d_to_corners(bboxes_3d)[:, [0, 3, 4, 7]][
+            ..., [0, 1]
+        ]
+        xs = bev_corners[..., 0] / bev_resolution + bev_w / 2
+        ys = -bev_corners[..., 1] / bev_resolution + bev_h / 2
+        for obj_idx, (x, y) in enumerate(zip(xs, ys)):
+            for p1, p2 in ((0, 1), (0, 2), (1, 3), (2, 3)):
+                if isinstance(color[0], (list, tuple)):
+                    tmp = color[obj_idx]
+                else:
+                    tmp = color
+                cv2.line(
+                    bev,
+                    (int(x[p1]), int(y[p1])),
+                    (int(x[p2]), int(y[p2])),
+                    tmp,
+                    thickness=thickness,
+                )
+    return bev.astype(np.uint8)
+
+
+def draw_lidar_bbox3d(bboxes_3d, imgs, lidar2imgs, color=(255, 0, 0)):
+    vis_imgs = []
+    for i, (img, lidar2img) in enumerate(zip(imgs, lidar2imgs)):
+        vis_imgs.append(
+            draw_lidar_bbox3d_on_img(bboxes_3d, img, lidar2img, color=color)
+        )
+
+    num_imgs = len(vis_imgs)
+    if num_imgs < 4 or num_imgs % 2 != 0:
+        vis_imgs = np.concatenate(vis_imgs, axis=1)
+    else:
+        vis_imgs = np.concatenate([
+            np.concatenate(vis_imgs[:num_imgs//2], axis=1),
+            np.concatenate(vis_imgs[num_imgs//2:], axis=1)
+        ], axis=0)
+
+    bev = draw_lidar_bbox3d_on_bev(bboxes_3d, vis_imgs.shape[0], color=color)
+    vis_imgs = np.concatenate([bev, vis_imgs], axis=1)
+    return vis_imgs

projects/mmdet3d_plugin/models/__init__.py

+from .sparsedrive import SparseDrive
+from .sparsedrive_head import SparseDriveHead
+from .blocks import (
+    DeformableFeatureAggregation,
+    DenseDepthNet,
+    AsymmetricFFN,
+)
+from .instance_bank import InstanceBank
+from .detection3d import (
+    SparseBox3DDecoder,
+    SparseBox3DTarget,
+    SparseBox3DRefinementModule,
+    SparseBox3DKeyPointsGenerator,
+    SparseBox3DEncoder,
+)
+from .map import *
+from .motion import *
+
+
+__all__ = [
+    "SparseDrive",
+    "SparseDriveHead",
+    "DeformableFeatureAggregation",
+    "DenseDepthNet",
+    "AsymmetricFFN",
+    "InstanceBank",
+    "SparseBox3DDecoder",
+    "SparseBox3DTarget",
+    "SparseBox3DRefinementModule",
+    "SparseBox3DKeyPointsGenerator",
+    "SparseBox3DEncoder",
+]

projects/mmdet3d_plugin/models/attention.py

+import warnings
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.functional import linear
+from torch.nn.init import xavier_uniform_, constant_
+
+from mmcv.utils import deprecated_api_warning
+from mmcv.runner import auto_fp16
+from mmcv.runner.base_module import BaseModule
+from mmcv.cnn.bricks.drop import build_dropout
+from mmcv.cnn.bricks.registry import ATTENTION
+import torch.utils.checkpoint as cp
+
+
+from einops import rearrange
+try:
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_kvpacked_func
+    print('Use flash_attn_unpadded_kvpacked_func')
+except:
+    from flash_attn.flash_attn_interface import  flash_attn_varlen_kvpacked_func as flash_attn_unpadded_kvpacked_func
+    print('Use flash_attn_varlen_kvpacked_func')
+from flash_attn.bert_padding import unpad_input, pad_input, index_first_axis
+
+
+def _in_projection_packed(q, k, v, w, b = None):
+    w_q, w_k, w_v = w.chunk(3)
+    if b is None:
+        b_q = b_k = b_v = None
+    else:
+        b_q, b_k, b_v = b.chunk(3)
+    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
+
+
+class FlashAttention(nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.1)
+    """
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+        self.fp16_enabled = True
+
+    @auto_fp16(apply_to=('q', 'kv'), out_fp32=True)
+    def forward(self, q, kv, 
+                causal=False, 
+                key_padding_mask=None):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            q: The tensor containing the query. (B, T, H, D) 
+            kv: The tensor containing the key, and value. (B, S, 2, H, D) 
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert q.dtype in [torch.float16, torch.bfloat16] and kv.dtype in [torch.float16, torch.bfloat16]
+        assert q.is_cuda and kv.is_cuda
+        assert q.shape[0] == kv.shape[0] and q.shape[-2] == kv.shape[-2] and q.shape[-1] == kv.shape[-1]
+
+        batch_size = q.shape[0]
+        seqlen_q, seqlen_k = q.shape[1], kv.shape[1]
+        if key_padding_mask is None:
+            q, kv = rearrange(q, 'b s ... -> (b s) ...'), rearrange(kv, 'b s ... -> (b s) ...')
+            max_sq, max_sk = seqlen_q, seqlen_k 
+            cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32,
+                                    device=q.device)
+            cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32,
+                                    device=kv.device)                    
+            output = flash_attn_unpadded_kvpacked_func(
+                q, kv, cu_seqlens_q, cu_seqlens_k, max_sq, max_sk,
+                self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+            output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+        else:
+            nheads = kv.shape[-2]
+            q = rearrange(q, 'b s ... -> (b s) ...')
+            max_sq = seqlen_q
+            cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32,
+                                    device=q.device)
+            x = rearrange(kv, 'b s two h d -> b s (two h d)')
+            x_unpad, indices, cu_seqlens_k, max_sk = unpad_input(x, key_padding_mask)
+            x_unpad = rearrange(x_unpad, 'nnz (two h d) -> nnz two h d', two=2, h=nheads)
+            output_unpad = flash_attn_unpadded_kvpacked_func(
+                q, x_unpad, cu_seqlens_q, cu_seqlens_k, max_sq, max_sk,
+                self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+            output = rearrange(output_unpad, '(b s) ... -> b s ...', b=batch_size)
+
+        return output, None
+
+
+class FlashMHA(nn.Module):
+
+    def __init__(self, embed_dim, num_heads, bias=True, batch_first=True, attention_dropout=0.0,
+                 causal=False, device=None, dtype=None, **kwargs) -> None:
+        assert batch_first
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.causal = causal
+        self.bias = bias
+
+        self.num_heads = num_heads
+        assert self.embed_dim % num_heads == 0, "self.kdim must be divisible by num_heads"
+        self.head_dim = self.embed_dim // num_heads
+        assert self.head_dim % 8 == 0 and self.head_dim <= 128, "Only support head_dim <= 128 and divisible by 8"
+
+        self.in_proj_weight = nn.Parameter(torch.empty((3 * embed_dim, embed_dim)))
+        if bias:
+            self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
+        self.inner_attn = FlashAttention(attention_dropout=attention_dropout, **factory_kwargs)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        xavier_uniform_(self.in_proj_weight)
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.)
+            constant_(self.out_proj.bias, 0.)
+        
+    def forward(self, q, k, v, key_padding_mask=None):
+        """x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim)
+        key_padding_mask: bool tensor of shape (batch, seqlen)
+        """
+        q, k, v = _in_projection_packed(q, k, v, self.in_proj_weight, self.in_proj_bias)
+        q = rearrange(q, 'b s (h d) -> b s h d', h=self.num_heads)
+        k = rearrange(k, 'b s (h d) -> b s h d', h=self.num_heads)
+        v = rearrange(v, 'b s (h d) -> b s h d', h=self.num_heads)
+        kv = torch.stack([k, v], dim=2)
+        
+        context, attn_weights = self.inner_attn(q, kv, key_padding_mask=key_padding_mask, causal=self.causal)
+        return self.out_proj(rearrange(context, 'b s h d -> b s (h d)')), attn_weights
+
+
+@ATTENTION.register_module()
+class MultiheadFlashAttention(BaseModule):
+    """A wrapper for ``torch.nn.MultiheadAttention``.
+    This module implements MultiheadAttention with identity connection,
+    and positional encoding  is also passed as input.
+    Args:
+        embed_dims (int): The embedding dimension.
+        num_heads (int): Parallel attention heads.
+        attn_drop (float): A Dropout layer on attn_output_weights.
+            Default: 0.0.
+        proj_drop (float): A Dropout layer after `nn.MultiheadAttention`.
+            Default: 0.0.
+        dropout_layer (agent:`ConfigDict`): The dropout_layer used
+            when adding the shortcut.
+        init_cfg (agent:`mmcv.ConfigDict`): The Config for initialization.
+            Default: None.
+        batch_first (bool): When it is True,  Key, Query and Value are shape of
+            (batch, n, embed_dim), otherwise (n, batch, embed_dim).
+             Default to False.
+    """
+
+    def __init__(self,
+                 embed_dims,
+                 num_heads,
+                 attn_drop=0.,
+                 proj_drop=0.,
+                 dropout_layer=dict(type='Dropout', drop_prob=0.),
+                 init_cfg=None,
+                 batch_first=True,
+                 **kwargs):
+        super(MultiheadFlashAttention, self).__init__(init_cfg)
+        if 'dropout' in kwargs:
+            warnings.warn(
+                'The arguments `dropout` in MultiheadAttention '
+                'has been deprecated, now you can separately '
+                'set `attn_drop`(float), proj_drop(float), '
+                'and `dropout_layer`(dict) ', DeprecationWarning)
+            attn_drop = kwargs['dropout']
+            dropout_layer['drop_prob'] = kwargs.pop('dropout')
+
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.batch_first = True
+        self.attn = FlashMHA(
+            embed_dim=embed_dims, 
+            num_heads=num_heads, 
+            attention_dropout=attn_drop, 
+            dtype=torch.float16, 
+            device='cuda',
+            **kwargs
+        )
+
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.dropout_layer = build_dropout(
+            dropout_layer) if dropout_layer else nn.Identity()
+
+    @deprecated_api_warning({'residual': 'identity'},
+                            cls_name='MultiheadAttention')
+    def forward(self,
+                query,
+                key=None,
+                value=None,
+                identity=None,
+                query_pos=None,
+                key_pos=None,
+                attn_mask=None,
+                key_padding_mask=None,
+                **kwargs):
+        """Forward function for `MultiheadAttention`.
+        **kwargs allow passing a more general data flow when combining
+        with other operations in `transformerlayer`.
+        Args:
+            query (Tensor): The input query with shape [num_queries, bs,
+                embed_dims] if self.batch_first is False, else
+                [bs, num_queries embed_dims].
+            key (Tensor): The key tensor with shape [num_keys, bs,
+                embed_dims] if self.batch_first is False, else
+                [bs, num_keys, embed_dims] .
+                If None, the ``query`` will be used. Defaults to None.
+            value (Tensor): The value tensor with same shape as `key`.
+                Same in `nn.MultiheadAttention.forward`. Defaults to None.
+                If None, the `key` will be used.
+            identity (Tensor): This tensor, with the same shape as x,
+                will be used for the identity link.
+                If None, `x` will be used. Defaults to None.
+            query_pos (Tensor): The positional encoding for query, with
+                the same shape as `x`. If not None, it will
+                be added to `x` before forward function. Defaults to None.
+            key_pos (Tensor): The positional encoding for `key`, with the
+                same shape as `key`. Defaults to None. If not None, it will
+                be added to `key` before forward function. If None, and
+                `query_pos` has the same shape as `key`, then `query_pos`
+                will be used for `key_pos`. Defaults to None.
+            attn_mask (Tensor): ByteTensor mask with shape [num_queries,
+                num_keys]. Same in `nn.MultiheadAttention.forward`.
+                Defaults to None.
+            key_padding_mask (Tensor): ByteTensor with shape [bs, num_keys].
+                Defaults to None.
+        Returns:
+            Tensor: forwarded results with shape
+            [num_queries, bs, embed_dims]
+            if self.batch_first is False, else
+            [bs, num_queries embed_dims].
+        """
+        assert attn_mask is None, 'attn mask not supported now.'
+        if key is None:
+            key = query
+        if value is None:
+            value = key
+        if identity is None:
+            identity = query
+        if key_pos is None:
+            if query_pos is not None:
+                # use query_pos if key_pos is not available
+                if query_pos.shape == key.shape:
+                    key_pos = query_pos
+                else:
+                    warnings.warn(f'position encoding of key is'
+                                  f'missing in {self.__class__.__name__}.')
+        if query_pos is not None:
+            query = query + query_pos
+        if key_pos is not None:
+            key = key + key_pos
+
+        # The dataflow('key', 'query', 'value') of ``FlashAttention`` is (batch, num_query, embed_dims).
+        if not self.batch_first:
+            query = query.transpose(0, 1)
+            key = key.transpose(0, 1)
+            value = value.transpose(0, 1)
+        
+        out = self.attn(
+            q=query,
+            k=key,
+            v=value,
+            key_padding_mask=key_padding_mask)[0]
+
+        if not self.batch_first:
+            out = out.transpose(0, 1)
+
+        return identity + self.dropout_layer(self.proj_drop(out))
+
+
+def gen_sineembed_for_position(pos_tensor, hidden_dim=256):
+    """Mostly copy-paste from https://github.com/IDEA-opensource/DAB-DETR/
+    """
+    half_hidden_dim = hidden_dim // 2
+    scale = 2 * math.pi
+    dim_t = torch.arange(half_hidden_dim, dtype=torch.float32, device=pos_tensor.device)
+    dim_t = 10000 ** (2 * (dim_t // 2) / half_hidden_dim)
+    x_embed = pos_tensor[..., 0] * scale
+    y_embed = pos_tensor[..., 1] * scale
+    pos_x = x_embed[..., None] / dim_t
+    pos_y = y_embed[..., None] / dim_t
+    pos_x = torch.stack((pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=-1).flatten(-2)
+    pos_y = torch.stack((pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=-1).flatten(-2)
+    pos = torch.cat((pos_y, pos_x), dim=-1)
+    return pos
+