indoor_loading.py 8.49 KB
Newer Older
1
import mmcv
liyinhao's avatar
liyinhao committed
2
3
import numpy as np

4
from mmdet.datasets.builder import PIPELINES
zhangwenwei's avatar
zhangwenwei committed
5
from mmdet.datasets.pipelines import LoadAnnotations
liyinhao's avatar
liyinhao committed
6
7


8
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
9
10
class NormalizePointsColor(object):
    """Normalize color of points
11
12
13
14

    Normalize color of the points.

    Args:
liyinhao's avatar
liyinhao committed
15
        color_mean (list[float]): Mean color of the point cloud.
16
17
18
19
20
21
    """

    def __init__(self, color_mean):
        self.color_mean = color_mean

    def __call__(self, results):
liyinhao's avatar
liyinhao committed
22
        points = results['points']
liyinhao's avatar
liyinhao committed
23
24
        assert points.shape[1] >= 6,\
            f'Expect points have channel >=6, got {points.shape[1]}'
25
26
27
28
        points[:, 3:6] = points[:, 3:6] - np.array(self.color_mean) / 256.0
        results['points'] = points
        return results

liyinhao's avatar
liyinhao committed
29
30
31
32
33
    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += '(color_mean={})'.format(self.color_mean)
        return repr_str

34
35

@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
36
37
class LoadPointsFromFile(object):
    """Load Points From File.
liyinhao's avatar
liyinhao committed
38

liyinhao's avatar
liyinhao committed
39
    Load sunrgbd and scannet points from file.
liyinhao's avatar
liyinhao committed
40
41

    Args:
zhangwenwei's avatar
zhangwenwei committed
42
        shift_height (bool): Whether to use shifted height.
liyinhao's avatar
liyinhao committed
43
44
        load_dim (int): The dimension of the loaded points.
            Default: 6.
liyinhao's avatar
liyinhao committed
45
        use_dim (list[int]): Which dimensions of the points to be used.
zhangwenwei's avatar
zhangwenwei committed
46
47
            Default: [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension
zhangwenwei's avatar
zhangwenwei committed
48
49
50
        file_client_args (dict): Config dict of file clients, refer to
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
            for more details.
liyinhao's avatar
liyinhao committed
51
52
    """

zhangwenwei's avatar
zhangwenwei committed
53
54
55
56
57
    def __init__(self,
                 load_dim=6,
                 use_dim=[0, 1, 2],
                 shift_height=False,
                 file_client_args=dict(backend='disk')):
zhangwenwei's avatar
zhangwenwei committed
58
59
60
        self.shift_height = shift_height
        if isinstance(use_dim, int):
            use_dim = list(range(use_dim))
liyinhao's avatar
liyinhao committed
61
62
63
        assert max(use_dim) < load_dim, \
            f'Expect all used dimensions < {load_dim}, got {use_dim}'

liyinhao's avatar
liyinhao committed
64
65
        self.load_dim = load_dim
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
66
67
        self.file_client_args = file_client_args.copy()
        self.file_client = None
liyinhao's avatar
liyinhao committed
68

zhangwenwei's avatar
zhangwenwei committed
69
    def _load_points(self, pts_filename):
zhangwenwei's avatar
zhangwenwei committed
70
71
72
73
74
75
76
77
78
79
80
        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
            mmcv.check_file_exist(pts_filename)
            if pts_filename.endswith('.npy'):
                points = np.load(pts_filename)
            else:
                points = np.fromfile(pts_filename, dtype=np.float32)
zhangwenwei's avatar
zhangwenwei committed
81
82
        return points

liyinhao's avatar
liyinhao committed
83
    def __call__(self, results):
liyinhao's avatar
liyinhao committed
84
        pts_filename = results['pts_filename']
zhangwenwei's avatar
zhangwenwei committed
85
        points = self._load_points(pts_filename)
liyinhao's avatar
liyinhao committed
86
87
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
liyinhao's avatar
liyinhao committed
88

zhangwenwei's avatar
zhangwenwei committed
89
        if self.shift_height:
liyinhao's avatar
liyinhao committed
90
91
92
93
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
            points = np.concatenate([points, np.expand_dims(height, 1)], 1)
        results['points'] = points
liyinhao's avatar
liyinhao committed
94
95
96
97
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
zhangwenwei's avatar
zhangwenwei committed
98
        repr_str += '(shift_height={})'.format(self.shift_height)
liyinhao's avatar
liyinhao committed
99
        repr_str += '(mean_color={})'.format(self.color_mean)
liyinhao's avatar
liyinhao committed
100
101
        repr_str += '(load_dim={})'.format(self.load_dim)
        repr_str += '(use_dim={})'.format(self.use_dim)
liyinhao's avatar
liyinhao committed
102
103
        return repr_str

liyinhao's avatar
liyinhao committed
104

