Spaces:
Paused
Paused
| import torch | |
| import torch.nn as nn | |
| from StructDiffusion.utils.pointnet import farthest_point_sample, index_points, square_distance | |
| # adapted from https://github.com/qq456cvb/Point-Transformers | |
| def sample_and_group(npoint, nsample, xyz, points): | |
| B, N, C = xyz.shape | |
| S = npoint | |
| fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint] | |
| new_xyz = index_points(xyz, fps_idx) | |
| new_points = index_points(points, fps_idx) | |
| dists = square_distance(new_xyz, xyz) # B x npoint x N | |
| idx = dists.argsort()[:, :, :nsample] # B x npoint x K | |
| grouped_points = index_points(points, idx) | |
| grouped_points_norm = grouped_points - new_points.view(B, S, 1, -1) | |
| new_points = torch.cat([grouped_points_norm, new_points.view(B, S, 1, -1).repeat(1, 1, nsample, 1)], dim=-1) | |
| return new_xyz, new_points | |
| class Local_op(nn.Module): | |
| def __init__(self, in_channels, out_channels): | |
| super().__init__() | |
| self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False) | |
| self.bn1 = nn.BatchNorm1d(out_channels) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| b, n, s, d = x.size() # torch.Size([32, 512, 32, 6]) | |
| x = x.permute(0, 1, 3, 2) | |
| x = x.reshape(-1, d, s) | |
| batch_size, _, N = x.size() | |
| x = self.relu(self.bn1(self.conv1(x))) # B, D, N | |
| x = torch.max(x, 2)[0] | |
| x = x.view(batch_size, -1) | |
| x = x.reshape(b, n, -1).permute(0, 2, 1) | |
| return x | |
| class SA_Layer(nn.Module): | |
| def __init__(self, channels): | |
| super().__init__() | |
| self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False) | |
| self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False) | |
| self.q_conv.weight = self.k_conv.weight | |
| self.v_conv = nn.Conv1d(channels, channels, 1) | |
| self.trans_conv = nn.Conv1d(channels, channels, 1) | |
| self.after_norm = nn.BatchNorm1d(channels) | |
| self.act = nn.ReLU() | |
| self.softmax = nn.Softmax(dim=-1) | |
| def forward(self, x): | |
| x_q = self.q_conv(x).permute(0, 2, 1) # b, n, c | |
| x_k = self.k_conv(x)# b, c, n | |
| x_v = self.v_conv(x) | |
| energy = x_q @ x_k # b, n, n | |
| attention = self.softmax(energy) | |
| attention = attention / (1e-9 + attention.sum(dim=1, keepdims=True)) | |
| x_r = x_v @ attention # b, c, n | |
| x_r = self.act(self.after_norm(self.trans_conv(x - x_r))) | |
| x = x + x_r | |
| return x | |
| class StackedAttention(nn.Module): | |
| def __init__(self, channels=64): | |
| super().__init__() | |
| self.conv1 = nn.Conv1d(channels, channels, kernel_size=1, bias=False) | |
| self.conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=False) | |
| self.bn1 = nn.BatchNorm1d(channels) | |
| self.bn2 = nn.BatchNorm1d(channels) | |
| self.sa1 = SA_Layer(channels) | |
| self.sa2 = SA_Layer(channels) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| # | |
| # b, 3, npoint, nsample | |
| # conv2d 3 -> 128 channels 1, 1 | |
| # b * npoint, c, nsample | |
| # permute reshape | |
| batch_size, _, N = x.size() | |
| x = self.relu(self.bn1(self.conv1(x))) # B, D, N | |
| x = self.relu(self.bn2(self.conv2(x))) | |
| x1 = self.sa1(x) | |
| x2 = self.sa2(x1) | |
| x = torch.cat((x1, x2), dim=1) | |
| return x | |
| class PointTransformerEncoderSmall(nn.Module): | |
| def __init__(self, output_dim=256, input_dim=6, mean_center=True): | |
| super(PointTransformerEncoderSmall, self).__init__() | |
| self.mean_center = mean_center | |
| # map the second dim of the input from input_dim to 64 | |
| self.conv1 = nn.Conv1d(input_dim, 64, kernel_size=1, bias=False) | |
| self.bn1 = nn.BatchNorm1d(64) | |
| self.gather_local_0 = Local_op(in_channels=128, out_channels=64) | |
| self.gather_local_1 = Local_op(in_channels=128, out_channels=64) | |
| self.pt_last = StackedAttention(channels=64) | |
| self.relu = nn.ReLU() | |
| self.conv_fuse = nn.Sequential(nn.Conv1d(192, 256, kernel_size=1, bias=False), | |
| nn.BatchNorm1d(256), | |
| nn.LeakyReLU(negative_slope=0.2)) | |
| self.linear1 = nn.Linear(256, 256, bias=False) | |
| self.bn6 = nn.BatchNorm1d(256) | |
| self.dp1 = nn.Dropout(p=0.5) | |
| self.linear2 = nn.Linear(256, 256) | |
| def forward(self, xyz, f=None): | |
| # xyz: B, N, 3 | |
| # f: B, N, D | |
| center = torch.mean(xyz, dim=1) | |
| if self.mean_center: | |
| xyz = xyz - center.view(-1, 1, 3).repeat(1, xyz.shape[1], 1) | |
| if f is None: | |
| x = self.pct(xyz) | |
| else: | |
| x = self.pct(torch.cat([xyz, f], dim=2)) # B, output_dim | |
| return center, x | |
| def pct(self, x): | |
| # x: B, N, D | |
| xyz = x[..., :3] | |
| x = x.permute(0, 2, 1) | |
| batch_size, _, _ = x.size() | |
| x = self.relu(self.bn1(self.conv1(x))) # B, D, N | |
| x = x.permute(0, 2, 1) | |
| new_xyz, new_feature = sample_and_group(npoint=128, nsample=32, xyz=xyz, points=x) | |
| feature_0 = self.gather_local_0(new_feature) | |
| feature = feature_0.permute(0, 2, 1) # B, nsamples, D | |
| new_xyz, new_feature = sample_and_group(npoint=32, nsample=16, xyz=new_xyz, points=feature) | |
| feature_1 = self.gather_local_1(new_feature) # B, D, nsamples | |
| x = self.pt_last(feature_1) # B, D * 2, nsamples | |
| x = torch.cat([x, feature_1], dim=1) # B, D * 3, nsamples | |
| x = self.conv_fuse(x) | |
| x = torch.max(x, 2)[0] | |
| x = x.view(batch_size, -1) | |
| x = self.relu(self.bn6(self.linear1(x))) | |
| x = self.dp1(x) | |
| x = self.linear2(x) | |
| return x | |
| class SampleAndGroup(nn.Module): | |
| def __init__(self, output_dim=64, input_dim=6, mean_center=True, npoints=(128, 32), nsamples=(32, 16)): | |
| super(SampleAndGroup, self).__init__() | |
| self.mean_center = mean_center | |
| self.npoints = npoints | |
| self.nsamples = nsamples | |
| # map the second dim of the input from input_dim to 64 | |
| self.conv1 = nn.Conv1d(input_dim, output_dim, kernel_size=1, bias=False) | |
| self.bn1 = nn.BatchNorm1d(output_dim) | |
| self.gather_local_0 = Local_op(in_channels=output_dim * 2, out_channels=output_dim) | |
| self.gather_local_1 = Local_op(in_channels=output_dim * 2, out_channels=output_dim) | |
| self.relu = nn.ReLU() | |
| def forward(self, xyz, f): | |
| # xyz: B, N, 3 | |
| # f: B, N, D | |
| center = torch.mean(xyz, dim=1) | |
| if self.mean_center: | |
| xyz = xyz - center.view(-1, 1, 3).repeat(1, xyz.shape[1], 1) | |
| x = self.sg(torch.cat([xyz, f], dim=2)) # B, nsamples, output_dim | |
| return center, x | |
| def sg(self, x): | |
| # x: B, N, D | |
| xyz = x[..., :3] | |
| x = x.permute(0, 2, 1) | |
| batch_size, _, _ = x.size() | |
| x = self.relu(self.bn1(self.conv1(x))) # B, D, N | |
| x = x.permute(0, 2, 1) | |
| new_xyz, new_feature = sample_and_group(npoint=self.npoints[0], nsample=self.nsamples[0], xyz=xyz, points=x) | |
| feature_0 = self.gather_local_0(new_feature) | |
| feature = feature_0.permute(0, 2, 1) # B, nsamples, D | |
| new_xyz, new_feature = sample_and_group(npoint=self.npoints[1], nsample=self.nsamples[1], xyz=new_xyz, points=feature) | |
| feature_1 = self.gather_local_1(new_feature) # B, D, nsamples | |
| x = feature_1.permute(0, 2, 1) # B, nsamples, D | |
| return x |