custom_3d.py 5.49 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
3
import os.path as osp
import tempfile

4
5
import mmcv
import numpy as np
wuyuefeng's avatar
wuyuefeng committed
6
from mmcv.utils import print_log
zhangwenwei's avatar
zhangwenwei committed
7
from torch.utils.data import Dataset
8
9
10
11
12
13

from mmdet.datasets import DATASETS
from .pipelines import Compose


@DATASETS.register_module()
zhangwenwei's avatar
zhangwenwei committed
14
class Custom3DDataset(Dataset):
15
16

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
17
                 data_root,
18
19
                 ann_file,
                 pipeline=None,
liyinhao's avatar
liyinhao committed
20
                 classes=None,
zhangwenwei's avatar
zhangwenwei committed
21
22
                 modality=None,
                 test_mode=False):
23
        super().__init__()
zhangwenwei's avatar
zhangwenwei committed
24
25
        self.data_root = data_root
        self.ann_file = ann_file
26
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
27
28
29
30
        self.modality = modality

        self.CLASSES = self.get_classes(classes)
        self.data_infos = self.load_annotations(self.ann_file)
31
32
33
34

        if pipeline is not None:
            self.pipeline = Compose(pipeline)

zhangwenwei's avatar
zhangwenwei committed
35
36
37
38
39
40
        # set group flag for the sampler
        if not self.test_mode:
            self._set_group_flag()

    def load_annotations(self, ann_file):
        return mmcv.load(ann_file)
41
42
43
44
45
46

    def get_data_info(self, index):
        info = self.data_infos[index]
        sample_idx = info['point_cloud']['lidar_idx']
        pts_filename = self._get_pts_filename(sample_idx)

liyinhao's avatar
liyinhao committed
47
48
49
50
        input_dict = dict(
            pts_filename=pts_filename,
            sample_idx=sample_idx,
            file_name=pts_filename)
51

zhangwenwei's avatar
zhangwenwei committed
52
        if not self.test_mode:
liyinhao's avatar
liyinhao committed
53
            annos = self.get_ann_info(index)
zhangwenwei's avatar
zhangwenwei committed
54
55
56
            input_dict['ann_info'] = annos
            if len(annos['gt_bboxes_3d']) == 0:
                return None
57
58
        return input_dict

zhangwenwei's avatar
zhangwenwei committed
59
60
61
62
    def pre_pipeline(self, results):
        results['bbox3d_fields'] = []
        results['pts_mask_fields'] = []
        results['pts_seg_fields'] = []
63

liyinhao's avatar
liyinhao committed
64
65
    def prepare_train_data(self, index):
        input_dict = self.get_data_info(index)
66
67
        if input_dict is None:
            return None
zhangwenwei's avatar
zhangwenwei committed
68
        self.pre_pipeline(input_dict)
69
        example = self.pipeline(input_dict)
zhangwenwei's avatar
zhangwenwei committed
70
        if example is None or len(example['gt_bboxes_3d']._data) == 0:
71
72
73
            return None
        return example

74
75
    def prepare_test_data(self, index):
        input_dict = self.get_data_info(index)
zhangwenwei's avatar
zhangwenwei committed
76
        self.pre_pipeline(input_dict)
77
78
        example = self.pipeline(input_dict)
        return example
79

liyinhao's avatar
liyinhao committed
80
81
    @classmethod
    def get_classes(cls, classes=None):
82
83
        """Get class names of current dataset.

liyinhao's avatar
liyinhao committed
84
85
86
87
88
89
        Args:
            classes (Sequence[str] | str | None): If classes is None, use
                default CLASSES defined by builtin dataset. If classes is a
                string, take it as a file name. The file contains the name of
                classes where each line contains one class name. If classes is
                a tuple or list, override the CLASSES defined by the dataset.
zhangwenwei's avatar
zhangwenwei committed
90
91
92

        Return:
            list[str]: return the list of class names
liyinhao's avatar
liyinhao committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
        """
        if classes is None:
            return cls.CLASSES

        if isinstance(classes, str):
            # take it as a file path
            class_names = mmcv.list_from_file(classes)
        elif isinstance(classes, (tuple, list)):
            class_names = classes
        else:
            raise ValueError(f'Unsupported type {type(classes)} of classes.')

        return class_names

liyinhao's avatar
liyinhao committed
107
108
109
110
111
112
113
114
115
116
    def format_results(self,
                       outputs,
                       pklfile_prefix=None,
                       submission_prefix=None):
        if pklfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            pklfile_prefix = osp.join(tmp_dir.name, 'results')
            out = f'{pklfile_prefix}.pkl'
        mmcv.dump(outputs, out)
        return outputs, tmp_dir
117

wuyuefeng's avatar
wuyuefeng committed
118
    def evaluate(self, results, metric=None, iou_thr=(0.25, 0.5), logger=None):
119
120
121
122
123
        """Evaluate.

        Evaluation in indoor protocol.

        Args:
liyinhao's avatar
liyinhao committed
124
            results (list[dict]): List of results.
wuyuefeng's avatar
wuyuefeng committed
125
126
            metric (str | list[str]): Metrics to be evaluated.
            iou_thr (list[float]): AP IoU thresholds.
127
128
        """
        from mmdet3d.core.evaluation import indoor_eval
liyinhao's avatar
liyinhao committed
129
130
131
132
133
        assert isinstance(
            results, list), f'Expect results to be list, got {type(results)}.'
        assert isinstance(
            results[0], dict
        ), f'Expect elements in results to be dict, got {type(results[0])}.'
134
        gt_annos = [info['annos'] for info in self.data_infos]
zhangwenwei's avatar
zhangwenwei committed
135
        label2cat = {i: cat_id for i, cat_id in enumerate(self.CLASSES)}
wuyuefeng's avatar
wuyuefeng committed
136
137
138
139
140
141
142
143
144
        ret_dict = indoor_eval(gt_annos, results, iou_thr, label2cat)

        result_str = str()
        for key, val in ret_dict.items():
            result_str += f'{key} : {val} \n'
        mAP_25, mAP_50 = ret_dict['mAP_0.25'], ret_dict['mAP_0.50']
        result_str += f'mAP(0.25): {mAP_25}    mAP(0.50): {mAP_50}'
        print_log('\n' + result_str, logger=logger)

liyinhao's avatar
liyinhao committed
145
        return ret_dict
zhangwenwei's avatar
zhangwenwei committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172

    def __len__(self):
        return len(self.data_infos)

    def _rand_another(self, idx):
        pool = np.where(self.flag == self.flag[idx])[0]
        return np.random.choice(pool)

    def __getitem__(self, idx):
        if self.test_mode:
            return self.prepare_test_data(idx)
        while True:
            data = self.prepare_train_data(idx)
            if data is None:
                idx = self._rand_another(idx)
                continue
            return data

    def _set_group_flag(self):
        """Set flag according to image aspect ratio.

        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        In 3D datasets, they are all the same, thus are all zeros

        """
        self.flag = np.zeros(len(self), dtype=np.uint8)