indoor_augment.py 7.95 KB
Newer Older
liyinhao's avatar
liyinhao committed
1
2
import numpy as np

3
from mmdet.datasets.builder import PIPELINES
liyinhao's avatar
liyinhao committed
4
5


6
@PIPELINES.register_module()
7
class IndoorFlipData(object):
liyinhao's avatar
liyinhao committed
8
    """Indoor Flip Data.
9

liyinhao's avatar
liyinhao committed
10
11
12
    Flip point cloud and ground truth boxes.
    The point cloud will ve flipped along the yz plane
    and the xz plane with a certain probability.
13
14

    Args:
liyinhao's avatar
liyinhao committed
15
16
17
        flip_ratio_yz (float): Probability of being flipped along yz plane.
            Default: 0.5.
        flip_ratio_xz (float): Probability of being flipped along xz plane.
liyinhao's avatar
liyinhao committed
18
            Default: 0.5.
19
20
    """

liyinhao's avatar
liyinhao committed
21
22
23
    def __init__(self, flip_ratio_yz=0.5, flip_ratio_xz=0.5):
        self.flip_ratio_yz = flip_ratio_yz
        self.flip_ratio_xz = flip_ratio_xz
24
25

    def __call__(self, results):
26
27
        points = results['points']
        gt_bboxes_3d = results['gt_bboxes_3d']
liyinhao's avatar
liyinhao committed
28
29
        aligned = True if gt_bboxes_3d.shape[1] == 6 else False
        if np.random.random() > self.flip_ratio_yz:
30
            # Flipping along the YZ plane
liyinhao's avatar
liyinhao committed
31
32
            points[:, 0] = -1 * points[:, 0]
            gt_bboxes_3d[:, 0] = -1 * gt_bboxes_3d[:, 0]
liyinhao's avatar
liyinhao committed
33
            if not aligned:
liyinhao's avatar
liyinhao committed
34
                gt_bboxes_3d[:, 6] = np.pi - gt_bboxes_3d[:, 6]
liyinhao's avatar
liyinhao committed
35
36

            results['flip'] = True
liyinhao's avatar
liyinhao committed
37
            results['gt_boxes'] = gt_bboxes_3d
38

liyinhao's avatar
liyinhao committed
39
        if aligned and np.random.random() > self.flip_ratio_xz:
40
            # Flipping along the XZ plane
liyinhao's avatar
liyinhao committed
41
42
            points[:, 1] = -1 * points[:, 1]
            gt_bboxes_3d[:, 1] = -1 * gt_bboxes_3d[:, 1]
liyinhao's avatar
liyinhao committed
43
            results['flip'] = True
liyinhao's avatar
liyinhao committed
44
45
            results['gt_bboxes_3d'] = gt_bboxes_3d
        results['points'] = points
46
47
48
49
50

        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
liyinhao's avatar
liyinhao committed
51
52
        repr_str += '(flip_ratio_yz={})'.format(self.flip_ratio_yz)
        repr_str += '(flip_ratio_xz={})'.format(self.flip_ratio_xz)
53
54
55
        return repr_str


liyinhao's avatar
liyinhao committed
56
@PIPELINES.register_module()
liyinhao's avatar
liyinhao committed
57
58
class IndoorPointsColorJitter(object):
    """Indoor Points Color Jitter.
liyinhao's avatar
liyinhao committed
59

liyinhao's avatar
liyinhao committed
60
61
    Randomly change the brightness and color of the point cloud, and
    drop out the points' colors with a certain range and probability.
liyinhao's avatar
liyinhao committed
62
63
64
65
66
67
68
69
70
71

    Args:
        color_mean (List[float]): Mean color of the point cloud.
            Default: [0.5, 0.5, 0.5].
        bright_range (List[float]): Range of brightness.
            Default: [0.8, 1.2].
        color_shift_range (List[float]): Range of color shift.
            Default: [0.95, 1.05].
        jitter_range (List[float]): Range of jittering.
            Default: [-0.025, 0.025].
72
        drop_prob (float): Probability to drop out points' color.
liyinhao's avatar
liyinhao committed
73
74
75
76
77
78
79
80
            Default: 0.3
    """

    def __init__(self,
                 color_mean=[0.5, 0.5, 0.5],
                 bright_range=[0.8, 1.2],
                 color_shift_range=[0.95, 1.05],
                 jitter_range=[-0.025, 0.025],
81
                 drop_prob=0.3):
