| 23 | |
| 24 | class PoseClassifier(nn.Module): |
| 25 | def __init__(self, pointnet_out_dim=128, pose_dim=6, hidden_dims=(512, 256, 128)): |
| 26 | super(PoseClassifier, self).__init__() |
| 27 | self.pointnet = PointNet(out_channels=(32, 64, pointnet_out_dim)) |
| 28 | input_dim = pointnet_out_dim + pose_dim |
| 29 | layers = [] |
| 30 | for hidden_dim in hidden_dims: |
| 31 | layers.append(nn.Linear(input_dim, hidden_dim)) |
| 32 | layers.append(nn.ReLU()) |
| 33 | input_dim = hidden_dim |
| 34 | layers.append(nn.Linear(input_dim, 1)) |
| 35 | self.classifier = nn.Sequential(*layers) |
| 36 | |
| 37 | def forward(self, point_cloud, poses): |
| 38 | # Point cloud feature extraction |