scannet_data_utils.py 12.4 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import os
zhangwenwei's avatar
zhangwenwei committed
3
4
from concurrent import futures as futures
from os import path as osp
5

6
7
8
import mmcv
import numpy as np

9

liyinhao's avatar
liyinhao committed
10
class ScanNetData(object):
liyinhao's avatar
liyinhao committed
11
    """ScanNet data.
liyinhao's avatar
liyinhao committed
12

liyinhao's avatar
liyinhao committed
13
    Generate scannet infos for scannet_converter.
liyinhao's avatar
liyinhao committed
14
15

    Args:
liyinhao's avatar
liyinhao committed
16
        root_path (str): Root path of the raw data.
17
        split (str, optional): Set split type of the data. Default: 'train'.
liyinhao's avatar
liyinhao committed
18
    """
19
20
21
22

    def __init__(self, root_path, split='train'):
        self.root_dir = root_path
        self.split = split
liyinhao's avatar
liyinhao committed
23
        self.split_dir = osp.join(root_path)
24
25
26
27
28
29
30
31
32
        self.classes = [
            'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
            'bookshelf', 'picture', 'counter', 'desk', 'curtain',
            'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
            'garbagebin'
        ]
        self.cat2label = {cat: self.classes.index(cat) for cat in self.classes}
        self.label2cat = {self.cat2label[t]: t for t in self.cat2label}
        self.cat_ids = np.array(
33
            [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
34
        self.cat_ids2class = {
35
            nyu40id: i
36
            for i, nyu40id in enumerate(list(self.cat_ids))
37
38
        }
        assert split in ['train', 'val', 'test']
liyinhao's avatar
liyinhao committed
39
40
        split_file = osp.join(self.root_dir, 'meta_data',
                              f'scannetv2_{split}.txt')
41
42
        mmcv.check_file_exist(split_file)
        self.sample_id_list = mmcv.list_from_file(split_file)
43
        self.test_mode = (split == 'test')
44
45
46
47

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

48
    def get_aligned_box_label(self, idx):
49
        box_file = osp.join(self.root_dir, 'scannet_instance_data',
50
                            f'{idx}_aligned_bbox.npy')
liyinhao's avatar
liyinhao committed
51
        mmcv.check_file_exist(box_file)
52
53
        return np.load(box_file)

54
55
56
57
58
59
60
61
62
63
64
65
    def get_unaligned_box_label(self, idx):
        box_file = osp.join(self.root_dir, 'scannet_instance_data',
                            f'{idx}_unaligned_bbox.npy')
        mmcv.check_file_exist(box_file)
        return np.load(box_file)

    def get_axis_align_matrix(self, idx):
        matrix_file = osp.join(self.root_dir, 'scannet_instance_data',
                               f'{idx}_axis_align_matrix.npy')
        mmcv.check_file_exist(matrix_file)
        return np.load(matrix_file)

66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
    def get_images(self, idx):
        paths = []
        path = osp.join(self.root_dir, 'posed_images', idx)
        for file in sorted(os.listdir(path)):
            if file.endswith('.jpg'):
                paths.append(osp.join('posed_images', idx, file))
        return paths

    def get_extrinsics(self, idx):
        extrinsics = []
        path = osp.join(self.root_dir, 'posed_images', idx)
        for file in sorted(os.listdir(path)):
            if file.endswith('.txt') and not file == 'intrinsic.txt':
                extrinsics.append(np.loadtxt(osp.join(path, file)))
        return extrinsics

    def get_intrinsics(self, idx):
        matrix_file = osp.join(self.root_dir, 'posed_images', idx,
                               'intrinsic.txt')
        mmcv.check_file_exist(matrix_file)
        return np.loadtxt(matrix_file)

liyinhao's avatar
liyinhao committed
88
    def get_infos(self, num_workers=4, has_label=True, sample_id_list=None):
liyinhao's avatar
liyinhao committed
89
        """Get data infos.
liyinhao's avatar
liyinhao committed
90
91
92
93

        This method gets information from the raw data.

        Args:
94
95
96
97
98
            num_workers (int, optional): Number of threads to be used.
                Default: 4.
            has_label (bool, optional): Whether the data has label.
                Default: True.
            sample_id_list (list[int], optional): Index list of the sample.
liyinhao's avatar
liyinhao committed
99
                Default: None.
liyinhao's avatar
liyinhao committed
100
101

        Returns:
liyinhao's avatar
liyinhao committed
102
            infos (list[dict]): Information of the raw data.
liyinhao's avatar
liyinhao committed
103
        """
104
105

        def process_single_scene(sample_idx):
liyinhao's avatar
liyinhao committed
106
            print(f'{self.split} sample_idx: {sample_idx}')
107
108
109
            info = dict()
            pc_info = {'num_features': 6, 'lidar_idx': sample_idx}
            info['point_cloud'] = pc_info
110
            pts_filename = osp.join(self.root_dir, 'scannet_instance_data',
liyinhao's avatar
liyinhao committed
111
112
113
114
115
116
                                    f'{sample_idx}_vert.npy')
            points = np.load(pts_filename)
            mmcv.mkdir_or_exist(osp.join(self.root_dir, 'points'))
            points.tofile(
                osp.join(self.root_dir, 'points', f'{sample_idx}.bin'))
            info['pts_path'] = osp.join('points', f'{sample_idx}.bin')
117

118
119
120
121
122
123
124
125
126
127
128
129
130
131
            # update with RGB image paths if exist
            if os.path.exists(osp.join(self.root_dir, 'posed_images')):
                info['intrinsics'] = self.get_intrinsics(sample_idx)
                all_extrinsics = self.get_extrinsics(sample_idx)
                all_img_paths = self.get_images(sample_idx)
                # some poses in ScanNet are invalid
                extrinsics, img_paths = [], []
                for extrinsic, img_path in zip(all_extrinsics, all_img_paths):
                    if np.all(np.isfinite(extrinsic)):
                        img_paths.append(img_path)
                        extrinsics.append(extrinsic)
                info['extrinsics'] = extrinsics
                info['img_paths'] = img_paths

132
133
134
135
136
137
138
139
140
            if not self.test_mode:
                pts_instance_mask_path = osp.join(
                    self.root_dir, 'scannet_instance_data',
                    f'{sample_idx}_ins_label.npy')
                pts_semantic_mask_path = osp.join(
                    self.root_dir, 'scannet_instance_data',
                    f'{sample_idx}_sem_label.npy')

                pts_instance_mask = np.load(pts_instance_mask_path).astype(
WRH's avatar
WRH committed
141
                    np.int64)
142
                pts_semantic_mask = np.load(pts_semantic_mask_path).astype(
WRH's avatar
WRH committed
143
                    np.int64)
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

                mmcv.mkdir_or_exist(osp.join(self.root_dir, 'instance_mask'))
                mmcv.mkdir_or_exist(osp.join(self.root_dir, 'semantic_mask'))

                pts_instance_mask.tofile(
                    osp.join(self.root_dir, 'instance_mask',
                             f'{sample_idx}.bin'))
                pts_semantic_mask.tofile(
                    osp.join(self.root_dir, 'semantic_mask',
                             f'{sample_idx}.bin'))

                info['pts_instance_mask_path'] = osp.join(
                    'instance_mask', f'{sample_idx}.bin')
                info['pts_semantic_mask_path'] = osp.join(
                    'semantic_mask', f'{sample_idx}.bin')
159
160
161

            if has_label:
                annotations = {}
162
163
164
165
                # box is of shape [k, 6 + class]
                aligned_box_label = self.get_aligned_box_label(sample_idx)
                unaligned_box_label = self.get_unaligned_box_label(sample_idx)
                annotations['gt_num'] = aligned_box_label.shape[0]
166
                if annotations['gt_num'] != 0:
167
168
169
                    aligned_box = aligned_box_label[:, :-1]  # k, 6
                    unaligned_box = unaligned_box_label[:, :-1]
                    classes = aligned_box_label[:, -1]  # k
170
                    annotations['name'] = np.array([
171
                        self.label2cat[self.cat_ids2class[classes[i]]]
172
173
                        for i in range(annotations['gt_num'])
                    ])
174
175
176
177
178
179
180
181
182
                    # default names are given to aligned bbox for compatibility
                    # we also save unaligned bbox info with marked names
                    annotations['location'] = aligned_box[:, :3]
                    annotations['dimensions'] = aligned_box[:, 3:6]
                    annotations['gt_boxes_upright_depth'] = aligned_box
                    annotations['unaligned_location'] = unaligned_box[:, :3]
                    annotations['unaligned_dimensions'] = unaligned_box[:, 3:6]
                    annotations[
                        'unaligned_gt_boxes_upright_depth'] = unaligned_box
183
184
185
                    annotations['index'] = np.arange(
                        annotations['gt_num'], dtype=np.int32)
                    annotations['class'] = np.array([
186
                        self.cat_ids2class[classes[i]]
187
188
                        for i in range(annotations['gt_num'])
                    ])
189
190
                axis_align_matrix = self.get_axis_align_matrix(sample_idx)
                annotations['axis_align_matrix'] = axis_align_matrix  # 4x4
191
192
193
194
195
196
197
198
                info['annos'] = annotations
            return info

        sample_id_list = sample_id_list if sample_id_list is not None \
            else self.sample_id_list
        with futures.ThreadPoolExecutor(num_workers) as executor:
            infos = executor.map(process_single_scene, sample_id_list)
        return list(infos)
199
200
201
202
203
204
205
206


class ScanNetSegData(object):
    """ScanNet dataset used to generate infos for semantic segmentation task.

    Args:
        data_root (str): Root path of the raw data.
        ann_file (str): The generated scannet infos.
207
208
209
210
211
        split (str, optional): Set split type of the data. Default: 'train'.
        num_points (int, optional): Number of points in each data input.
            Default: 8192.
        label_weight_func (function, optional): Function to compute the
            label weight. Default: None.
212
213
214
215
216
217
218
219
220
221
222
    """

    def __init__(self,
                 data_root,
                 ann_file,
                 split='train',
                 num_points=8192,
                 label_weight_func=None):
        self.data_root = data_root
        self.data_infos = mmcv.load(ann_file)
        self.split = split
223
        assert split in ['train', 'val', 'test']
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
        self.num_points = num_points

        self.all_ids = np.arange(41)  # all possible ids
        self.cat_ids = np.array([
            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36,
            39
        ])  # used for seg task
        self.ignore_index = len(self.cat_ids)

        self.cat_id2class = np.ones((self.all_ids.shape[0],), dtype=np.int) * \
            self.ignore_index
        for i, cat_id in enumerate(self.cat_ids):
            self.cat_id2class[cat_id] = i

        # label weighting function is taken from
        # https://github.com/charlesq34/pointnet2/blob/master/scannet/scannet_dataset.py#L24
        self.label_weight_func = (lambda x: 1.0 / np.log(1.2 + x)) if \
            label_weight_func is None else label_weight_func

    def get_seg_infos(self):
244
245
        if self.split == 'test':
            return
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
        scene_idxs, label_weight = self.get_scene_idxs_and_label_weight()
        save_folder = osp.join(self.data_root, 'seg_info')
        mmcv.mkdir_or_exist(save_folder)
        np.save(
            osp.join(save_folder, f'{self.split}_resampled_scene_idxs.npy'),
            scene_idxs)
        np.save(
            osp.join(save_folder, f'{self.split}_label_weight.npy'),
            label_weight)
        print(f'{self.split} resampled scene index and label weight saved')

    def _convert_to_label(self, mask):
        """Convert class_id in loaded segmentation mask to label."""
        if isinstance(mask, str):
            if mask.endswith('npy'):
                mask = np.load(mask)
            else:
WRH's avatar
WRH committed
263
                mask = np.fromfile(mask, dtype=np.int64)
264
265
266
267
        label = self.cat_id2class[mask]
        return label

    def get_scene_idxs_and_label_weight(self):
268
        """Compute scene_idxs for data sampling and label weight for loss
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
        calculation.

        We sample more times for scenes with more points. Label_weight is
        inversely proportional to number of class points.
        """
        num_classes = len(self.cat_ids)
        num_point_all = []
        label_weight = np.zeros((num_classes + 1, ))  # ignore_index
        for data_info in self.data_infos:
            label = self._convert_to_label(
                osp.join(self.data_root, data_info['pts_semantic_mask_path']))
            num_point_all.append(label.shape[0])
            class_count, _ = np.histogram(label, range(num_classes + 2))
            label_weight += class_count

        # repeat scene_idx for num_scene_point // num_sample_point times
        sample_prob = np.array(num_point_all) / float(np.sum(num_point_all))
        num_iter = int(np.sum(num_point_all) / float(self.num_points))
        scene_idxs = []
        for idx in range(len(self.data_infos)):
289
            scene_idxs.extend([idx] * int(round(sample_prob[idx] * num_iter)))
290
291
292
293
294
295
296
297
        scene_idxs = np.array(scene_idxs).astype(np.int32)

        # calculate label weight, adopted from PointNet++
        label_weight = label_weight[:-1].astype(np.float32)
        label_weight = label_weight / label_weight.sum()
        label_weight = self.label_weight_func(label_weight).astype(np.float32)

        return scene_idxs, label_weight