voxel_encoder.py 20.3 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
import torch
3
from mmcv.cnn import build_norm_layer
4
from mmcv.ops import DynamicScatter
5
from mmcv.runner import force_fp32
zhangwenwei's avatar
zhangwenwei committed
6
7
from torch import nn

8
from mmdet3d.registry import MODELS
zhangwenwei's avatar
zhangwenwei committed
9
from .. import builder
zhangwenwei's avatar
zhangwenwei committed
10
from .utils import VFELayer, get_paddings_indicator
zhangwenwei's avatar
zhangwenwei committed
11
12


13
@MODELS.register_module()
zhangwenwei's avatar
zhangwenwei committed
14
class HardSimpleVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
15
    """Simple voxel feature encoder used in SECOND.
zhangwenwei's avatar
zhangwenwei committed
16

zhangwenwei's avatar
zhangwenwei committed
17
    It simply averages the values of points in a voxel.
18
19

    Args:
20
        num_features (int, optional): Number of features to use. Default: 4.
zhangwenwei's avatar
zhangwenwei committed
21
    """
zhangwenwei's avatar
zhangwenwei committed
22

23
    def __init__(self, num_features=4):
zhangwenwei's avatar
zhangwenwei committed
24
        super(HardSimpleVFE, self).__init__()
25
        self.num_features = num_features
26
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
27

28
    @force_fp32(out_fp16=True)
zhangwenwei's avatar
zhangwenwei committed
29
    def forward(self, features, num_points, coors):
zhangwenwei's avatar
zhangwenwei committed
30
        """Forward function.
zhangwenwei's avatar
zhangwenwei committed
31
32

        Args:
wangtai's avatar
wangtai committed
33
            features (torch.Tensor): Point features in shape
zhangwenwei's avatar
zhangwenwei committed
34
35
36
37
38
39
40
41
42
                (N, M, 3(4)). N is the number of voxels and M is the maximum
                number of points inside a single voxel.
            num_points (torch.Tensor): Number of points in each voxel,
                 shape (N, ).
            coors (torch.Tensor): Coordinates of voxels.

        Returns:
            torch.Tensor: Mean of points inside each voxel in shape (N, 3(4))
        """
43
        points_mean = features[:, :, :self.num_features].sum(
zhangwenwei's avatar
zhangwenwei committed
44
45
46
47
            dim=1, keepdim=False) / num_points.type_as(features).view(-1, 1)
        return points_mean.contiguous()


48
@MODELS.register_module()
zhangwenwei's avatar
zhangwenwei committed
49
class DynamicSimpleVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
50
    """Simple dynamic voxel feature encoder used in DV-SECOND.
zhangwenwei's avatar
zhangwenwei committed
51
52
53
54
55
56
57
58

    It simply averages the values of points in a voxel.
    But the number of points in a voxel is dynamic and varies.

    Args:
        voxel_size (tupe[float]): Size of a single voxel
        point_cloud_range (tuple[float]): Range of the point cloud and voxels
    """
zhangwenwei's avatar
zhangwenwei committed
59
60
61
62

    def __init__(self,
                 voxel_size=(0.2, 0.2, 4),
                 point_cloud_range=(0, -40, -3, 70.4, 40, 1)):
zhangwenwei's avatar
zhangwenwei committed
63
        super(DynamicSimpleVFE, self).__init__()
zhangwenwei's avatar
zhangwenwei committed
64
        self.scatter = DynamicScatter(voxel_size, point_cloud_range, True)
65
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
66
67

    @torch.no_grad()
68
    @force_fp32(out_fp16=True)
zhangwenwei's avatar
zhangwenwei committed
69
    def forward(self, features, coors):
zhangwenwei's avatar
zhangwenwei committed
70
        """Forward function.
zhangwenwei's avatar
zhangwenwei committed
71
72

        Args:
wangtai's avatar
wangtai committed
73
            features (torch.Tensor): Point features in shape
zhangwenwei's avatar
zhangwenwei committed
74
75
76
77
78
79
80
                (N, 3(4)). N is the number of points.
            coors (torch.Tensor): Coordinates of voxels.

        Returns:
            torch.Tensor: Mean of points inside each voxel in shape (M, 3(4)).
                M is the number of voxels.
        """
