| 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) |