| import torch | |
| import torch.nn as nn | |
| from . import functional as F | |
| from .ball_query import BallQuery | |
| from .shared_mlp import SharedMLP | |
| __all__ = ['PointNetAModule', 'PointNetSAModule', 'PointNetFPModule'] | |
| class PointNetAModule(nn.Module): | |
| def __init__(self, in_channels, out_channels, include_coordinates=True): | |
| super().__init__() | |
| if not isinstance(out_channels, (list, tuple)): | |
| out_channels = [[out_channels]] | |
| elif not isinstance(out_channels[0], (list, tuple)): | |
| out_channels = [out_channels] | |
| mlps = [] | |
| total_out_channels = 0 | |
| for _out_channels in out_channels: | |
| mlps.append( | |
| SharedMLP( | |
| in_channels=in_channels + (3 if include_coordinates else 0), | |
| out_channels=_out_channels, dim=1) | |
| ) | |
| total_out_channels += _out_channels[-1] | |
| self.include_coordinates = include_coordinates | |
| self.out_channels = total_out_channels | |
| self.mlps = nn.ModuleList(mlps) | |
| def forward(self, inputs): | |
| features, coords = inputs | |
| if self.include_coordinates: | |
| features = torch.cat([features, coords], dim=1) | |
| coords = torch.zeros((coords.size(0), 3, 1), device=coords.device) | |
| if len(self.mlps) > 1: | |
| features_list = [] | |
| for mlp in self.mlps: | |
| features_list.append(mlp(features).max(dim=-1, keepdim=True).values) | |
| return torch.cat(features_list, dim=1), coords | |
| else: | |
| return self.mlps[0](features).max(dim=-1, keepdim=True).values, coords | |
| def extra_repr(self): | |
| return f'out_channels={self.out_channels}, include_coordinates={self.include_coordinates}' | |
| class PointNetSAModule(nn.Module): | |
| def __init__(self, num_centers, radius, num_neighbors, in_channels, out_channels, include_coordinates=True): | |
| super().__init__() | |
| if not isinstance(radius, (list, tuple)): | |
| radius = [radius] | |
| if not isinstance(num_neighbors, (list, tuple)): | |
| num_neighbors = [num_neighbors] * len(radius) | |
| assert len(radius) == len(num_neighbors) | |
| if not isinstance(out_channels, (list, tuple)): | |
| out_channels = [[out_channels]] * len(radius) | |
| elif not isinstance(out_channels[0], (list, tuple)): | |
| out_channels = [out_channels] * len(radius) | |
| assert len(radius) == len(out_channels) | |
| groupers, mlps = [], [] | |
| total_out_channels = 0 | |
| for _radius, _out_channels, _num_neighbors in zip(radius, out_channels, num_neighbors): | |
| groupers.append( | |
| BallQuery(radius=_radius, num_neighbors=_num_neighbors, include_coordinates=include_coordinates) | |
| ) | |
| mlps.append( | |
| SharedMLP( | |
| in_channels=in_channels + (3 if include_coordinates else 0), | |
| out_channels=_out_channels, dim=2) | |
| ) | |
| total_out_channels += _out_channels[-1] | |
| self.num_centers = num_centers | |
| self.out_channels = total_out_channels | |
| self.groupers = nn.ModuleList(groupers) | |
| self.mlps = nn.ModuleList(mlps) | |
| def forward(self, inputs): | |
| features, coords = inputs | |
| centers_coords = F.furthest_point_sample(coords, self.num_centers) | |
| features_list = [] | |
| for grouper, mlp in zip(self.groupers, self.mlps): | |
| features_list.append(mlp(grouper(coords, centers_coords, features)).max(dim=-1).values) | |
| if len(features_list) > 1: | |
| return torch.cat(features_list, dim=1), centers_coords | |
| else: | |
| return features_list[0], centers_coords | |
| def extra_repr(self): | |
| return f'num_centers={self.num_centers}, out_channels={self.out_channels}' | |
| class PointNetFPModule(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super().__init__() | |
| self.mlp = SharedMLP(in_channels=in_channels, out_channels=out_channels, dim=1) | |
| def forward(self, inputs): | |
| if len(inputs) == 3: | |
| points_coords, centers_coords, centers_features = inputs | |
| points_features = None | |
| else: | |
| points_coords, centers_coords, centers_features, points_features = inputs | |
| interpolated_features = F.nearest_neighbor_interpolate(points_coords, centers_coords, centers_features) | |
| if points_features is not None: | |
| interpolated_features = torch.cat( | |
| [interpolated_features, points_features], dim=1 | |
| ) | |
| return self.mlp(interpolated_features), points_coords | |