zhangwenwei's avatar
zhangwenwei committed
81
82
83
84
85
86
        # This function is used from the start of the voxelnet
        # num_points: [concated_num_points]
        features, features_coors = self.scatter(features, coors)
        return features, features_coors


87
@MODELS.register_module()
zhangwenwei's avatar
zhangwenwei committed
88
class DynamicVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
89
    """Dynamic Voxel feature encoder used in DV-SECOND.
zhangwenwei's avatar
zhangwenwei committed
90
91
92
93
94
95

    It encodes features of voxels and their points. It could also fuse
    image feature into voxel features in a point-wise manner.
    The number of points inside the voxel varies.

    Args:
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
        in_channels (int, optional): Input channels of VFE. Defaults to 4.
        feat_channels (list(int), optional): Channels of features in VFE.
        with_distance (bool, optional): Whether to use the L2 distance of
            points to the origin point. Defaults to False.
        with_cluster_center (bool, optional): Whether to use the distance
            to cluster center of points inside a voxel. Defaults to False.
        with_voxel_center (bool, optional): Whether to use the distance
            to center of voxel for each points inside a voxel.
            Defaults to False.
        voxel_size (tuple[float], optional): Size of a single voxel.
            Defaults to (0.2, 0.2, 4).
        point_cloud_range (tuple[float], optional): The range of points
            or voxels. Defaults to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict, optional): Config dict of normalization layers.
        mode (str, optional): The mode when pooling features of points
            inside a voxel. Available options include 'max' and 'avg'.
            Defaults to 'max'.
        fusion_layer (dict, optional): The config dict of fusion
            layer used in multi-modal detectors. Defaults to None.
        return_point_feats (bool, optional): Whether to return the features
            of each points. Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
117
    """
zhangwenwei's avatar
zhangwenwei committed
118
119

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
120
121
                 in_channels=4,
                 feat_channels=[],
zhangwenwei's avatar
zhangwenwei committed
122
123
124
125
126
127
128
129
130
131
                 with_distance=False,
                 with_cluster_center=False,
                 with_voxel_center=False,
                 voxel_size=(0.2, 0.2, 4),
                 point_cloud_range=(0, -40, -3, 70.4, 40, 1),
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 mode='max',
                 fusion_layer=None,
                 return_point_feats=False):
        super(DynamicVFE, self).__init__()
zhangwenwei's avatar
zhangwenwei committed
132
133
        assert mode in ['avg', 'max']
        assert len(feat_channels) > 0
zhangwenwei's avatar
zhangwenwei committed
134
        if with_cluster_center:
zhangwenwei's avatar
zhangwenwei committed
135
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
136
        if with_voxel_center:
zhangwenwei's avatar
zhangwenwei committed
137
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
138
        if with_distance:
139
            in_channels += 1
zhangwenwei's avatar
zhangwenwei committed
140
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
141
142
143
144
        self._with_distance = with_distance
        self._with_cluster_center = with_cluster_center
        self._with_voxel_center = with_voxel_center
        self.return_point_feats = return_point_feats
145
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
146
147
148
149
150
151
152
153
154
155
156

        # Need pillar (voxel) size and x/y offset in order to calculate offset
        self.vx = voxel_size[0]
        self.vy = voxel_size[1]
        self.vz = voxel_size[2]
        self.x_offset = self.vx / 2 + point_cloud_range[0]
        self.y_offset = self.vy / 2 + point_cloud_range[1]
        self.z_offset = self.vz / 2 + point_cloud_range[2]
        self.point_cloud_range = point_cloud_range
        self.scatter = DynamicScatter(voxel_size, point_cloud_range, True)

zhangwenwei's avatar
zhangwenwei committed
157
        feat_channels = [self.in_channels] + list(feat_channels)
zhangwenwei's avatar
zhangwenwei committed
158
        vfe_layers = []