liyinhao's avatar
liyinhao committed
82
83
84
85
        self.color_mean = color_mean
        self.bright_range = bright_range
        self.color_shift_range = color_shift_range
        self.jitter_range = jitter_range
86
        self.drop_prob = drop_prob
liyinhao's avatar
liyinhao committed
87
88

    def __call__(self, results):
89
        points = results['points']
liyinhao's avatar
liyinhao committed
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        assert points.shape[1] >= 6
        rgb_color = points[:, 3:6] + self.color_mean
        # brightness change for each channel
        rgb_color *= np.random.uniform(self.bright_range[0],
                                       self.bright_range[1], 3)
        # color shift for each channel
        rgb_color += np.random.uniform(self.color_shift_range[0],
                                       self.color_shift_range[1], 3)
        # jittering on each pixel
        rgb_color += np.expand_dims(
            np.random.uniform(self.jitter_range[0], self.jitter_range[1]), -1)
        rgb_color = np.clip(rgb_color, 0, 1)
        # randomly drop out points' colors
        rgb_color *= np.expand_dims(
104
            np.random.random(points.shape[0]) > self.drop_prob, -1)
liyinhao's avatar
liyinhao committed
105
106
107
108
109
110
111
112
113
114
        points[:, 3:6] = rgb_color - self.color_mean
        results['points'] = points
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += '(color_mean={})'.format(self.color_mean)
        repr_str += '(bright_range={})'.format(self.bright_range)
        repr_str += '(color_shift_range={})'.format(self.color_shift_range)
        repr_str += '(jitter_range={})'.format(self.jitter_range)
115
        repr_str += '(drop_prob={})'.format(self.drop_prob)
liyinhao's avatar
liyinhao committed
116
117


118
119
120
121
122
123
124
# TODO: merge outdoor indoor transform.
# TODO: try transform noise.
@PIPELINES.register_module()
class IndoorGlobalRotScale(object):
    """Indoor Global Rotate Scale.

    Augment sunrgbd and scannet data with global rotating and scaling.
125
126

    Args:
liyinhao's avatar
liyinhao committed
127
        use_height (bool): Whether to use height.
liyinhao's avatar
liyinhao committed
128
129
130
131
132
            Default: True.
        rot_range (List[float]): Range of rotation.
            Default: None.
        scale_range (List[float]): Range of scale.
            Default: None.
133
134
    """

liyinhao's avatar
liyinhao committed
135
    def __init__(self, use_height=True, rot_range=None, scale_range=None):
liyinhao's avatar
liyinhao committed
136
137
138
139
        self.use_height = use_height
        self.rot_range = rot_range
        self.scale_range = scale_range

liyinhao's avatar
liyinhao committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
    def _rotz(self, t):
        """Rotate About Z.

        Rotation about the z-axis.

        Args:
            t (float): Angle of rotation.

        Returns:
            rot_mat (ndarray): Matrix of rotation.
        """
        c = np.cos(t)
        s = np.sin(t)
        rot_mat = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
        return rot_mat