105
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
106
107
class LoadAnnotations3D(LoadAnnotations):
    """Load Annotations3D.
liyinhao's avatar
liyinhao committed
108

zhangwenwei's avatar
zhangwenwei committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
    Load instance mask and semantic mask of points and
    encapsulate the items into related fields.

    Args:
        with_bbox_3d (bool, optional): Whether to load 3D boxes.
            Defaults to True.
        with_label_3d (bool, optional): Whether to load 3D labels.
            Defaults to True.
        with_mask_3d (bool, optional): Whether to load 3D instance masks.
            for points. Defaults to False.
        with_seg_3d (bool, optional): Whether to load 3D semantic masks.
            for points. Defaults to False.
        with_bbox (bool, optional): Whether to load 2D boxes.
            Defaults to False.
        with_label (bool, optional): Whether to load 2D labels.
            Defaults to False.
        with_mask (bool, optional): Whether to load 2D instance masks.
            Defaults to False.
        with_seg (bool, optional): Whether to load 2D semantic masks.
            Defaults to False.
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
zhangwenwei's avatar
zhangwenwei committed
131
132
133
        file_client_args (dict): Config dict of file clients, refer to
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
            for more details.
liyinhao's avatar
liyinhao committed
134
135
    """

zhangwenwei's avatar
zhangwenwei committed
136
137
138
139
140
141
142
143
144
    def __init__(self,
                 with_bbox_3d=True,
                 with_label_3d=True,
                 with_mask_3d=False,
                 with_seg_3d=False,
                 with_bbox=False,
                 with_label=False,
                 with_mask=False,
                 with_seg=False,
zhangwenwei's avatar
zhangwenwei committed
145
146
147
148
149
150
151
152
153
                 poly2mask=True,
                 file_client_args=dict(backend='disk')):
        super().__init__(
            with_bbox,
            with_label,
            with_mask,
            with_seg,
            poly2mask,
            file_client_args=file_client_args)
zhangwenwei's avatar
zhangwenwei committed
154
155
156
157
158
159
160
        self.with_bbox_3d = with_bbox_3d
        self.with_label_3d = with_label_3d
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d

    def _load_bboxes_3d(self, results):
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
161
        results['bbox3d_fields'].append('gt_bboxes_3d')
zhangwenwei's avatar
zhangwenwei committed
162
        return results
liyinhao's avatar
liyinhao committed
163

zhangwenwei's avatar
zhangwenwei committed
164
165
166
    def _load_labels_3d(self, results):
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results
liyinhao's avatar
liyinhao committed
167

zhangwenwei's avatar
zhangwenwei committed
168
169
    def _load_masks_3d(self, results):
        pts_instance_mask_path = results['ann_info']['pts_instance_mask_path']
zhangwenwei's avatar
zhangwenwei committed
170
171
172
173
174
175
176
177
178
179
180

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            mask_bytes = self.file_client.get(pts_instance_mask_path)
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int)
        except ConnectionError:
            mmcv.check_file_exist(pts_instance_mask_path)
            pts_instance_mask = np.fromfile(
                pts_instance_mask_path, dtype=np.long)

liyinhao's avatar
liyinhao committed
181
        results['pts_instance_mask'] = pts_instance_mask
zhangwenwei's avatar
zhangwenwei committed
182
        results['pts_mask_fields'].append('pts_instance_mask')
zhangwenwei's avatar
zhangwenwei committed
183
184
185
186
        return results

    def _load_semantic_seg_3d(self, results):
        pts_semantic_mask_path = results['ann_info']['pts_semantic_mask_path']
zhangwenwei's avatar
zhangwenwei committed
187
188
189
190
191

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            mask_bytes = self.file_client.get(pts_semantic_mask_path)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
192
193
            # add .copy() to fix read-only bug
            pts_semantic_mask = np.frombuffer(mask_bytes, dtype=np.int).copy()
zhangwenwei's avatar
zhangwenwei committed
194
195
196
197
198
        except ConnectionError:
            mmcv.check_file_exist(pts_semantic_mask_path)
            pts_semantic_mask = np.fromfile(
                pts_semantic_mask_path, dtype=np.long)

liyinhao's avatar
liyinhao committed
199
        results['pts_semantic_mask'] = pts_semantic_mask
zhangwenwei's avatar
zhangwenwei committed
200
        results['pts_seg_fields'].append('pts_semantic_mask')
zhangwenwei's avatar
zhangwenwei committed
201
202
203
204
205
206
207
208
209
210
211
212
213
214
        return results

    def __call__(self, results):
        results = super().__call__(results)
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
            if results is None:
                return None
        if self.with_label_3d:
            results = self._load_labels_3d(results)
        if self.with_mask_3d:
            results = self._load_masks_3d(results)
        if self.with_seg_3d:
            results = self._load_semantic_seg_3d(results)
liyinhao's avatar
liyinhao committed
215
216
217
218

        return results

    def __repr__(self):
zhangwenwei's avatar
zhangwenwei committed
219
220
221
222
223
224
225
226
227
228
229
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d},\n'
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d},\n'
        repr_str += f'{indent_str}with_mask_3d={self.with_mask_3d},\n'
        repr_str += f'{indent_str}with_seg_3d={self.with_seg_3d},\n'
        repr_str += f'{indent_str}with_bbox={self.with_bbox},\n'
        repr_str += f'{indent_str}with_label={self.with_label},\n'
        repr_str += f'{indent_str}with_mask={self.with_mask},\n'
        repr_str += f'{indent_str}with_seg={self.with_seg},\n'
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
liyinhao's avatar
liyinhao committed
230
        return repr_str