zhangwenwei's avatar
zhangwenwei committed
159
160
161
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
zhangwenwei's avatar
zhangwenwei committed
162
163
164
165
166
167
168
            if i > 0:
                in_filters *= 2
            norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)
            vfe_layers.append(
                nn.Sequential(
                    nn.Linear(in_filters, out_filters, bias=False), norm_layer,
                    nn.ReLU(inplace=True)))
169
        self.vfe_layers = nn.ModuleList(vfe_layers)
zhangwenwei's avatar
zhangwenwei committed
170
171
172
173
174
175
176
177
178
179
        self.num_vfe = len(vfe_layers)
        self.vfe_scatter = DynamicScatter(voxel_size, point_cloud_range,
                                          (mode != 'max'))
        self.cluster_scatter = DynamicScatter(
            voxel_size, point_cloud_range, average_points=True)
        self.fusion_layer = None
        if fusion_layer is not None:
            self.fusion_layer = builder.build_fusion_layer(fusion_layer)

    def map_voxel_center_to_point(self, pts_coors, voxel_mean, voxel_coors):
zhangwenwei's avatar
zhangwenwei committed
180
181
182
183
184
185
186
187
188
189
        """Map voxel features to its corresponding points.

        Args:
            pts_coors (torch.Tensor): Voxel coordinate of each point.
            voxel_mean (torch.Tensor): Voxel features to be mapped.
            voxel_coors (torch.Tensor): Coordinates of valid voxels

        Returns:
            torch.Tensor: Features or centers of each point.
        """
zhangwenwei's avatar
zhangwenwei committed
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
        # Step 1: scatter voxel into canvas
        # Calculate necessary things for canvas creation
        canvas_z = int(
            (self.point_cloud_range[5] - self.point_cloud_range[2]) / self.vz)
        canvas_y = int(
            (self.point_cloud_range[4] - self.point_cloud_range[1]) / self.vy)
        canvas_x = int(
            (self.point_cloud_range[3] - self.point_cloud_range[0]) / self.vx)
        # canvas_channel = voxel_mean.size(1)
        batch_size = pts_coors[-1, 0] + 1
        canvas_len = canvas_z * canvas_y * canvas_x * batch_size
        # Create the canvas for this sample
        canvas = voxel_mean.new_zeros(canvas_len, dtype=torch.long)
        # Only include non-empty pillars
        indices = (
            voxel_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            voxel_coors[:, 1] * canvas_y * canvas_x +
            voxel_coors[:, 2] * canvas_x + voxel_coors[:, 3])
        # Scatter the blob back to the canvas
        canvas[indices.long()] = torch.arange(
            start=0, end=voxel_mean.size(0), device=voxel_mean.device)

        # Step 2: get voxel mean for each point
        voxel_index = (
            pts_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            pts_coors[:, 1] * canvas_y * canvas_x +
            pts_coors[:, 2] * canvas_x + pts_coors[:, 3])
        voxel_inds = canvas[voxel_index.long()]
        center_per_point = voxel_mean[voxel_inds, ...]
        return center_per_point

221
    @force_fp32(out_fp16=True)
zhangwenwei's avatar
zhangwenwei committed
222
223
224
225
226
    def forward(self,
                features,
                coors,
                points=None,
                img_feats=None,
zhangwenwei's avatar
zhangwenwei committed
227
                img_metas=None):
