voxel_encoder.py 20.4 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
jshilong's avatar
jshilong committed
6
from torch import Tensor, nn
zhangwenwei's avatar
zhangwenwei committed
7

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

jshilong's avatar
jshilong committed
23
    def __init__(self, num_features: int = 4) -> None:
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)
jshilong's avatar
jshilong committed
29
30
    def forward(self, features: Tensor, num_points: Tensor, coors: Tensor,
                *args, **kwargs) -> Tensor:
zhangwenwei's avatar
zhangwenwei committed
31
        """Forward function.
zhangwenwei's avatar
zhangwenwei committed
32
33

        Args:
wangtai's avatar
wangtai committed
34
            features (torch.Tensor): Point features in shape
zhangwenwei's avatar
zhangwenwei committed
35
36
37
38
39
40
41
42
43
                (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))
        """
44
        points_mean = features[:, :, :self.num_features].sum(
zhangwenwei's avatar
zhangwenwei committed
45
46
47
48
            dim=1, keepdim=False) / num_points.type_as(features).view(-1, 1)
        return points_mean.contiguous()


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

    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
60
61
62
63

    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
64
        super(DynamicSimpleVFE, self).__init__()
zhangwenwei's avatar
zhangwenwei committed
65
        self.scatter = DynamicScatter(voxel_size, point_cloud_range, True)
66
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
67
68

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

        Args:
wangtai's avatar
wangtai committed
74
            features (torch.Tensor): Point features in shape
zhangwenwei's avatar
zhangwenwei committed
75
76
77
78
79
80
81
                (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
82
83
84
85
86
87
        # 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


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

    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:
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
        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
118
    """
zhangwenwei's avatar
zhangwenwei committed
119
120

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
121
122
                 in_channels=4,
                 feat_channels=[],
zhangwenwei's avatar
zhangwenwei committed
123
124
125
126
127
128
129
130
131
132
                 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
133
134
        assert mode in ['avg', 'max']
        assert len(feat_channels) > 0
zhangwenwei's avatar
zhangwenwei committed
135
        if with_cluster_center:
zhangwenwei's avatar
zhangwenwei committed
136
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
137
        if with_voxel_center:
zhangwenwei's avatar
zhangwenwei committed
138
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
139
        if with_distance:
140
            in_channels += 1
zhangwenwei's avatar
zhangwenwei committed
141
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
142
143
144
145
        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
146
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
147
148
149
150
151
152
153
154
155
156
157

        # 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
158
        feat_channels = [self.in_channels] + list(feat_channels)
zhangwenwei's avatar
zhangwenwei committed
159
        vfe_layers = []
zhangwenwei's avatar
zhangwenwei committed
160
161
162
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
zhangwenwei's avatar
zhangwenwei committed
163
164
165
166
167
168
169
            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)))
170
        self.vfe_layers = nn.ModuleList(vfe_layers)
zhangwenwei's avatar
zhangwenwei committed
171
172
173
174
175
176
177
178
179
180
        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
181
182
183
184
185
186
187
188
189
190
        """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
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
221
222
223
224
225
226
        # 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

    def forward(self,
                features,
                coors,
                points=None,
                img_feats=None,
jshilong's avatar
jshilong committed
227
228
229
                img_metas=None,
                *args,
                **kwargs):
zhangwenwei's avatar
zhangwenwei committed
230
        """Forward functions.
zhangwenwei's avatar
zhangwenwei committed
231
232
233
234
235
236

        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.
237
            img_feats (list[torch.Tensor], optional): Image features used for
zhangwenwei's avatar
zhangwenwei committed
238
                multi-modality fusion. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
239
            img_metas (dict, optional): [description]. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
240
241
242
243
244

        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
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
276
277
        """
        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
278
                                                img_metas)
zhangwenwei's avatar
zhangwenwei committed
279
280
281
282
283
284
285
286
287
288
289
290
            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


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

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

    Args:
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
        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
319
    """
zhangwenwei's avatar
zhangwenwei committed
320
321

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
322
323
                 in_channels=4,
                 feat_channels=[],
zhangwenwei's avatar
zhangwenwei committed
324
325
326
327
328
329
330
331
332
333
                 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
334
        assert len(feat_channels) > 0
zhangwenwei's avatar
zhangwenwei committed
335
        if with_cluster_center:
zhangwenwei's avatar
zhangwenwei committed
336
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
337
        if with_voxel_center:
zhangwenwei's avatar
zhangwenwei committed
338
            in_channels += 3
zhangwenwei's avatar
zhangwenwei committed
339
        if with_distance:
340
            in_channels += 1
zhangwenwei's avatar
zhangwenwei committed
341
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
342
343
344
345
        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
346
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
347
348
349
350
351
352
353
354
355
356
357

        # 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
358
        feat_channels = [self.in_channels] + list(feat_channels)
zhangwenwei's avatar
zhangwenwei committed
359
        vfe_layers = []
zhangwenwei's avatar
zhangwenwei committed
360
361
362
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
zhangwenwei's avatar
zhangwenwei committed
363
364
365
366
            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
367
            if i == (len(feat_channels) - 2):
zhangwenwei's avatar
zhangwenwei committed
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
                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)

389
    @force_fp32(out_fp16=True)
zhangwenwei's avatar
zhangwenwei committed
390
391
392
393
394
    def forward(self,
                features,
                num_points,
                coors,
                img_feats=None,
jshilong's avatar
jshilong committed
395
396
397
                img_metas=None,
                *args,
                **kwargs):
zhangwenwei's avatar
zhangwenwei committed
398
        """Forward functions.
zhangwenwei's avatar
zhangwenwei committed
399
400
401
402
403

        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).
404
            img_feats (list[torch.Tensor], optional): Image features used for
zhangwenwei's avatar
zhangwenwei committed
405
                multi-modality fusion. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
406
            img_metas (dict, optional): [description]. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
407
408
409
410
411

        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
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
449
450
451
452
        """
        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
453

zhangwenwei's avatar
zhangwenwei committed
454
455
        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
456
                                                coors, img_feats, img_metas)
zhangwenwei's avatar
zhangwenwei committed
457

zhangwenwei's avatar
zhangwenwei committed
458
459
460
        return voxel_feats

    def fusion_with_mask(self, features, mask, voxel_feats, coors, img_feats,
zhangwenwei's avatar
zhangwenwei committed
461
                         img_metas):
zhangwenwei's avatar
zhangwenwei committed
462
463
464
465
466
467
468
469
470
        """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
471
            img_metas (list(dict)): Meta information of image and points.
zhangwenwei's avatar
zhangwenwei committed
472
473
474
475

        Returns:
            torch.Tensor: Fused features of each voxel.
        """
zhangwenwei's avatar
zhangwenwei committed
476
477
478
479
480
481
482
483
484
        # 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
485
                                        img_metas)
zhangwenwei's avatar
zhangwenwei committed
486

zhangwenwei's avatar
zhangwenwei committed
487
488
489
490
491
        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
492

zhangwenwei's avatar
zhangwenwei committed
493
        return out