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

import mmcv
import numpy as np
zhangwenwei's avatar
zhangwenwei committed
7
from scipy import io as sio
8
9


10
def random_sampling(points, num_points, replace=None, return_choices=False):
liyinhao's avatar
liyinhao committed
11
    """Random sampling.
liyinhao's avatar
liyinhao committed
12

13
    Sampling point cloud to a certain number of points.
liyinhao's avatar
liyinhao committed
14
15

    Args:
16
        points (ndarray): Point cloud.
17
        num_points (int): The number of samples.
liyinhao's avatar
liyinhao committed
18
19
20
21
        replace (bool): Whether the sample is with or without replacement.
        return_choices (bool): Whether to return choices.

    Returns:
22
        points (ndarray): Point cloud after sampling.
23
    """
liyinhao's avatar
liyinhao committed
24

25
    if replace is None:
26
27
        replace = (points.shape[0] < num_points)
    choices = np.random.choice(points.shape[0], num_points, replace=replace)
28
    if return_choices:
29
        return points[choices], choices
30
    else:
31
        return points[choices]
32
33


liyinhao's avatar
liyinhao committed
34
class SUNRGBDInstance(object):
35
36
37
38
39
40
41
42
43
44
45

    def __init__(self, line):
        data = line.split(' ')
        data[1:] = [float(x) for x in data[1:]]
        self.classname = data[0]
        self.xmin = data[1]
        self.ymin = data[2]
        self.xmax = data[1] + data[3]
        self.ymax = data[2] + data[4]
        self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])
        self.centroid = np.array([data[5], data[6], data[7]])
46
47
48
49
50
        self.width = data[8]
        self.length = data[9]
        self.height = data[10]
        # data[9] is x_size (length), data[8] is y_size (width), data[10] is
        # z_size (height) in our depth coordinate system,
51
52
        # l corresponds to the size along the x axis
        self.size = np.array([data[9], data[8], data[10]]) * 2
53
54
55
        self.orientation = np.zeros((3, ))
        self.orientation[0] = data[11]
        self.orientation[1] = data[12]
56
57
58
59
        self.heading_angle = np.arctan2(self.orientation[1],
                                        self.orientation[0])
        self.box3d = np.concatenate(
            [self.centroid, self.size, self.heading_angle[None]])
60
61


liyinhao's avatar
liyinhao committed
62
class SUNRGBDData(object):
liyinhao's avatar
liyinhao committed
63
    """SUNRGBD data.
liyinhao's avatar
liyinhao committed
64

liyinhao's avatar
liyinhao committed
65
    Generate scannet infos for sunrgbd_converter.
liyinhao's avatar
liyinhao committed
66
67
68

    Args:
        root_path (str): Root path of the raw data.
69
70
        split (str, optional): Set split type of the data. Default: 'train'.
        use_v1 (bool, optional): Whether to use v1. Default: False.
liyinhao's avatar
liyinhao committed
71
    """
72
73
74
75

    def __init__(self, root_path, split='train', use_v1=False):
        self.root_dir = root_path
        self.split = split
liyinhao's avatar
liyinhao committed
76
        self.split_dir = osp.join(root_path, 'sunrgbd_trainval')
77
78
79
80
81
82
83
        self.classes = [
            'bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
            'night_stand', 'bookshelf', 'bathtub'
        ]
        self.cat2label = {cat: self.classes.index(cat) for cat in self.classes}
        self.label2cat = {
            label: self.classes[label]
liyinhao's avatar
liyinhao committed
84
            for label in range(len(self.classes))
85
86
        }
        assert split in ['train', 'val', 'test']
liyinhao's avatar
liyinhao committed
87
        split_file = osp.join(self.split_dir, f'{split}_data_idx.txt')
88
89
        mmcv.check_file_exist(split_file)
        self.sample_id_list = map(int, mmcv.list_from_file(split_file))
liyinhao's avatar
liyinhao committed
90
91
92
        self.image_dir = osp.join(self.split_dir, 'image')
        self.calib_dir = osp.join(self.split_dir, 'calib')
        self.depth_dir = osp.join(self.split_dir, 'depth')
93
        if use_v1:
liyinhao's avatar
liyinhao committed
94
            self.label_dir = osp.join(self.split_dir, 'label_v1')
95
        else:
liyinhao's avatar
liyinhao committed
96
            self.label_dir = osp.join(self.split_dir, 'label')
97
98
99
100
101

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

    def get_image(self, idx):
liyinhao's avatar
liyinhao committed
102
        img_filename = osp.join(self.image_dir, f'{idx:06d}.jpg')
liyinhao's avatar
liyinhao committed
103
        return mmcv.imread(img_filename)
104
105
106
107
108
109

    def get_image_shape(self, idx):
        image = self.get_image(idx)
        return np.array(image.shape[:2], dtype=np.int32)

    def get_depth(self, idx):
liyinhao's avatar
liyinhao committed
110
        depth_filename = osp.join(self.depth_dir, f'{idx:06d}.mat')
111
112
113
114
        depth = sio.loadmat(depth_filename)['instance']
        return depth

    def get_calibration(self, idx):