zhangwenwei's avatar
zhangwenwei committed
228
        """Forward functions.
zhangwenwei's avatar
zhangwenwei committed
229
230
231
232
233
234

        Args:
            features (torch.Tensor): Features of voxels, shape is NxC.
            coors (torch.Tensor): Coordinates of voxels, shape is  Nx(1+NDim).
            points (list[torch.Tensor], optional): Raw points used to guide the
                multi-modality fusion. Defaults to None.
235
            img_feats (list[torch.Tensor], optional): Image features used for
zhangwenwei's avatar
zhangwenwei committed
236
                multi-modality fusion. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
237
            img_metas (dict, optional): [description]. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
238
239
240
241
242

        Returns:
            tuple: If `return_point_feats` is False, returns voxel features and
                its coordinates. If `return_point_feats` is True, returns
                feature of each points inside voxels.
zhangwenwei's avatar
zhangwenwei committed
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        """
        features_ls = [features]
        # Find distance of x, y, and z from cluster center
        if self._with_cluster_center:
            voxel_mean, mean_coors = self.cluster_scatter(features, coors)
            points_mean = self.map_voxel_center_to_point(
                coors, voxel_mean, mean_coors)
            # TODO: maybe also do cluster for reflectivity
            f_cluster = features[:, :3] - points_mean[:, :3]
            features_ls.append(f_cluster)

        # Find distance of x, y, and z from pillar center
        if self._with_voxel_center:
            f_center = features.new_zeros(size=(features.size(0), 3))
            f_center[:, 0] = features[:, 0] - (
                coors[:, 3].type_as(features) * self.vx + self.x_offset)
            f_center[:, 1] = features[:, 1] - (
                coors[:, 2].type_as(features) * self.vy + self.y_offset)
            f_center[:, 2] = features[:, 2] - (
                coors[:, 1].type_as(features) * self.vz + self.z_offset)
            features_ls.append(f_center)

        if self._with_distance:
            points_dist = torch.norm(features[:, :3], 2, 1, keepdim=True)
            features_ls.append(points_dist)

        # Combine together feature decorations
        features = torch.cat(features_ls, dim=-1)
        for i, vfe in enumerate(self.vfe_layers):
            point_feats = vfe(features)
            if (i == len(self.vfe_layers) - 1 and self.fusion_layer is not None
                    and img_feats is not None):
                point_feats = self.fusion_layer(img_feats, points, point_feats,
zhangwenwei's avatar
zhangwenwei committed
276
                                                img_metas)
zhangwenwei's avatar
zhangwenwei committed
277
278
279
280
281
282
283
284
285
286
287
288
            voxel_feats, voxel_coors = self.vfe_scatter(point_feats, coors)
            if i != len(self.vfe_layers) - 1:
                # need to concat voxel feats if it is not the last vfe
                feat_per_point = self.map_voxel_center_to_point(
                    coors, voxel_feats, voxel_coors)
                features = torch.cat([point_feats, feat_per_point], dim=1)

        if self.return_point_feats:
            return point_feats
        return voxel_feats, voxel_coors


289
@MODELS.register_module()
zhangwenwei's avatar
zhangwenwei committed
290
class HardVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
291
    """Voxel feature encoder used in DV-SECOND.
zhangwenwei's avatar
zhangwenwei committed
292
293
294
295
296

    It encodes features of voxels and their points. It could also fuse
    image feature into voxel features in a point-wise manner.

    Args:
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
        in_channels (int, optional): Input channels of VFE. Defaults to 4.
        feat_channels (list(int), optional): Channels of features in VFE.
        with_distance (bool, optional): Whether to use the L2 distance
            of points to the origin point. Defaults to False.
        with_cluster_center (bool, optional): Whether to use the distance
            to cluster center of points inside a voxel. Defaults to False.
        with_voxel_center (bool, optional): Whether to use the distance to
            center of voxel for each points inside a voxel. Defaults to False.
        voxel_size (tuple[float], optional): Size of a single voxel.
            Defaults to (0.2, 0.2, 4).
        point_cloud_range (tuple[float], optional): The range of points
            or voxels. Defaults to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict, optional): Config dict of normalization layers.
        mode (str, optional): The mode when pooling features of points inside a
            voxel. Available options include 'max' and 'avg'.
            Defaults to 'max'.
        fusion_layer (dict, optional): The config dict of fusion layer
            used in multi-modal detectors. Defaults to None.
        return_point_feats (bool, optional): Whether to return the
            features of each points. Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
317
    """
zhangwenwei's avatar
zhangwenwei committed
318
319

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
320
321
                 in_channels=4,
                 feat_channels=[],
zhangwenwei's avatar
zhangwenwei committed
322
323
324
325
326
327
328
329
330
331
                 with_distance=False,
                 with_cluster_center=False,
                 with_voxel_center=False,
                 voxel_size=(0.2, 0.2, 4),
                 point_cloud_range=(0, -40, -3, 70.4, 40, 1),
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 mode='max',
                 fusion_layer=None,
                 return_point_feats=False):
        super(HardVFE, self).__init__()
