| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import functools |
|
|
| from .pv_module import SharedMLP, PVConv |
|
|
| def create_pointnet_components( |
| blocks, in_channels, with_se=False, normalize=True, eps=0, |
| width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=False, device='cuda'): |
| r, vr = width_multiplier, voxel_resolution_multiplier |
| layers, concat_channels = [], 0 |
| for out_channels, num_blocks, voxel_resolution in blocks: |
| out_channels = int(r * out_channels) |
| if voxel_resolution is None: |
| block = functools.partial(SharedMLP, device=device) |
| else: |
| block = functools.partial( |
| PVConv, kernel_size=3, resolution=int(vr * voxel_resolution), |
| with_se=with_se, normalize=normalize, eps=eps, scale_pvcnn=scale_pvcnn, device=device) |
| for _ in range(num_blocks): |
| layers.append(block(in_channels, out_channels)) |
| in_channels = out_channels |
| concat_channels += out_channels |
| return layers, in_channels, concat_channels |
|
|
| class PCMerger(nn.Module): |
| |
| def __init__(self, in_channels=204, device="cuda"): |
| super(PCMerger, self).__init__() |
| self.mlp_normal = SharedMLP(3, [128, 128], device=device) |
| self.mlp_rgb = SharedMLP(3, [128, 128], device=device) |
| self.mlp_sam = SharedMLP(204 - 6, [128, 128], device=device) |
|
|
| def forward(self, feat, mv_feat, pc2pc_idx): |
| mv_feat_normal = self.mlp_normal(mv_feat[:, :3, :]) |
| mv_feat_rgb = self.mlp_rgb(mv_feat[:, 3:6, :]) |
| mv_feat_sam = self.mlp_sam(mv_feat[:, 6:, :]) |
|
|
| mv_feat_normal = mv_feat_normal.permute(0, 2, 1) |
| mv_feat_rgb = mv_feat_rgb.permute(0, 2, 1) |
| mv_feat_sam = mv_feat_sam.permute(0, 2, 1) |
| feat = feat.permute(0, 2, 1) |
| |
| for i in range(mv_feat.shape[0]): |
| mask = (pc2pc_idx[i] != -1).reshape(-1) |
| idx = pc2pc_idx[i][mask].reshape(-1) |
| feat[i][mask] += mv_feat_normal[i][idx] + mv_feat_rgb[i][idx] + mv_feat_sam[i][idx] |
| |
| return feat.permute(0, 2, 1) |
|
|
|
|
| class PVCNNEncoder(nn.Module): |
| def __init__(self, pvcnn_feat_dim, device='cuda', in_channels=3, use_2d_feat=False): |
| super(PVCNNEncoder, self).__init__() |
| self.device = device |
| self.blocks = ((pvcnn_feat_dim, 1, 32), (128, 2, 16), (256, 1, 8)) |
| self.use_2d_feat=use_2d_feat |
| if in_channels == 6: |
| self.append_channel = 2 |
| elif in_channels == 3: |
| self.append_channel = 1 |
| else: |
| raise NotImplementedError |
| layers, channels_point, concat_channels_point = create_pointnet_components( |
| blocks=self.blocks, in_channels=in_channels + self.append_channel, with_se=False, normalize=False, |
| width_multiplier=1, voxel_resolution_multiplier=1, scale_pvcnn=True, |
| device=device |
| ) |
| self.encoder = nn.ModuleList(layers) |
| if self.use_2d_feat: |
| self.merger = PCMerger() |
|
|
| |
|
|
| def forward(self, input_pc, mv_feat=None, pc2pc_idx=None): |
| features = input_pc.permute(0, 2, 1) * 2 |
| coords = features[:, :3, :] |
| out_features_list = [] |
| voxel_feature_list = [] |
| zero_padding = torch.zeros(features.shape[0], self.append_channel, features.shape[-1], device=features.device, dtype=features.dtype) |
| features = torch.cat([features, zero_padding], dim=1) |
|
|
| for i in range(len(self.encoder)): |
| features, _, voxel_feature = self.encoder[i]((features, coords)) |
| if i == 0 and mv_feat is not None: |
| features = self.merger(features, mv_feat.permute(0, 2, 1), pc2pc_idx) |
| out_features_list.append(features) |
| voxel_feature_list.append(voxel_feature) |
| return voxel_feature_list, out_features_list |
