"aten/THC/THCGreedy.h" did not exist on "b1589f14c72ef813ecd961a301c1865b58684b4c"
indoor_metric.py 6.54 KB
Newer Older
jshilong's avatar
jshilong committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangshilong's avatar
zhangshilong committed
2
from collections import OrderedDict
3
from typing import Dict, List, Optional, Sequence, Union
jshilong's avatar
jshilong committed
4

zhangshilong's avatar
zhangshilong committed
5
import numpy as np
6
from mmdet.evaluation import eval_map
jshilong's avatar
jshilong committed
7
8
9
from mmengine.evaluator import BaseMetric
from mmengine.logging import MMLogger

zhangshilong's avatar
zhangshilong committed
10
from mmdet3d.evaluation import indoor_eval
jshilong's avatar
jshilong committed
11
from mmdet3d.registry import METRICS
zhangshilong's avatar
zhangshilong committed
12
from mmdet3d.structures import get_box_type
jshilong's avatar
jshilong committed
13
14
15
16


@METRICS.register_module()
class IndoorMetric(BaseMetric):
zhangshilong's avatar
zhangshilong committed
17
    """Indoor scene evaluation metric.
jshilong's avatar
jshilong committed
18
19

    Args:
20
21
22
23
24
25
        iou_thr (float or List[float]): List of iou threshold when calculate
            the metric. Defaults to [0.25, 0.5].
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
            'gpu'. Defaults to 'cpu'.
        prefix (str, optional): The prefix that will be added in the metric
jshilong's avatar
jshilong committed
26
            names to disambiguate homonymous metrics of different evaluators.
27
28
            If prefix is not provided in the argument, self.default_prefix will
            be used instead. Defaults to None.
jshilong's avatar
jshilong committed
29
30
31
32
33
    """

    def __init__(self,
                 iou_thr: List[float] = [0.25, 0.5],
                 collect_device: str = 'cpu',
34
                 prefix: Optional[str] = None) -> None:
jshilong's avatar
jshilong committed
35
36
        super(IndoorMetric, self).__init__(
            prefix=prefix, collect_device=collect_device)
37
        self.iou_thr = [iou_thr] if isinstance(iou_thr, float) else iou_thr
jshilong's avatar
jshilong committed
38

39
    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
jshilong's avatar
jshilong committed
40
41
        """Process one batch of data samples and predictions.

42
43
        The processed results should be stored in ``self.results``, which will
        be used to compute the metrics when all batches have been processed.
jshilong's avatar
jshilong committed
44
45

        Args:
46
            data_batch (dict): A batch of data from the dataloader.
47
            data_samples (Sequence[dict]): A batch of outputs from the model.
jshilong's avatar
jshilong committed
48
        """
49
50
51
        for data_sample in data_samples:
            pred_3d = data_sample['pred_instances_3d']
            eval_ann_info = data_sample['eval_ann_info']
jshilong's avatar
jshilong committed
52
53
54
55
56
57
            cpu_pred_3d = dict()
            for k, v in pred_3d.items():
                if hasattr(v, 'to'):
                    cpu_pred_3d[k] = v.to('cpu')
                else:
                    cpu_pred_3d[k] = v
58
            self.results.append((eval_ann_info, cpu_pred_3d))
jshilong's avatar
jshilong committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

    def compute_metrics(self, results: list) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.

        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()
        ann_infos = []
        pred_results = []

        for eval_ann, sinlge_pred_results in results:
            ann_infos.append(eval_ann)
            pred_results.append(sinlge_pred_results)

78
        # some checkpoints may not record the key "box_type_3d"
jshilong's avatar
jshilong committed
79
        box_type_3d, box_mode_3d = get_box_type(
80
            self.dataset_meta.get('box_type_3d', 'depth'))
jshilong's avatar
jshilong committed
81
82
83
84
85

        ret_dict = indoor_eval(
            ann_infos,
            pred_results,
            self.iou_thr,
86
            self.dataset_meta['classes'],
jshilong's avatar
jshilong committed
87
88
89
90
            logger=logger,
            box_mode_3d=box_mode_3d)

        return ret_dict
zhangshilong's avatar
zhangshilong committed
91
92
93
94
95
96
97


@METRICS.register_module()
class Indoor2DMetric(BaseMetric):
    """indoor 2d predictions evaluation metric.

    Args:
98
99
100
101
102
103
        iou_thr (float or List[float]): List of iou threshold when calculate
            the metric. Defaults to [0.5].
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be 'cpu' or
            'gpu'. Defaults to 'cpu'.
        prefix (str, optional): The prefix that will be added in the metric
zhangshilong's avatar
zhangshilong committed
104
            names to disambiguate homonymous metrics of different evaluators.
105
106
            If prefix is not provided in the argument, self.default_prefix will
            be used instead. Defaults to None.
zhangshilong's avatar
zhangshilong committed
107
108
109
    """

    def __init__(self,
110
                 iou_thr: Union[float, List[float]] = [0.5],
zhangshilong's avatar
zhangshilong committed
111
                 collect_device: str = 'cpu',
112
                 prefix: Optional[str] = None):
zhangshilong's avatar
zhangshilong committed
113
114
        super(Indoor2DMetric, self).__init__(
            prefix=prefix, collect_device=collect_device)
115
        self.iou_thr = [iou_thr] if isinstance(iou_thr, float) else iou_thr
zhangshilong's avatar
zhangshilong committed
116

117
    def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
zhangshilong's avatar
zhangshilong committed
118
119
        """Process one batch of data samples and predictions.

120
121
        The processed results should be stored in ``self.results``, which will
        be used to compute the metrics when all batches have been processed.
zhangshilong's avatar
zhangshilong committed
122
123

        Args:
124
            data_batch (dict): A batch of data from the dataloader.
125
            data_samples (Sequence[dict]): A batch of outputs from the model.
zhangshilong's avatar
zhangshilong committed
126
        """
127
128
129
        for data_sample in data_samples:
            pred = data_sample['pred_instances']
            eval_ann_info = data_sample['eval_ann_info']
zhangshilong's avatar
zhangshilong committed
130
            ann = dict(
131
132
                labels=eval_ann_info['gt_bboxes_labels'],
                bboxes=eval_ann_info['gt_bboxes'])
zhangshilong's avatar
zhangshilong committed
133
134
135
136
137
138

            pred_bboxes = pred['bboxes'].cpu().numpy()
            pred_scores = pred['scores'].cpu().numpy()
            pred_labels = pred['labels'].cpu().numpy()

            dets = []
139
            for label in range(len(self.dataset_meta['classes'])):
zhangshilong's avatar
zhangshilong committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
                index = np.where(pred_labels == label)[0]
                pred_bbox_scores = np.hstack(
                    [pred_bboxes[index], pred_scores[index].reshape((-1, 1))])
                dets.append(pred_bbox_scores)

            self.results.append((ann, dets))

    def compute_metrics(self, results: list) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.

        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()
        annotations, preds = zip(*results)
        eval_results = OrderedDict()
160
        for iou_thr_2d_single in self.iou_thr:
zhangshilong's avatar
zhangshilong committed
161
162
163
164
165
            mean_ap, _ = eval_map(
                preds,
                annotations,
                scale_ranges=None,
                iou_thr=iou_thr_2d_single,
166
                dataset=self.dataset_meta['classes'],
zhangshilong's avatar
zhangshilong committed
167
168
169
                logger=logger)
            eval_results['mAP_' + str(iou_thr_2d_single)] = mean_ap
        return eval_results