custom_3d.py 5.77 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
                 modality=None,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
22
                 filter_empty_gt=True,
zhangwenwei's avatar
zhangwenwei committed
23
                 test_mode=False):
24
        super().__init__()
zhangwenwei's avatar
zhangwenwei committed
25
26
        self.data_root = data_root
        self.ann_file = ann_file
27
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
28
        self.modality = modality
wuyuefeng's avatar
Votenet  
wuyuefeng committed
29
        self.filter_empty_gt = filter_empty_gt
zhangwenwei's avatar
zhangwenwei committed
30
31
32

        self.CLASSES = self.get_classes(classes)
        self.data_infos = self.load_annotations(self.ann_file)
33
34
35
36

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

zhangwenwei's avatar
zhangwenwei committed
37
38
39
40
41
42
        # 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)
43
44
45
46

    def get_data_info(self, index):
        info = self.data_infos[index]
        sample_idx = info['point_cloud']['lidar_idx']
liyinhao's avatar
liyinhao committed
47
        pts_filename = osp.join(self.data_root, info['pts_path'])
48

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

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

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

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

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

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

liyinhao's avatar
liyinhao committed
87
88
89
90
91
92
        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
93
94
95

        Return:
            list[str]: return the list of class names
liyinhao's avatar
liyinhao committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
        """
        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
110
111
112
113
114
115
116
117
118
119
    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
120

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

        Evaluation in indoor protocol.

        Args:
liyinhao's avatar
liyinhao committed
127
            results (list[dict]): List of results.
wuyuefeng's avatar
wuyuefeng committed
128
129
            metric (str | list[str]): Metrics to be evaluated.
            iou_thr (list[float]): AP IoU thresholds.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
130

131
132
        """
        from mmdet3d.core.evaluation import indoor_eval
liyinhao's avatar
liyinhao committed
133
134
        assert isinstance(
            results, list), f'Expect results to be list, got {type(results)}.'
wuyuefeng's avatar
Votenet  
wuyuefeng committed
135
136
        assert len(results) > 0, f'Expect length of results > 0.'
        assert len(results) == len(self.data_infos)
liyinhao's avatar
liyinhao committed
137
138
139
        assert isinstance(
            results[0], dict
        ), f'Expect elements in results to be dict, got {type(results[0])}.'
140
        gt_annos = [info['annos'] for info in self.data_infos]
zhangwenwei's avatar
zhangwenwei committed
141
        label2cat = {i: cat_id for i, cat_id in enumerate(self.CLASSES)}
wuyuefeng's avatar
wuyuefeng committed
142
143
144
145
146
147
148
149
150
        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
151
        return ret_dict
zhangwenwei's avatar
zhangwenwei committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

    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)