voxel_encoder.py 20 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
import torch
2
from mmcv.cnn import build_norm_layer
3
from mmcv.runner import force_fp32
zhangwenwei's avatar
zhangwenwei committed
4
5
6
7
from torch import nn

from mmdet3d.ops import DynamicScatter
from .. import builder
8
from ..builder import VOXEL_ENCODERS
zhangwenwei's avatar
zhangwenwei committed
9
from .utils import VFELayer, get_paddings_indicator
zhangwenwei's avatar
zhangwenwei committed
10
11


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

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

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

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

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

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


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

    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
58
59
60
61

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

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

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


86
@VOXEL_ENCODERS.register_module()
zhangwenwei's avatar
zhangwenwei committed
87
class DynamicVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
88
    """Dynamic Voxel feature encoder used in DV-SECOND.
zhangwenwei's avatar
zhangwenwei committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

    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:
        in_channels (int): Input channels of VFE. Defaults to 4.
        feat_channels (list(int)): Channels of features in VFE.
        with_distance (bool): Whether to use the L2 distance of points to the
            origin point. Default False.
        with_cluster_center (bool): Whether to use the distance to cluster
            center of points inside a voxel. Default to False.
        with_voxel_center (bool): Whether to use the distance to center of
            voxel for each points inside a voxel. Default to False.
        voxel_size (tuple[float]): Size of a single voxel. Default to
            (0.2, 0.2, 4).
        point_cloud_range (tuple[float]): The range of points or voxels.
            Default to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict): Config dict of normalization layers.
        mode (str): The mode when pooling features of points inside a voxel.
            Available options include 'max' and 'avg'. Default to 'max'.
        fusion_layer (dict | None): The config dict of fusion layer used in
            multi-modal detectors. Default to None.
        return_point_feats (bool): Whether to return the features of each
            points. Default to False.
    """
zhangwenwei's avatar
zhangwenwei committed
115
116

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

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

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

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

        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
240
241
242
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
        """
        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
273
                                                img_metas)
zhangwenwei's avatar
zhangwenwei committed
274
275
276
277
278
279
280
281
282
283
284
285
            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


286
@VOXEL_ENCODERS.register_module()
zhangwenwei's avatar
zhangwenwei committed
287
class HardVFE(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
288
    """Voxel feature encoder used in DV-SECOND.
zhangwenwei's avatar
zhangwenwei committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313

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

    Args:
        in_channels (int): Input channels of VFE. Defaults to 4.
        feat_channels (list(int)): Channels of features in VFE.
        with_distance (bool): Whether to use the L2 distance of points to the
            origin point. Default False.
        with_cluster_center (bool): Whether to use the distance to cluster
            center of points inside a voxel. Default to False.
        with_voxel_center (bool): Whether to use the distance to center of
            voxel for each points inside a voxel. Default to False.
        voxel_size (tuple[float]): Size of a single voxel. Default to
            (0.2, 0.2, 4).
        point_cloud_range (tuple[float]): The range of points or voxels.
            Default to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict): Config dict of normalization layers.
        mode (str): The mode when pooling features of points inside a voxel.
            Available options include 'max' and 'avg'. Default to 'max'.
        fusion_layer (dict | None): The config dict of fusion layer used in
            multi-modal detectors. Default to None.
        return_point_feats (bool): Whether to return the features of each
            points. Default to False.
    """
zhangwenwei's avatar
zhangwenwei committed
314
315

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

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

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

        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).
            img_feats (list[torch.Tensor], optional): Image fetures used for
                multi-modality fusion. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
398
            img_metas (dict, optional): [description]. Defaults to None.
zhangwenwei's avatar
zhangwenwei committed
399
400
401
402
403

        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
404
405
406
407
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
        """
        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
445

zhangwenwei's avatar
zhangwenwei committed
446
447
        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
448
                                                coors, img_feats, img_metas)
zhangwenwei's avatar
zhangwenwei committed
449

zhangwenwei's avatar
zhangwenwei committed
450
451
452
        return voxel_feats

    def fusion_with_mask(self, features, mask, voxel_feats, coors, img_feats,
zhangwenwei's avatar
zhangwenwei committed
453
                         img_metas):
zhangwenwei's avatar
zhangwenwei committed
454
455
456
457
458
459
460
461
462
        """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
463
            img_metas (list(dict)): Meta information of image and points.
zhangwenwei's avatar
zhangwenwei committed
464
465
466
467

        Returns:
            torch.Tensor: Fused features of each voxel.
        """
zhangwenwei's avatar
zhangwenwei committed
468
469
470
471
472
473
474
475
476
        # 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
477
                                        img_metas)
zhangwenwei's avatar
zhangwenwei committed
478

zhangwenwei's avatar
zhangwenwei committed
479
480
481
482
483
        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
484

zhangwenwei's avatar
zhangwenwei committed
485
        return out