lyft_metric.py 16.6 KB
Newer Older
VVsssssk's avatar
VVsssssk committed
1
2
3
4
5
6
# Copyright (c) OpenMMLab. All rights reserved.
import os
import tempfile
from os import path as osp
from typing import Dict, List, Optional, Sequence, Tuple, Union

7
import mmengine
VVsssssk's avatar
VVsssssk committed
8
9
10
11
import numpy as np
import pandas as pd
from lyft_dataset_sdk.lyftdataset import LyftDataset as Lyft
from lyft_dataset_sdk.utils.data_classes import Box as LyftBox
12
from mmengine import load
VVsssssk's avatar
VVsssssk committed
13
14
15
16
from mmengine.evaluator import BaseMetric
from mmengine.logging import MMLogger
from pyquaternion import Quaternion

zhangshilong's avatar
zhangshilong committed
17
from mmdet3d.evaluation import lyft_eval
VVsssssk's avatar
VVsssssk committed
18
19
20
21
22
23
24
25
26
27
from mmdet3d.registry import METRICS


@METRICS.register_module()
class LyftMetric(BaseMetric):
    """Lyft evaluation metric.

    Args:
        data_root (str): Path of dataset root.
        ann_file (str): Path of annotation file.
28
29
30
        metric (str or List[str]): Metrics to be evaluated. Defaults to 'bbox'.
        modality (dict): Modality to specify the sensor data used as input.
            Defaults to dict(use_camera=False, use_lidar=True).
VVsssssk's avatar
VVsssssk committed
31
32
        prefix (str, optional): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
            If prefix is not provided in the argument, self.default_prefix will
            be used instead. Defaults to None.
        jsonfile_prefix (str, optional): The prefix of json files including the
            file path and the prefix of filename, e.g., "a/b/prefix". If not
            specified, a temp file will be created. Defaults to None.
        format_only (bool): Format the output results without perform
            evaluation. It is useful when you want to format the result to a
            specific format and submit it to the test server.
            Defaults to False.
        csv_savepath (str, optional): The path for saving csv files. It
            includes the file path and the csv filename, e.g.,
            "a/b/filename.csv". If not specified, the result will not be
            converted to csv file. Defaults to None.
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
VVsssssk's avatar
VVsssssk committed
48
            'gpu'. Defaults to 'cpu'.
49
50
        backend_args (dict, optional): Arguments to instantiate the
            corresponding backend. Defaults to None.
VVsssssk's avatar
VVsssssk committed
51
52
    """

53
54
55
56
57
58
59
60
61
62
    def __init__(self,
                 data_root: str,
                 ann_file: str,
                 metric: Union[str, List[str]] = 'bbox',
                 modality=dict(
                     use_camera=False,
                     use_lidar=True,
                 ),
                 prefix: Optional[str] = None,
                 jsonfile_prefix: str = None,
63
                 format_only: bool = False,
64
65
66
                 csv_savepath: str = None,
                 collect_device: str = 'cpu',
                 backend_args: Optional[dict] = None) -> None:
VVsssssk's avatar
VVsssssk committed
67
68
69
70
71
72
73
        self.default_prefix = 'Lyft metric'
        super(LyftMetric, self).__init__(
            collect_device=collect_device, prefix=prefix)
        self.ann_file = ann_file
        self.data_root = data_root
        self.modality = modality
        self.jsonfile_prefix = jsonfile_prefix
74
75
76
77
78
79
        self.format_only = format_only
        if self.format_only:
            assert csv_savepath is not None, 'csv_savepath must be not None '
            'when format_only is True, otherwise the result files will be '
            'saved to a temp directory which will be cleaned up at the end.'

80
        self.backend_args = backend_args
VVsssssk's avatar
VVsssssk committed
81
82
83
        self.csv_savepath = csv_savepath
        self.metrics = metric if isinstance(metric, list) else [metric]

84
85
    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and data_samples.
VVsssssk's avatar
VVsssssk committed
86

87
88
        The processed results should be stored in ``self.results``, which will
        be used to compute the metrics when all batches have been processed.
VVsssssk's avatar
VVsssssk committed
89
90

        Args:
91
            data_batch (dict): A batch of data from the dataloader.
92
            data_samples (Sequence[dict]): A batch of outputs from the model.
