point_fusion.py 12 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 ConvModule
4
from mmengine.model import BaseModule
zhangwenwei's avatar
zhangwenwei committed
5
6
from torch import nn as nn
from torch.nn import functional as F
zhangwenwei's avatar
zhangwenwei committed
7

8
from mmdet3d.registry import MODELS
zhangshilong's avatar
zhangshilong committed
9
10
from mmdet3d.structures.bbox_3d import (get_proj_mat_by_coord_type,
                                        points_cam2img)
11
from . import apply_3d_transformation
zhangwenwei's avatar
zhangwenwei committed
12
13


14
15
16
17
18
19
20
21
22
23
24
25
26
def point_sample(img_meta,
                 img_features,
                 points,
                 proj_mat,
                 coord_type,
                 img_scale_factor,
                 img_crop_offset,
                 img_flip,
                 img_pad_shape,
                 img_shape,
                 aligned=True,
                 padding_mode='zeros',
                 align_corners=True):
zhangwenwei's avatar
zhangwenwei committed
27
    """Obtain image features using points.
zhangwenwei's avatar
zhangwenwei committed
28

zhangwenwei's avatar
zhangwenwei committed
29
    Args:
30
        img_meta (dict): Meta info.
wangtai's avatar
wangtai committed
31
32
        img_features (torch.Tensor): 1 x C x H x W image features.
        points (torch.Tensor): Nx3 point cloud in LiDAR coordinates.
33
34
        proj_mat (torch.Tensor): 4x4 transformation matrix.
        coord_type (str): 'DEPTH' or 'CAMERA' or 'LIDAR'.
35
        img_scale_factor (torch.Tensor): Scale factor with shape of
wangtai's avatar
wangtai committed
36
            (w_scale, h_scale).
37
        img_crop_offset (torch.Tensor): Crop offset used to crop
wangtai's avatar
wangtai committed
38
            image during data augmentation with shape of (w_offset, h_offset).
zhangwenwei's avatar
zhangwenwei committed
39
40
        img_flip (bool): Whether the image is flipped.
        img_pad_shape (tuple[int]): int tuple indicates the h & w after
wangtai's avatar
wangtai committed
41
            padding, this is necessary to obtain features in feature map.
zhangwenwei's avatar
zhangwenwei committed
42
        img_shape (tuple[int]): int tuple indicates the h & w before padding
wangtai's avatar
wangtai committed
43
            after scaling, this is necessary for flipping coordinates.
zhangwenwei's avatar
zhangwenwei committed
44
45
46
47
48
49
50
51
        aligned (bool, optional): Whether use bilinear interpolation when
            sampling image features for each point. Defaults to True.
        padding_mode (str, optional): Padding mode when padding values for
            features of out-of-image points. Defaults to 'zeros'.
        align_corners (bool, optional): Whether to align corners when
            sampling image features for each point. Defaults to True.

    Returns:
wangtai's avatar
wangtai committed
52
        torch.Tensor: NxC image features sampled by point coordinates.
zhangwenwei's avatar
zhangwenwei committed
53
    """
54
55

    # apply transformation based on info in img_meta
56
57
    points = apply_3d_transformation(
        points, coord_type, img_meta, reverse=True)
zhangwenwei's avatar
zhangwenwei committed
58

59
    # project points to image coordinate
60
    pts_2d = points_cam2img(points, proj_mat)
zhangwenwei's avatar
zhangwenwei committed
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

    # img transformation: scale -> crop -> flip
    # the image is resized by img_scale_factor
    img_coors = pts_2d[:, 0:2] * img_scale_factor  # Nx2
    img_coors -= img_crop_offset

    # grid sample, the valid grid range should be in [-1,1]
    coor_x, coor_y = torch.split(img_coors, 1, dim=1)  # each is Nx1

    if img_flip:
        # by default we take it as horizontal flip
        # use img_shape before padding for flip
        orig_h, orig_w = img_shape
        coor_x = orig_w - coor_x

    h, w = img_pad_shape
    coor_y = coor_y / h * 2 - 1
    coor_x = coor_x / w * 2 - 1
    grid = torch.cat([coor_x, coor_y],
                     dim=1).unsqueeze(0).unsqueeze(0)  # Nx2 -> 1x1xNx2

    # align_corner=True provides higher performance
    mode = 'bilinear' if aligned else 'nearest'
    point_features = F.grid_sample(
        img_features,
        grid,
        mode=mode,
        padding_mode=padding_mode,
        align_corners=align_corners)  # 1xCx1xN feats

    return point_features.squeeze().t()