zhangwenwei's avatar
zhangwenwei committed
332
        assert len(feat_channels) > 0
zhangwenwei's avatar
zhangwenwei committed
333
        if with_cluster_center:
zhangwenwei's avatar
zhangwenwei committed
334
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
335
        if with_voxel_center:
zhangwenwei's avatar
zhangwenwei committed
336
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
337
        if with_distance:
338
            in_channels += 1
zhangwenwei's avatar
zhangwenwei committed
339
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
340
341
342
343
        self._with_distance = with_distance
        self._with_cluster_center = with_cluster_center
        self._with_voxel_center = with_voxel_center
        self.return_point_feats = return_point_feats
344
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
345
346
347
348
349
350
351
352
353
354
355

        # Need pillar (voxel) size and x/y offset to calculate pillar offset
        self.vx = voxel_size[0]
        self.vy = voxel_size[1]
        self.vz = voxel_size[2]
        self.x_offset = self.vx / 2 + point_cloud_range[0]
        self.y_offset = self.vy / 2 + point_cloud_range[1]
        self.z_offset = self.vz / 2 + point_cloud_range[2]
        self.point_cloud_range = point_cloud_range
        self.scatter = DynamicScatter(voxel_size, point_cloud_range, True)

zhangwenwei's avatar
zhangwenwei committed
356
        feat_channels = [self.in_channels] + list(feat_channels)
zhangwenwei's avatar
zhangwenwei committed
357
        vfe_layers = []
zhangwenwei's avatar
zhangwenwei committed
358
359
360
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
zhangwenwei's avatar
zhangwenwei committed
361
362
363
364
            if i > 0:
                in_filters *= 2
            # TODO: pass norm_cfg to VFE
            # norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)
zhangwenwei's avatar
zhangwenwei committed
365
            if i == (len(feat_channels) - 2):
zhangwenwei's avatar
zhangwenwei committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
                cat_max = False
                max_out = True
                if fusion_layer:
                    max_out = False
            else:
                max_out = True
                cat_max = True
            vfe_layers.append(
                VFELayer(
                    in_filters,
                    out_filters,
                    norm_cfg=norm_cfg,
                    max_out=max_out,
                    cat_max=cat_max))
            self.vfe_layers = nn.ModuleList(vfe_layers)
        self.num_vfe = len(vfe_layers)

        self.fusion_layer = None
        if fusion_layer is not None:
            self.fusion_layer = builder.build_fusion_layer(fusion_layer)

387
    @force_fp32(out_fp16=True)
zhangwenwei's avatar
zhangwenwei committed
388
389
390
391
392
    def forward(self,
                features,
                num_points,
                coors,
                img_feats=None,
zhangwenwei's avatar
zhangwenwei committed
393
                img_metas=None):
