| 81 | return x |
| 82 | |
| 83 | class pointnet_encoder(nn.Module): |
| 84 | def __init__(self): |
| 85 | super(pointnet_encoder, self).__init__() |
| 86 | self.channel = 3 |
| 87 | self.stn = STN3d(self.channel) |
| 88 | self.conv1 = torch.nn.Conv1d(self.channel, 64, 1) |
| 89 | self.conv2 = torch.nn.Conv1d(64, 128, 1) |
| 90 | self.conv3 = torch.nn.Conv1d(128, 128, 1) |
| 91 | self.conv4 = torch.nn.Conv1d(128, 512, 1) |
| 92 | self.conv5 = torch.nn.Conv1d(512, 1024, 1) |
| 93 | self.bn1 = nn.BatchNorm1d(64) |
| 94 | self.bn2 = nn.BatchNorm1d(128) |
| 95 | self.bn3 = nn.BatchNorm1d(128) |
| 96 | self.bn4 = nn.BatchNorm1d(512) |
| 97 | self.bn5 = nn.BatchNorm1d(1024) |
| 98 | self.fstn = STNkd(k = 128) |
| 99 | |
| 100 | def forward(self, point_cloud, return_global): |
| 101 | point_cloud = point_cloud.transpose(2, 1) |
| 102 | B, D, N = point_cloud.size() |
| 103 | assert(D == 3) |
| 104 | #trans = self.stn(point_cloud) |
| 105 | #point_cloud = point_cloud.transpose(2, 1) |
| 106 | #point_cloud = torch.bmm(point_cloud, trans) |
| 107 | #point_cloud = point_cloud.transpose(2, 1) |
| 108 | |
| 109 | out1 = F.relu(self.bn1(self.conv1(point_cloud))) |
| 110 | out2 = F.relu(self.bn2(self.conv2(out1))) |
| 111 | out3 = F.relu(self.bn3(self.conv3(out2))) |
| 112 | |
| 113 | #trans_feat = self.fstn(out3) |
| 114 | #x = out3.transpose(2, 1) |
| 115 | #net_transformed = torch.bmm(x, trans_feat) |
| 116 | #net_transformed = net_transformed.transpose(2, 1) |
| 117 | net_transformed = out3 |
| 118 | |
| 119 | out4 = F.relu(self.bn4(self.conv4(net_transformed))) |
| 120 | out5 = self.bn5(self.conv5(out4)) |
| 121 | out_max = torch.max(out5, 2, keepdim=True)[0] |
| 122 | out_max = out_max.view(-1, 1024) |
| 123 | |
| 124 | if return_global: |
| 125 | return out_max |
| 126 | else: |
| 127 | expand = out_max.view(-1, 1024, 1).repeat(1, 1, N) |
| 128 | concat = torch.cat([point_cloud, expand, out1, out2, out3, out4, out5], 1) |
| 129 | return concat, out_max |