94
@MODELS.register_module()
95
class PointFusion(BaseModule):
zhangwenwei's avatar
zhangwenwei committed
96
    """Fuse image features from multi-scale features.
zhangwenwei's avatar
zhangwenwei committed
97
98
99
100
101
102
103
104

    Args:
        img_channels (list[int] | int): Channels of image features.
            It could be a list if the input is multi-scale image features.
        pts_channels (int): Channels of point features
        mid_channels (int): Channels of middle layers
        out_channels (int): Channels of output fused features
        img_levels (int, optional): Number of image levels. Defaults to 3.
105
106
        coord_type (str): 'DEPTH' or 'CAMERA' or 'LIDAR'.
            Defaults to 'LIDAR'.
zhangwenwei's avatar
zhangwenwei committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        conv_cfg (dict, optional): Dict config of conv layers of middle
            layers. Defaults to None.
        norm_cfg (dict, optional): Dict config of norm layers of middle
            layers. Defaults to None.
        act_cfg (dict, optional): Dict config of activatation layers.
            Defaults to None.
        activate_out (bool, optional): Whether to apply relu activation
            to output features. Defaults to True.
        fuse_out (bool, optional): Whether apply conv layer to the fused
            features. Defaults to False.
        dropout_ratio (int, float, optional): Dropout ratio of image
            features to prevent overfitting. Defaults to 0.
        aligned (bool, optional): Whether apply aligned feature fusion.
            Defaults to True.
        align_corners (bool, optional): Whether to align corner when
            sampling features according to points. Defaults to True.
        padding_mode (str, optional): Mode used to pad the features of
            points that do not have corresponding image features.
            Defaults to 'zeros'.
        lateral_conv (bool, optional): Whether to apply lateral convs
            to image features. Defaults to True.
zhangwenwei's avatar
zhangwenwei committed
128
129
130
131
132
133
134
135
    """

    def __init__(self,
                 img_channels,
                 pts_channels,
                 mid_channels,
                 out_channels,
                 img_levels=3,
136
                 coord_type='LIDAR',
zhangwenwei's avatar
zhangwenwei committed
137
138
                 conv_cfg=None,
                 norm_cfg=None,
zhangwenwei's avatar
zhangwenwei committed
139
                 act_cfg=None,
140
                 init_cfg=None,
zhangwenwei's avatar
zhangwenwei committed
141
142
143
144
145
146
147
                 activate_out=True,
                 fuse_out=False,
                 dropout_ratio=0,
                 aligned=True,
                 align_corners=True,
                 padding_mode='zeros',
                 lateral_conv=True):
148
        super(PointFusion, self).__init__(init_cfg=init_cfg)
zhangwenwei's avatar
zhangwenwei committed
149
150
151
152
153
154
155
156
157
        if isinstance(img_levels, int):
            img_levels = [img_levels]
        if isinstance(img_channels, int):
            img_channels = [img_channels] * len(img_levels)
        assert isinstance(img_levels, list)
        assert isinstance(img_channels, list)
        assert len(img_channels) == len(img_levels)

        self.img_levels = img_levels
158
        self.coord_type = coord_type
zhangwenwei's avatar
zhangwenwei committed
159
        self.act_cfg = act_cfg
zhangwenwei's avatar
zhangwenwei committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
        self.activate_out = activate_out
        self.fuse_out = fuse_out
        self.dropout_ratio = dropout_ratio
        self.img_channels = img_channels
        self.aligned = aligned
        self.align_corners = align_corners
        self.padding_mode = padding_mode

        self.lateral_convs = None
        if lateral_conv:
            self.lateral_convs = nn.ModuleList()
            for i in range(len(img_channels)):
                l_conv = ConvModule(
                    img_channels[i],
                    mid_channels,
                    3,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
zhangwenwei's avatar
zhangwenwei committed
179
                    act_cfg=self.act_cfg,
zhangwenwei's avatar
zhangwenwei committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
                    inplace=False)
                self.lateral_convs.append(l_conv)
            self.img_transform = nn.Sequential(
                nn.Linear(mid_channels * len(img_channels), out_channels),
                nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
            )
        else:
            self.img_transform = nn.Sequential(
                nn.Linear(sum(img_channels), out_channels),
                nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
            )
        self.pts_transform = nn.Sequential(
            nn.Linear(pts_channels, out_channels),
            nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
        )

        if self.fuse_out:
            self.fuse_conv = nn.Sequential(
                nn.Linear(mid_channels, out_channels),
                # For pts the BN is initialized differently by default
                # TODO: check whether this is necessary
                nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
                nn.ReLU(inplace=False))

204
205
206
207
208
        if init_cfg is None:
            self.init_cfg = [
                dict(type='Xavier', layer='Conv2d', distribution='uniform'),
                dict(type='Xavier', layer='Linear', distribution='uniform')
            ]
zhangwenwei's avatar
zhangwenwei committed
209

zhangwenwei's avatar
zhangwenwei committed
210
    def forward(self, img_feats, pts, pts_feats, img_metas):