zhangwenwei's avatar
zhangwenwei committed
394
        """Forward functions.
zhangwenwei's avatar
zhangwenwei committed
395
396
397
398
399

        Args:
            features (torch.Tensor): Features of voxels, shape is MxNxC.
            num_points (torch.Tensor): Number of points in each voxel.
            coors (torch.Tensor): Coordinates of voxels, shape is Mx(1+NDim).
400
            img_feats (list[torch.Tensor], optional): Image features used for
zhangwenwei's avatar
zhangwenwei committed
401
                multi-modality fusion. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
402
            img_metas (dict, optional): [description]. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
403
404
405
406
407

        Returns:
            tuple: If `return_point_feats` is False, returns voxel features and
                its coordinates. If `return_point_feats` is True, returns
                feature of each points inside voxels.
zhangwenwei's avatar
zhangwenwei committed
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
        """
        features_ls = [features]
        # Find distance of x, y, and z from cluster center
        if self._with_cluster_center:
            points_mean = (
                features[:, :, :3].sum(dim=1, keepdim=True) /
                num_points.type_as(features).view(-1, 1, 1))
            # TODO: maybe also do cluster for reflectivity
            f_cluster = features[:, :, :3] - points_mean
            features_ls.append(f_cluster)

        # Find distance of x, y, and z from pillar center
        if self._with_voxel_center:
            f_center = features.new_zeros(
                size=(features.size(0), features.size(1), 3))
            f_center[:, :, 0] = features[:, :, 0] - (
                coors[:, 3].type_as(features).unsqueeze(1) * self.vx +
                self.x_offset)
            f_center[:, :, 1] = features[:, :, 1] - (
                coors[:, 2].type_as(features).unsqueeze(1) * self.vy +
                self.y_offset)
            f_center[:, :, 2] = features[:, :, 2] - (
                coors[:, 1].type_as(features).unsqueeze(1) * self.vz +
                self.z_offset)
            features_ls.append(f_center)

        if self._with_distance:
            points_dist = torch.norm(features[:, :, :3], 2, 2, keepdim=True)
            features_ls.append(points_dist)

        # Combine together feature decorations
        voxel_feats = torch.cat(features_ls, dim=-1)
        # The feature decorations were calculated without regard to whether
        # pillar was empty.
        # Need to ensure that empty voxels remain set to zeros.
        voxel_count = voxel_feats.shape[1]
        mask = get_paddings_indicator(num_points, voxel_count, axis=0)
        voxel_feats *= mask.unsqueeze(-1).type_as(voxel_feats)

        for i, vfe in enumerate(self.vfe_layers):
            voxel_feats = vfe(voxel_feats)
zhangwenwei's avatar
zhangwenwei committed
449

zhangwenwei's avatar
zhangwenwei committed
450
451
        if (self.fusion_layer is not None and img_feats is not None):
            voxel_feats = self.fusion_with_mask(features, mask, voxel_feats,
zhangwenwei's avatar
zhangwenwei committed
452
                                                coors, img_feats, img_metas)
zhangwenwei's avatar
zhangwenwei committed
453

zhangwenwei's avatar
zhangwenwei committed
454
455
456
        return voxel_feats

    def fusion_with_mask(self, features, mask, voxel_feats, coors, img_feats,
zhangwenwei's avatar
zhangwenwei committed
457
                         img_metas):
zhangwenwei's avatar
zhangwenwei committed
458
459
460
461
462
463
464
465
466
        """Fuse image and point features with mask.

        Args:
            features (torch.Tensor): Features of voxel, usually it is the
                values of points in voxels.
            mask (torch.Tensor): Mask indicates valid features in each voxel.
            voxel_feats (torch.Tensor): Features of voxels.
            coors (torch.Tensor): Coordinates of each single voxel.
            img_feats (list[torch.Tensor]): Multi-scale feature maps of image.
zhangwenwei's avatar
zhangwenwei committed
467
            img_metas (list(dict)): Meta information of image and points.
zhangwenwei's avatar
zhangwenwei committed
468
469
470
471

        Returns:
            torch.Tensor: Fused features of each voxel.
        """
zhangwenwei's avatar
zhangwenwei committed
472
473
474
475
476
477
478
479
480
        # the features is consist of a batch of points
        batch_size = coors[-1, 0] + 1
        points = []
        for i in range(batch_size):
            single_mask = (coors[:, 0] == i)
            points.append(features[single_mask][mask[single_mask]])

        point_feats = voxel_feats[mask]
        point_feats = self.fusion_layer(img_feats, points, point_feats,
zhangwenwei's avatar
zhangwenwei committed
481
                                        img_metas)
zhangwenwei's avatar
zhangwenwei committed
482

zhangwenwei's avatar
zhangwenwei committed
483
484
485
486
487
        voxel_canvas = voxel_feats.new_zeros(
            size=(voxel_feats.size(0), voxel_feats.size(1),
                  point_feats.size(-1)))
        voxel_canvas[mask] = point_feats
        out = torch.max(voxel_canvas, dim=1)[0]
zhangwenwei's avatar
zhangwenwei committed
488

zhangwenwei's avatar
zhangwenwei committed
489
        return out