# nuScenes dev-kit.
# Code written by Holger Caesar & Oscar Beijbom, 2018.
# Licensed under the Creative Commons [see licence.txt]

import argparse
import json
import os
import random
import time
from typing import Tuple, Dict, Any

import numpy as np

from nuscenes import NuScenes
from nuscenes.eval.detection.algo import accumulate, calc_ap, calc_tp
from nuscenes.eval.detection.config import config_factory
from nuscenes.eval.detection.constants import TP_METRICS
from nuscenes.eval.detection.data_classes import DetectionConfig, MetricDataList, DetectionMetrics, EvalBoxes
from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes
from nuscenes.eval.detection.render import summary_plot, class_pr_curve, class_tp_curve, dist_pr_curve, visualize_sample


class NuScenesEval:
    """
    This is the official nuScenes detection evaluation code.
    Results are written to the provided output_dir.

    nuScenes uses the following 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 an 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://github.com/nutonomy/nuscenes-devkit for more details.
    """
    def __init__(self,
                 nusc: NuScenes,
                 config: DetectionConfig,
                 result_path: str,
                 eval_set: str,
                 output_dir: str = None,
                 verbose: bool = True):
        """
        Initialize a NuScenesEval 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 or val.
        :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

        # 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.
        self.pred_boxes, self.meta = load_prediction(self.result_path, self.cfg.max_boxes_per_sample, verbose=verbose)
        self.gt_boxes = load_gt(self.nusc, self.eval_set, verbose=verbose)

        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, MetricDataList]:
        """
        Performs the actual evaluation.
        :return: A tuple of high-level and the raw metric data.
        """

        start_time = time.time()

        # -----------------------------------
        # Step 1: Accumulate metric data for all classes and distance thresholds.
        # -----------------------------------
        if self.verbose:
            print('Accumulating metric data')
        metric_data_list = MetricDataList()
        for class_name in self.cfg.class_names:
            for dist_th in self.cfg.dist_ths:
                md = accumulate(self.gt_boxes, self.pred_boxes, class_name, self.cfg.dist_fcn, dist_th)
                metric_data_list.set(class_name, dist_th, md)

        # -----------------------------------
        # Step 2: Calculate metrics from the data.
        # -----------------------------------
        if self.verbose:
            print('Calculating metrics')
        metrics = DetectionMetrics(self.cfg)
        for class_name in self.cfg.class_names:
            for dist_th in self.cfg.dist_ths:
                metric_data = metric_data_list[(class_name, dist_th)]
                ap = calc_ap(metric_data, self.cfg.min_recall, self.cfg.min_precision)
                metrics.add_label_ap(class_name, dist_th, ap)

            for metric_name in TP_METRICS:
                metric_data = metric_data_list[(class_name, self.cfg.dist_th_tp)]
                if class_name in ['traffic_cone'] and metric_name in ['attr_err', 'vel_err', 'orient_err']:
                    tp = np.nan
                elif class_name in ['barrier'] and metric_name in ['attr_err', 'vel_err']:
                    tp = np.nan
                else:
                    tp = calc_tp(metric_data, self.cfg.min_recall, metric_name)
                metrics.add_label_tp(class_name, metric_name, tp)

        metrics.add_runtime(time.time() - start_time)

        return metrics, metric_data_list

    def render(self, metrics: DetectionMetrics, md_list: MetricDataList) -> None:
        """
        Renders various PR and TP curves.
        :param metrics: DetectionMetrics instance.
        :param md_list: MetricDataList instance.
        """

        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(43)
            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()

        # Render PR and TP curves.
        if render_curves:
            self.render(metrics, metric_data_list)

        # Dump the metric data, meta and metrics to disk.
        if self.verbose:
            print('Saving metrics to: %s' % self.output_dir)
        metrics_summary = metrics.serialize()
        metrics_summary['meta'] = self.meta.copy()
        with open(os.path.join(self.output_dir, 'metrics_summary.json'), 'w') as f:
            json.dump(metrics_summary, f, indent=2)
        with open(os.path.join(self.output_dir, 'metrics_details.json'), 'w') as f:
            json.dump(metric_data_list.serialize(), f, indent=2)

        # Print high-level metrics.
        print('mAP: %.4f' % (metrics_summary['mean_ap']))
        err_name_mapping = {
            'trans_err': 'mATE',
            'scale_err': 'mASE',
            'orient_err': 'mAOE',
            'vel_err': 'mAVE',
            'attr_err': 'mAAE'
        }
        for tp_name, tp_val in metrics_summary['tp_errors'].items():
            print('%s: %.4f' % (err_name_mapping[tp_name], tp_val))
        print('NDS: %.4f' % (metrics_summary['nd_score']))
        print('Eval time: %.1fs' % metrics_summary['eval_time'])

        return metrics_summary


if __name__ == "__main__":

    # Settings.
    parser = argparse.ArgumentParser(description='Evaluate nuScenes result submission.',
                                     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_name', type=str, default='cvpr_2019',
                        help='Name of the configuration to use for evaluation, e.g. cvpr_2019.')
    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_name_ = args.config_name
    plot_examples_ = args.plot_examples
    render_curves_ = bool(args.render_curves)
    verbose_ = bool(args.verbose)

    cfg_ = config_factory(config_name_)
    nusc_ = NuScenes(version=version_, verbose=verbose_, dataroot=dataroot_)
    nusc_eval = NuScenesEval(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_)

[((230, 13, 231, 92), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((251, 19, 251, 55), 'os.path.expanduser', 'os.path.expanduser', ({(251, 38, 251, 54): 'args.result_path'}, {}), '(args.result_path)', False, 'import os\n'), ((252, 18, 252, 53), 'os.path.expanduser', 'os.path.expanduser', ({(252, 37, 252, 52): 'args.output_dir'}, {}), '(args.output_dir)', False, 'import os\n'), ((261, 11, 261, 39), 'nuscenes.eval.detection.config.config_factory', 'config_factory', ({(261, 26, 261, 38): 'config_name_'}, {}), '(config_name_)', False, 'from nuscenes.eval.detection.config import config_factory\n'), ((262, 12, 262, 76), 'nuscenes.NuScenes', 'NuScenes', (), '', False, 'from nuscenes import NuScenes\n'), ((67, 24, 67, 62), 'os.path.join', 'os.path.join', ({(67, 37, 67, 52): 'self.output_dir', (67, 54, 67, 61): '"""plots"""'}, {}), "(self.output_dir, 'plots')", False, 'import os\n'), ((74, 37, 74, 118), 'nuscenes.eval.detection.loaders.load_prediction', 'load_prediction', (), '', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((75, 24, 75, 74), 'nuscenes.eval.detection.loaders.load_gt', 'load_gt', (), '', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((81, 26, 81, 64), 'nuscenes.eval.detection.loaders.add_center_dist', 'add_center_dist', ({(81, 42, 81, 46): 'nusc', (81, 48, 81, 63): 'self.pred_boxes'}, {}), '(nusc, self.pred_boxes)', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((82, 24, 82, 60), 'nuscenes.eval.detection.loaders.add_center_dist', 'add_center_dist', ({(82, 40, 82, 44): 'nusc', (82, 46, 82, 59): 'self.gt_boxes'}, {}), '(nusc, self.gt_boxes)', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((87, 26, 87, 105), 'nuscenes.eval.detection.loaders.filter_eval_boxes', 'filter_eval_boxes', (), '', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((90, 24, 90, 101), 'nuscenes.eval.detection.loaders.filter_eval_boxes', 'filter_eval_boxes', (), '', False, 'from nuscenes.eval.detection.loaders import load_prediction, load_gt, add_center_dist, filter_eval_boxes\n'), ((100, 21, 100, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((107, 27, 107, 43), 'nuscenes.eval.detection.data_classes.MetricDataList', 'MetricDataList', ({}, {}), '()', False, 'from nuscenes.eval.detection.data_classes import DetectionConfig, MetricDataList, DetectionMetrics, EvalBoxes\n'), ((118, 18, 118, 44), 'nuscenes.eval.detection.data_classes.DetectionMetrics', 'DetectionMetrics', ({(118, 35, 118, 43): 'self.cfg'}, {}), '(self.cfg)', False, 'from nuscenes.eval.detection.data_classes import DetectionConfig, MetricDataList, DetectionMetrics, EvalBoxes\n'), ((68, 15, 68, 45), 'os.path.isdir', 'os.path.isdir', ({(68, 29, 68, 44): 'self.output_dir'}, {}), '(self.output_dir)', False, 'import os\n'), ((69, 12, 69, 40), 'os.makedirs', 'os.makedirs', ({(69, 24, 69, 39): 'self.output_dir'}, {}), '(self.output_dir)', False, 'import os\n'), ((70, 15, 70, 43), 'os.path.isdir', 'os.path.isdir', ({(70, 29, 70, 42): 'self.plot_dir'}, {}), '(self.plot_dir)', False, 'import os\n'), ((71, 12, 71, 38), 'os.makedirs', 'os.makedirs', ({(71, 24, 71, 37): 'self.plot_dir'}, {}), '(self.plot_dir)', False, 'import os\n'), ((147, 19, 147, 61), 'os.path.join', 'os.path.join', ({(147, 32, 147, 45): 'self.plot_dir', (147, 47, 147, 60): "(name + '.pdf')"}, {}), "(self.plot_dir, name + '.pdf')", False, 'import os\n'), ((175, 12, 175, 27), 'random.seed', 'random.seed', ({(175, 24, 175, 26): '(43)'}, {}), '(43)', False, 'import random\n'), ((177, 12, 177, 41), 'random.shuffle', 'random.shuffle', ({(177, 27, 177, 40): 'sample_tokens'}, {}), '(sample_tokens)', False, 'import random\n'), ((181, 26, 181, 67), 'os.path.join', 'os.path.join', ({(181, 39, 181, 54): 'self.output_dir', (181, 56, 181, 66): '"""examples"""'}, {}), "(self.output_dir, 'examples')", False, 'import os\n'), ((206, 12, 206, 51), 'json.dump', 'json.dump', (), '', False, 'import json\n'), ((110, 21, 110, 103), 'nuscenes.eval.detection.algo.accumulate', 'accumulate', ({(110, 32, 110, 45): 'self.gt_boxes', (110, 47, 110, 62): 'self.pred_boxes', (110, 64, 110, 74): 'class_name', (110, 76, 110, 93): 'self.cfg.dist_fcn', (110, 95, 110, 102): 'dist_th'}, {}), '(self.gt_boxes, self.pred_boxes, class_name, self.cfg.dist_fcn,\n    dist_th)', False, 'from nuscenes.eval.detection.algo import accumulate, calc_ap, calc_tp\n'), ((122, 21, 122, 86), 'nuscenes.eval.detection.algo.calc_ap', 'calc_ap', ({(122, 29, 122, 40): 'metric_data', (122, 42, 122, 61): 'self.cfg.min_recall', (122, 63, 122, 85): 'self.cfg.min_precision'}, {}), '(metric_data, self.cfg.min_recall, self.cfg.min_precision)', False, 'from nuscenes.eval.detection.algo import accumulate, calc_ap, calc_tp\n'), ((135, 28, 135, 39), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((182, 19, 182, 45), 'os.path.isdir', 'os.path.isdir', ({(182, 33, 182, 44): 'example_dir'}, {}), '(example_dir)', False, 'import os\n'), ((183, 16, 183, 37), 'os.mkdir', 'os.mkdir', ({(183, 25, 183, 36): 'example_dir'}, {}), '(example_dir)', False, 'import os\n'), ((205, 18, 205, 71), 'os.path.join', 'os.path.join', ({(205, 31, 205, 46): 'self.output_dir', (205, 48, 205, 70): '"""metrics_summary.json"""'}, {}), "(self.output_dir, 'metrics_summary.json')", False, 'import os\n'), ((207, 18, 207, 71), 'os.path.join', 'os.path.join', ({(207, 31, 207, 46): 'self.output_dir', (207, 48, 207, 70): '"""metrics_details.json"""'}, {}), "(self.output_dir, 'metrics_details.json')", False, 'import os\n'), ((132, 25, 132, 79), 'nuscenes.eval.detection.algo.calc_tp', 'calc_tp', ({(132, 33, 132, 44): 'metric_data', (132, 46, 132, 65): 'self.cfg.min_recall', (132, 67, 132, 78): 'metric_name'}, {}), '(metric_data, self.cfg.min_recall, metric_name)', False, 'from nuscenes.eval.detection.algo import accumulate, calc_ap, calc_tp\n'), ((187, 79, 187, 90), 'nuscenes.eval.detection.data_classes.EvalBoxes', 'EvalBoxes', ({}, {}), '()', False, 'from nuscenes.eval.detection.data_classes import DetectionConfig, MetricDataList, DetectionMetrics, EvalBoxes\n')]