VVsssssk's avatar
VVsssssk committed
93
        """
94
        for data_sample in data_samples:
VVsssssk's avatar
VVsssssk committed
95
            result = dict()
96
97
98
99
100
101
102
103
104
105
            pred_3d = data_sample['pred_instances_3d']
            pred_2d = data_sample['pred_instances']
            for attr_name in pred_3d:
                pred_3d[attr_name] = pred_3d[attr_name].to('cpu')
            result['pred_instances_3d'] = pred_3d
            for attr_name in pred_2d:
                pred_2d[attr_name] = pred_2d[attr_name].to('cpu')
            result['pred_instances'] = pred_2d
            sample_idx = data_sample['sample_idx']
            result['sample_idx'] = sample_idx
106
            self.results.append(result)
VVsssssk's avatar
VVsssssk committed
107

108
    def compute_metrics(self, results: List[dict]) -> Dict[str, float]:
VVsssssk's avatar
VVsssssk committed
109
110
111
        """Compute the metrics from processed results.

        Args:
112
            results (List[dict]): The processed results of the whole dataset.
VVsssssk's avatar
VVsssssk committed
113
114
115
116
117
118
119

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        logger: MMLogger = MMLogger.get_current_instance()

120
        classes = self.dataset_meta['classes']
VVsssssk's avatar
VVsssssk committed
121
122
        self.version = self.dataset_meta['version']

123
        # load annotations
124
        self.data_infos = load(
125
            self.ann_file, backend_args=self.backend_args)['data_list']
VVsssssk's avatar
VVsssssk committed
126
        result_dict, tmp_dir = self.format_results(results, classes,
127
128
                                                   self.jsonfile_prefix,
                                                   self.csv_savepath)
VVsssssk's avatar
VVsssssk committed
129
130

        metric_dict = {}
131
132
133
134
135
136

        if self.format_only:
            logger.info(
                f'results are saved in {osp.dirname(self.csv_savepath)}')
            return metric_dict

VVsssssk's avatar
VVsssssk committed
137
138
139
140
141
142
143
144
145
146
        for metric in self.metrics:
            ap_dict = self.lyft_evaluate(
                result_dict, metric=metric, logger=logger)
            for result in ap_dict:
                metric_dict[result] = ap_dict[result]

        if tmp_dir is not None:
            tmp_dir.cleanup()
        return metric_dict

147
148
149
150
151
152
153
    def format_results(
        self,
        results: List[dict],
        classes: Optional[List[str]] = None,
        jsonfile_prefix: Optional[str] = None,
        csv_savepath: Optional[str] = None
    ) -> Tuple[dict, Union[tempfile.TemporaryDirectory, None]]:
VVsssssk's avatar
VVsssssk committed
154
155
156
        """Format the results to json (standard format for COCO evaluation).

        Args:
157
158
159
            results (List[dict]): Testing results of the dataset.
            classes (List[str], optional): A list of class name.
                Defaults to None.
VVsssssk's avatar
VVsssssk committed
160
161
162
            jsonfile_prefix (str, optional): The prefix of json files. It
                includes the file path and the prefix of filename, e.g.,
                "a/b/prefix". If not specified, a temp file will be created.
163
164
165
166
167
                Defaults to None.
            csv_savepath (str, optional): The path for saving csv files. It
                includes the file path and the csv filename, e.g.,
                "a/b/filename.csv". If not specified, the result will not be
                converted to csv file. Defaults to None.
VVsssssk's avatar
VVsssssk committed
168
169

        Returns:
170
171
172
173
            tuple: Returns (result_dict, tmp_dir), where ``result_dict`` is a
            dict containing the json filepaths, ``tmp_dir`` is the temporal
            directory created for saving json files when ``jsonfile_prefix`` is
            not specified.
VVsssssk's avatar
VVsssssk committed
174
175
176
177
178
179
180
181
182
        """
        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
        result_dict = dict()
183
        sample_idx_list = [result['sample_idx'] for result in results]
VVsssssk's avatar
VVsssssk committed
184
185
186
187
188
189
190
191
192

        for name in results[0]:
            if 'pred' in name and '3d' in name and name[0] != '_':
                print(f'\nFormating bboxes of {name}')
                # format result of model output in Det3dDataSample,
                # include 'pred_instances_3d','pts_pred_instances_3d',
                # 'img_pred_instances_3d'
                results_ = [out[name] for out in results]
                tmp_file_ = osp.join(jsonfile_prefix, name)
193
194
195
                result_dict[name] = self._format_bbox(results_,
                                                      sample_idx_list, classes,
                                                      tmp_file_)
VVsssssk's avatar
VVsssssk committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
        if csv_savepath is not None:
            if 'pred_instances_3d' in result_dict:
                self.json2csv(result_dict['pred_instances_3d'], csv_savepath)
            elif 'pts_pred_instances_3d' in result_dict:
                self.json2csv(result_dict['pts_pred_instances_3d'],
                              csv_savepath)
        return result_dict, tmp_dir

    def json2csv(self, json_path: str, csv_savepath: str) -> None:
        """Convert the json file to csv format for submission.

        Args:
            json_path (str): Path of the result json file.
            csv_savepath (str): Path to save the csv file.
        """