liyinhao's avatar
liyinhao committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
    def _rotate_aligned_boxes(self, input_boxes, rot_mat):
        """Rotate Aligned Boxes.

        Rotate function for the aligned boxes.

        Args:
            input_boxes (ndarray): 3D boxes.
            rot_mat (ndarray): Rotation matrix.

        Returns:
            rotated_boxes (ndarry): 3D boxes after rotation.
        """
        centers, lengths = input_boxes[:, 0:3], input_boxes[:, 3:6]
        new_centers = np.dot(centers, np.transpose(rot_mat))

        dx, dy = lengths[:, 0] / 2.0, lengths[:, 1] / 2.0
        new_x = np.zeros((dx.shape[0], 4))
        new_y = np.zeros((dx.shape[0], 4))

        for i, crnr in enumerate([(-1, -1), (1, -1), (1, 1), (-1, 1)]):
            crnrs = np.zeros((dx.shape[0], 3))
            crnrs[:, 0] = crnr[0] * dx
            crnrs[:, 1] = crnr[1] * dy
            crnrs = np.dot(crnrs, np.transpose(rot_mat))
            new_x[:, i] = crnrs[:, 0]
            new_y[:, i] = crnrs[:, 1]

        new_dx = 2.0 * np.max(new_x, 1)
        new_dy = 2.0 * np.max(new_y, 1)
        new_lengths = np.stack((new_dx, new_dy, lengths[:, 2]), axis=1)

        return np.concatenate([new_centers, new_lengths], axis=1)
188
189

    def __call__(self, results):
190
191
        points = results['points']
        gt_bboxes_3d = results['gt_bboxes_3d']
liyinhao's avatar
liyinhao committed
192
        name = 'scannet' if gt_bboxes_3d.shape[1] == 6 else 'sunrgbd'
liyinhao's avatar
liyinhao committed
193

liyinhao's avatar
liyinhao committed
194
        if self.rot_range is not None:
195
            assert len(self.rot_range) == 2
liyinhao's avatar
liyinhao committed
196
            rot_angle = np.random.uniform(self.rot_range[0], self.rot_range[1])
liyinhao's avatar
liyinhao committed
197
198
            rot_mat = self._rotz(rot_angle)
            points[:, :3] = np.dot(points[:, :3], rot_mat.T)
liyinhao's avatar
liyinhao committed
199
            if name == 'scannet':
liyinhao's avatar
liyinhao committed
200
201
                gt_bboxes_3d = self._rotate_aligned_boxes(
                    gt_bboxes_3d, rot_mat)
liyinhao's avatar
liyinhao committed
202
            else:
liyinhao's avatar
liyinhao committed
203
                gt_bboxes_3d[:, :3] = np.dot(gt_bboxes_3d[:, :3], rot_mat.T)
liyinhao's avatar
liyinhao committed
204
                gt_bboxes_3d[:, 6] -= rot_angle
liyinhao's avatar
liyinhao committed
205

liyinhao's avatar
liyinhao committed
206
        if self.scale_range is not None:
207
            assert len(self.scale_range) == 2
liyinhao's avatar
liyinhao committed
208
            # Augment point cloud scale
liyinhao's avatar
liyinhao committed
209
210
            scale_ratio = np.random.uniform(self.scale_range[0],
                                            self.scale_range[1])
211

liyinhao's avatar
liyinhao committed
212
213
            points[:, :3] *= scale_ratio
            gt_bboxes_3d[:, :3] *= scale_ratio
liyinhao's avatar
liyinhao committed
214
            gt_bboxes_3d[:, 3:6] *= scale_ratio
liyinhao's avatar
liyinhao committed
215
            if self.use_height:
216
                points[:, -1] *= scale_ratio
liyinhao's avatar
liyinhao committed
217

liyinhao's avatar
liyinhao committed
218
219
        results['points'] = points
        results['gt_bboxes_3d'] = gt_bboxes_3d
220
221
222
223
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
liyinhao's avatar
liyinhao committed
224
        repr_str += '(use_height={})'.format(self.use_height)
liyinhao's avatar
liyinhao committed
225
226
        repr_str += '(rot_range={})'.format(self.rot_range)
        repr_str += '(scale_range={})'.format(self.scale_range)
227
        return repr_str