|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
from mmcv.cnn import build_norm_layer |
|
|
from mmcv.runner import auto_fp16 |
|
|
from mmcv.utils import Registry |
|
|
from torch import nn |
|
|
from torch.nn import functional as F |
|
|
|
|
|
RADAR_ENCODERS = Registry("radar_encoder") |
|
|
|
|
|
|
|
|
def build_radar_encoder(cfg): |
|
|
"""Build backbone.""" |
|
|
return RADAR_ENCODERS.build(cfg) |
|
|
|
|
|
|
|
|
class RFELayer(nn.Module): |
|
|
"""Radar Feature Encoder layer. |
|
|
|
|
|
Args: |
|
|
in_channels (int): Number of input channels. |
|
|
out_channels (int): Number of output channels. |
|
|
norm_cfg (dict): Config dict of normalization layers |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
in_channels, |
|
|
out_channels, |
|
|
norm_cfg=dict(type="BN1d", eps=1e-3, momentum=0.01), |
|
|
): |
|
|
super(RFELayer, self).__init__() |
|
|
self.fp16_enabled = False |
|
|
self.norm = build_norm_layer(norm_cfg, out_channels)[1] |
|
|
self.linear = nn.Linear(in_channels, out_channels, bias=False) |
|
|
|
|
|
@auto_fp16(apply_to=("inputs"), out_fp32=True) |
|
|
def forward(self, inputs): |
|
|
"""Forward function. |
|
|
|
|
|
Args: |
|
|
inputs (torch.Tensor): Points features of shape (B, M, C). |
|
|
M is the number of points in |
|
|
C is the number of channels of point features. |
|
|
|
|
|
Returns: |
|
|
the same shape |
|
|
""" |
|
|
|
|
|
x = self.linear(inputs) |
|
|
|
|
|
x = self.norm(x.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous() |
|
|
out = F.relu(x) |
|
|
|
|
|
return out |
|
|
|
|
|
|
|
|
@RADAR_ENCODERS.register_module() |
|
|
class RadarPointEncoder(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
in_channels, |
|
|
out_channels, |
|
|
norm_cfg=dict(type="BN1d", eps=1e-3, momentum=0.01), |
|
|
): |
|
|
super(RadarPointEncoder, self).__init__() |
|
|
self.in_channels = in_channels |
|
|
self.out_channels = out_channels |
|
|
in_chn = in_channels |
|
|
|
|
|
layers = [] |
|
|
for out_chn in out_channels: |
|
|
layer = RFELayer(in_chn, out_chn) |
|
|
layers.append(layer) |
|
|
in_chn = out_chn |
|
|
|
|
|
self.feat_layers = nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, points): |
|
|
""" |
|
|
points: [B, N, C]. N: as max |
|
|
masks: [B, N, 1] |
|
|
|
|
|
ret: |
|
|
out: [B, N, C+1], last channel as 0-1 mask |
|
|
""" |
|
|
masks = points[:, :, [-1]] |
|
|
x = points[:, :, :-1] |
|
|
for feat_layer in self.feat_layers: |
|
|
x = feat_layer(x) |
|
|
|
|
|
out = x * masks |
|
|
|
|
|
out = torch.cat((x, masks), dim=-1) |
|
|
return out |
|
|
|
|
|
|
|
|
@RADAR_ENCODERS.register_module() |
|
|
class RadarPointEncoderXYAttn(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
in_channels, |
|
|
out_channels, |
|
|
norm_cfg=dict(type="BN1d", eps=1e-3, momentum=0.01), |
|
|
): |
|
|
super(RadarPointEncoderXYAttn, self).__init__() |
|
|
self.in_channels = in_channels |
|
|
self.out_channels = out_channels |
|
|
in_chn = in_channels |
|
|
|
|
|
layers = [] |
|
|
for out_chn in out_channels: |
|
|
layer = RFELayer(in_chn, out_chn) |
|
|
layers.append(layer) |
|
|
in_chn = out_chn |
|
|
|
|
|
trans_encoder_layer = nn.TransformerEncoderLayer( |
|
|
in_chn, 8, dim_feedforward=2048, dropout=0.1, activation="relu" |
|
|
) |
|
|
|
|
|
self.transformer_encoder = torch.nn.TransformerEncoder( |
|
|
trans_encoder_layer, |
|
|
num_layers=2, |
|
|
) |
|
|
|
|
|
self.feat_layers = nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, points): |
|
|
""" |
|
|
points: [B, N, C]. N: as max |
|
|
masks: [B, N, 1] |
|
|
|
|
|
ret: |
|
|
out: [B, N, C+1], last channel as 0-1 mask |
|
|
""" |
|
|
masks = points[:, :, [-1]] |
|
|
x = points[:, :, :-1] |
|
|
xy = points[:, :, :2] |
|
|
|
|
|
for feat_layer in self.feat_layers: |
|
|
x = feat_layer(x) |
|
|
|
|
|
x = x.permute(1, 0, 2) |
|
|
x = self.transformer_encoder( |
|
|
x, src_key_padding_mask=masks.squeeze(dim=-1).type(torch.bool) |
|
|
) |
|
|
|
|
|
x = x.permute(1, 0, 2) |
|
|
|
|
|
out = x * masks |
|
|
|
|
|
out = torch.cat((x, masks), dim=-1) |
|
|
|
|
|
out = torch.cat((xy, out), dim=-1) |
|
|
return out |
|
|
|
|
|
|
|
|
@RADAR_ENCODERS.register_module() |
|
|
class RadarPointEncoderXY(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
in_channels, |
|
|
out_channels, |
|
|
norm_cfg=dict(type="BN1d", eps=1e-3, momentum=0.01), |
|
|
): |
|
|
super(RadarPointEncoderXY, self).__init__() |
|
|
self.in_channels = in_channels |
|
|
self.out_channels = out_channels |
|
|
in_chn = in_channels |
|
|
|
|
|
layers = [] |
|
|
for out_chn in out_channels: |
|
|
layer = RFELayer(in_chn, out_chn) |
|
|
layers.append(layer) |
|
|
in_chn = out_chn |
|
|
|
|
|
self.feat_layers = nn.Sequential(*layers) |
|
|
|
|
|
def forward(self, points): |
|
|
""" |
|
|
points: [B, N, C]. N: as max |
|
|
masks: [B, N, 1] |
|
|
|
|
|
ret: |
|
|
out: [B, N, C+3], |
|
|
last channel as 0-1 mask |
|
|
first two channes as xy cords |
|
|
""" |
|
|
masks = points[:, :, [-1]] |
|
|
x = points[:, :, :-1] |
|
|
xy = points[:, :, :2] |
|
|
|
|
|
for feat_layer in self.feat_layers: |
|
|
x = feat_layer(x) |
|
|
|
|
|
out = x * masks |
|
|
|
|
|
out = torch.cat((x, masks), dim=-1) |
|
|
|
|
|
out = torch.cat((xy, out), dim=-1) |
|
|
return out |
|
|
|