liyinhao's avatar
liyinhao committed
115
        calib_filepath = osp.join(self.calib_dir, f'{idx:06d}.txt')
116
117
        lines = [line.rstrip() for line in open(calib_filepath)]
        Rt = np.array([float(x) for x in lines[0].split(' ')])
118
        Rt = np.reshape(Rt, (3, 3), order='F').astype(np.float32)
119
        K = np.array([float(x) for x in lines[1].split(' ')])
120
        K = np.reshape(K, (3, 3), order='F').astype(np.float32)
121
122
123
        return K, Rt

    def get_label_objects(self, idx):
liyinhao's avatar
liyinhao committed
124
        label_filename = osp.join(self.label_dir, f'{idx:06d}.txt')
125
        lines = [line.rstrip() for line in open(label_filename)]
liyinhao's avatar
liyinhao committed
126
        objects = [SUNRGBDInstance(line) for line in lines]
127
128
        return objects

liyinhao's avatar
liyinhao committed
129
    def get_infos(self, num_workers=4, has_label=True, sample_id_list=None):
liyinhao's avatar
liyinhao committed
130
        """Get data infos.
liyinhao's avatar
liyinhao committed
131
132
133
134

        This method gets information from the raw data.

        Args:
135
136
137
138
139
            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
140
                Default: None.
liyinhao's avatar
liyinhao committed
141
142

        Returns:
liyinhao's avatar
liyinhao committed
143
            infos (list[dict]): Information of the raw data.
liyinhao's avatar
liyinhao committed
144
        """
145
146

        def process_single_scene(sample_idx):
liyinhao's avatar
liyinhao committed
147
            print(f'{self.split} sample_idx: {sample_idx}')
148
            # convert depth to points
liyinhao's avatar
liyinhao committed
149
            SAMPLE_NUM = 50000
150
151
            # TODO: Check whether can move the point
            #  sampling process during training.
152
153
154
155
156
157
158
            pc_upright_depth = self.get_depth(sample_idx)
            pc_upright_depth_subsampled = random_sampling(
                pc_upright_depth, SAMPLE_NUM)

            info = dict()
            pc_info = {'num_features': 6, 'lidar_idx': sample_idx}
            info['point_cloud'] = pc_info
liyinhao's avatar
liyinhao committed
159
160
161
162
163
164

            mmcv.mkdir_or_exist(osp.join(self.root_dir, 'points'))
            pc_upright_depth_subsampled.tofile(
                osp.join(self.root_dir, 'points', f'{sample_idx:06d}.bin'))

            info['pts_path'] = osp.join('points', f'{sample_idx:06d}.bin')
165
            img_path = osp.join('image', f'{sample_idx:06d}.jpg')
166
167
            image_info = {
                'image_idx': sample_idx,
168
169
                'image_shape': self.get_image_shape(sample_idx),
                'image_path': img_path
170
171
172
173
174
175
176
177
178
179
180
181
            }
            info['image'] = image_info

            K, Rt = self.get_calibration(sample_idx)
            calib_info = {'K': K, 'Rt': Rt}
            info['calib'] = calib_info

            if has_label:
                obj_list = self.get_label_objects(sample_idx)
                annotations = {}
                annotations['gt_num'] = len([
                    obj.classname for obj in obj_list
182
                    if obj.classname in self.cat2label.keys()
183
184
185
186
                ])
                if annotations['gt_num'] != 0:
                    annotations['name'] = np.array([
                        obj.classname for obj in obj_list
187
                        if obj.classname in self.cat2label.keys()
188
189
190
                    ])
                    annotations['bbox'] = np.concatenate([
                        obj.box2d.reshape(1, 4) for obj in obj_list
191
                        if obj.classname in self.cat2label.keys()
192
193
194
195
                    ],
                                                         axis=0)
                    annotations['location'] = np.concatenate([
                        obj.centroid.reshape(1, 3) for obj in obj_list
196
                        if obj.classname in self.cat2label.keys()
197
198
199
                    ],
                                                             axis=0)
                    annotations['dimensions'] = 2 * np.array([
200
                        [obj.length, obj.width, obj.height] for obj in obj_list
201
                        if obj.classname in self.cat2label.keys()
Yezhen Cong's avatar
Yezhen Cong committed
202
                    ])  # lwh (depth) format
203
204
                    annotations['rotation_y'] = np.array([
                        obj.heading_angle for obj in obj_list
205
                        if obj.classname in self.cat2label.keys()
206
207
208
209
                    ])
                    annotations['index'] = np.arange(
                        len(obj_list), dtype=np.int32)
                    annotations['class'] = np.array([
210
211
                        self.cat2label[obj.classname] for obj in obj_list
                        if obj.classname in self.cat2label.keys()
212
213
214
215
                    ])
                    annotations['gt_boxes_upright_depth'] = np.stack(
                        [
                            obj.box3d for obj in obj_list
216
                            if obj.classname in self.cat2label.keys()
217
218
219
220
221
222
223
224
225
226
                        ],
                        axis=0)  # (K,8)
                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)