211
        results = mmengine.load(json_path)['results']
VVsssssk's avatar
VVsssssk committed
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        sample_list_path = osp.join(self.data_root, 'sample_submission.csv')
        data = pd.read_csv(sample_list_path)
        Id_list = list(data['Id'])
        pred_list = list(data['PredictionString'])
        cnt = 0
        print('Converting the json to csv...')
        for token in results.keys():
            cnt += 1
            predictions = results[token]
            prediction_str = ''
            for i in range(len(predictions)):
                prediction_str += \
                    str(predictions[i]['score']) + ' ' + \
                    str(predictions[i]['translation'][0]) + ' ' + \
                    str(predictions[i]['translation'][1]) + ' ' + \
                    str(predictions[i]['translation'][2]) + ' ' + \
                    str(predictions[i]['size'][0]) + ' ' + \
                    str(predictions[i]['size'][1]) + ' ' + \
                    str(predictions[i]['size'][2]) + ' ' + \
                    str(Quaternion(list(predictions[i]['rotation']))
                        .yaw_pitch_roll[0]) + ' ' + \
                    predictions[i]['name'] + ' '
            prediction_str = prediction_str[:-1]
            idx = Id_list.index(token)
            pred_list[idx] = prediction_str
        df = pd.DataFrame({'Id': Id_list, 'PredictionString': pred_list})
238
        mmengine.mkdir_or_exist(os.path.dirname(csv_savepath))
VVsssssk's avatar
VVsssssk committed
239
240
241
242
        df.to_csv(csv_savepath, index=False)

    def _format_bbox(self,
                     results: List[dict],
243
244
245
                     sample_idx_list: List[int],
                     classes: Optional[List[str]] = None,
                     jsonfile_prefix: Optional[str] = None) -> str:
VVsssssk's avatar
VVsssssk committed
246
247
248
        """Convert the results to the standard format.

        Args:
249
250
251
252
            results (List[dict]): Testing results of the dataset.
            sample_idx_list (List[int]): List of result sample idx.
            classes (List[str], optional): A list of class name.
                Defaults to None.
VVsssssk's avatar
VVsssssk committed
253
            jsonfile_prefix (str, optional): The prefix of the output jsonfile.
254
255
                You can specify the output directory/filename by modifying the
                jsonfile_prefix. Defaults to None.
VVsssssk's avatar
VVsssssk committed
256
257
258
259
260
261
262

        Returns:
            str: Path of the output json file.
        """
        lyft_annos = {}

        print('Start to convert detection format...')
263
        for i, det in enumerate(mmengine.track_iter_progress(results)):
VVsssssk's avatar
VVsssssk committed
264
265
            annos = []
            boxes = output_to_lyft_box(det)
266
267
268
269
            sample_idx = sample_idx_list[i]
            sample_token = self.data_infos[sample_idx]['token']
            boxes = lidar_lyft_box_to_global(self.data_infos[sample_idx],
                                             boxes)
VVsssssk's avatar
VVsssssk committed
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
            for i, box in enumerate(boxes):
                name = classes[box.label]
                lyft_anno = dict(
                    sample_token=sample_token,
                    translation=box.center.tolist(),
                    size=box.wlh.tolist(),
                    rotation=box.orientation.elements.tolist(),
                    name=name,
                    score=box.score)
                annos.append(lyft_anno)
            lyft_annos[sample_token] = annos
        lyft_submissions = {
            'meta': self.modality,
            'results': lyft_annos,
        }

286
        mmengine.mkdir_or_exist(jsonfile_prefix)
VVsssssk's avatar
VVsssssk committed
287
288
        res_path = osp.join(jsonfile_prefix, 'results_lyft.json')
        print('Results writes to', res_path)
289
        mmengine.dump(lyft_submissions, res_path)
VVsssssk's avatar
VVsssssk committed
290
291
292
293
294
        return res_path

    def lyft_evaluate(self,
                      result_dict: dict,
                      metric: str = 'bbox',
295
                      logger: Optional[MMLogger] = None) -> Dict[str, float]:
VVsssssk's avatar
VVsssssk committed
296
297
298
299
        """Evaluation in Lyft protocol.

        Args:
            result_dict (dict): Formatted results of the dataset.
300
301
302
            metric (str): Metrics to be evaluated. Defaults to 'bbox'.
            logger (MMLogger, optional): Logger used for printing related
                information during evaluation. Defaults to None.
VVsssssk's avatar
VVsssssk committed
303
304

        Returns:
305
            Dict[str, float]: Evaluation results.
VVsssssk's avatar
VVsssssk committed
306
307
308
        """
        metric_dict = dict()
        for name in result_dict:
309
            print(f'Evaluating bboxes of {name}')
VVsssssk's avatar
VVsssssk committed
310
311
            ret_dict = self._evaluate_single(
                result_dict[name], logger=logger, result_name=name)
312
            metric_dict.update(ret_dict)
VVsssssk's avatar
VVsssssk committed
313
314
315
316
317
318
319
320
321
322
        return metric_dict

    def _evaluate_single(self,
                         result_path: str,
                         logger: MMLogger = None,
                         result_name: str = 'pts_bbox') -> dict:
        """Evaluation for a single model in Lyft protocol.

        Args:
            result_path (str): Path of the result file.
323
324
            logger (MMLogger, optional): Logger used for printing related
                information during evaluation. Defaults to None.
VVsssssk's avatar
VVsssssk committed
325
            result_name (str): Result name in the metric prefix.
326
                Defaults to 'pts_bbox'.
VVsssssk's avatar
VVsssssk committed
327
328

        Returns:
329
            Dict[str, float]: Dictionary of evaluation details.
VVsssssk's avatar
VVsssssk committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
        """
        output_dir = osp.join(*osp.split(result_path)[:-1])
        lyft = Lyft(
            data_path=osp.join(self.data_root, self.version),
            json_path=osp.join(self.data_root, self.version, self.version),
            verbose=True)
        eval_set_map = {
            'v1.01-train': 'val',
        }
        metrics = lyft_eval(lyft, self.data_root, result_path,
                            eval_set_map[self.version], output_dir, logger)

        # record metrics
        detail = dict()
        metric_prefix = f'{result_name}_Lyft'

        for i, name in enumerate(metrics['class_names']):
            AP = float(metrics['mAPs_cate'][i])
            detail[f'{metric_prefix}/{name}_AP'] = AP

        detail[f'{metric_prefix}/mAP'] = metrics['Final mAP']
        return detail


def output_to_lyft_box(detection: dict) -> List[LyftBox]:
    """Convert the output to the box class in the Lyft.

    Args:
        detection (dict): Detection results.

    Returns:
361
        List[:obj:`LyftBox`]: List of standard LyftBoxes.
VVsssssk's avatar
VVsssssk committed
362
    """
zhangshilong's avatar
zhangshilong committed
363
    bbox3d = detection['bbox_3d']
VVsssssk's avatar
VVsssssk committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
    scores = detection['scores_3d'].numpy()
    labels = detection['labels_3d'].numpy()

    box_gravity_center = bbox3d.gravity_center.numpy()
    box_dims = bbox3d.dims.numpy()
    box_yaw = bbox3d.yaw.numpy()

    # our LiDAR coordinate system -> Lyft box coordinate system
    lyft_box_dims = box_dims[:, [1, 0, 2]]

    box_list = []
    for i in range(len(bbox3d)):
        quat = Quaternion(axis=[0, 0, 1], radians=box_yaw[i])
        box = LyftBox(
            box_gravity_center[i],
            lyft_box_dims[i],
            quat,
            label=labels[i],
            score=scores[i])
        box_list.append(box)
    return box_list


def lidar_lyft_box_to_global(info: dict,
                             boxes: List[LyftBox]) -> List[LyftBox]:
    """Convert the box from ego to global coordinate.

    Args:
392
393
394
        info (dict): Info for a specific sample data, including the calibration
            information.
        boxes (List[:obj:`LyftBox`]): List of predicted LyftBoxes.
VVsssssk's avatar
VVsssssk committed
395
396

    Returns:
397
398
        List[:obj:`LyftBox`]: List of standard LyftBoxes in the global
        coordinate.
VVsssssk's avatar
VVsssssk committed
399
400
401
402
403
404
405
406
407
408
409
410
411
    """
    box_list = []
    for box in boxes:
        # Move box to ego vehicle coord system
        lidar2ego = np.array(info['lidar_points']['lidar2ego'])
        box.rotate(Quaternion(matrix=lidar2ego, rtol=1e-05, atol=1e-07))
        box.translate(lidar2ego[:3, 3])
        # Move box to global coord system
        ego2global = np.array(info['ego2global'])
        box.rotate(Quaternion(matrix=ego2global, rtol=1e-05, atol=1e-07))
        box.translate(ego2global[:3, 3])
        box_list.append(box)
    return box_list