zhangwenwei's avatar
zhangwenwei committed
211
        """Forward function.
zhangwenwei's avatar
zhangwenwei committed
212
213

        Args:
wangtai's avatar
wangtai committed
214
215
216
217
218
            img_feats (list[torch.Tensor]): Image features.
            pts: [list[torch.Tensor]]: A batch of points with shape N x 3.
            pts_feats (torch.Tensor): A tensor consist of point features of the
                total batch.
            img_metas (list[dict]): Meta information of images.
zhangwenwei's avatar
zhangwenwei committed
219

zhangwenwei's avatar
zhangwenwei committed
220
        Returns:
wangtai's avatar
wangtai committed
221
            torch.Tensor: Fused features of each point.
zhangwenwei's avatar
zhangwenwei committed
222
        """
zhangwenwei's avatar
zhangwenwei committed
223
        img_pts = self.obtain_mlvl_feats(img_feats, pts, img_metas)
zhangwenwei's avatar
zhangwenwei committed
224
225
226
227
228
229
230
231
232
233
234
235
236
        img_pre_fuse = self.img_transform(img_pts)
        if self.training and self.dropout_ratio > 0:
            img_pre_fuse = F.dropout(img_pre_fuse, self.dropout_ratio)
        pts_pre_fuse = self.pts_transform(pts_feats)

        fuse_out = img_pre_fuse + pts_pre_fuse
        if self.activate_out:
            fuse_out = F.relu(fuse_out)
        if self.fuse_out:
            fuse_out = self.fuse_conv(fuse_out)

        return fuse_out

zhangwenwei's avatar
zhangwenwei committed
237
    def obtain_mlvl_feats(self, img_feats, pts, img_metas):
238
239
240
241
242
243
244
245
246
247
248
        """Obtain multi-level features for each point.

        Args:
            img_feats (list(torch.Tensor)): Multi-scale image features produced
                by image backbone in shape (N, C, H, W).
            pts (list[torch.Tensor]): Points of each sample.
            img_metas (list[dict]): Meta information for each sample.

        Returns:
            torch.Tensor: Corresponding image features of each point.
        """
zhangwenwei's avatar
zhangwenwei committed
249
250
251
252
253
254
255
256
257
        if self.lateral_convs is not None:
            img_ins = [
                lateral_conv(img_feats[i])
                for i, lateral_conv in zip(self.img_levels, self.lateral_convs)
            ]
        else:
            img_ins = img_feats
        img_feats_per_point = []
        # Sample multi-level features
zhangwenwei's avatar
zhangwenwei committed
258
        for i in range(len(img_metas)):
zhangwenwei's avatar
zhangwenwei committed
259
260
261
262
            mlvl_img_feats = []
            for level in range(len(self.img_levels)):
                mlvl_img_feats.append(
                    self.sample_single(img_ins[level][i:i + 1], pts[i][:, :3],
zhangwenwei's avatar
zhangwenwei committed
263
                                       img_metas[i]))
zhangwenwei's avatar
zhangwenwei committed
264
265
266
267
268
269
270
            mlvl_img_feats = torch.cat(mlvl_img_feats, dim=-1)
            img_feats_per_point.append(mlvl_img_feats)

        img_pts = torch.cat(img_feats_per_point, dim=0)
        return img_pts

    def sample_single(self, img_feats, pts, img_meta):
271
272
273
274
        """Sample features from single level image feature map.

        Args:
            img_feats (torch.Tensor): Image feature map in shape
275
                (1, C, H, W).
276
277
278
279
280
281
            pts (torch.Tensor): Points of a single sample.
            img_meta (dict): Meta information of the single sample.

        Returns:
            torch.Tensor: Single level image features of each point.
        """
282
        # TODO: image transformation also extracted
zhangwenwei's avatar
zhangwenwei committed
283
        img_scale_factor = (
zhangwenwei's avatar
zhangwenwei committed
284
            pts.new_tensor(img_meta['scale_factor'][:2])
zhangwenwei's avatar
zhangwenwei committed
285
286
287
288
289
            if 'scale_factor' in img_meta.keys() else 1)
        img_flip = img_meta['flip'] if 'flip' in img_meta.keys() else False
        img_crop_offset = (
            pts.new_tensor(img_meta['img_crop_offset'])
            if 'img_crop_offset' in img_meta.keys() else 0)
290
        proj_mat = get_proj_mat_by_coord_type(img_meta, self.coord_type)
zhangwenwei's avatar
zhangwenwei committed
291
        img_pts = point_sample(
292
293
294
295
296
297
298
            img_meta=img_meta,
            img_features=img_feats,
            points=pts,
            proj_mat=pts.new_tensor(proj_mat),
            coord_type=self.coord_type,
            img_scale_factor=img_scale_factor,
            img_crop_offset=img_crop_offset,
zhangwenwei's avatar
zhangwenwei committed
299
            img_flip=img_flip,
300
            img_pad_shape=img_meta['input_shape'][:2],
zhangwenwei's avatar
zhangwenwei committed
301
302
303
304
305
306
            img_shape=img_meta['img_shape'][:2],
            aligned=self.aligned,
            padding_mode=self.padding_mode,
            align_corners=self.align_corners,
        )
        return img_pts