| 22 | return x |
| 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 |
| 39 | point_cloud_features = self.pointnet(point_cloud) # (batch_size, pointnet_out_dim) |
| 40 | |
| 41 | # Repeat point cloud features for each pose |
| 42 | repeated_features = point_cloud_features.unsqueeze(1).repeat(1, poses.size(1), 1) # (batch_size, num_poses, pointnet_out_dim) |
| 43 | |
| 44 | # Concatenate pose features with point cloud features |
| 45 | combined_features = torch.cat((repeated_features, poses), dim=-1) # (batch_size, num_poses, pointnet_out_dim + pose_dim) |
| 46 | |
| 47 | # Flatten the input for the classifier |
| 48 | combined_features = combined_features.view(-1, combined_features.size(-1)) # (batch_size * num_poses, pointnet_out_dim + pose_dim) |
| 49 | |
| 50 | # Classification |
| 51 | scores = self.classifier(combined_features) # (batch_size * num_poses, 1) |
| 52 | scores = scores.view(-1, poses.size(1)) # (batch_size, num_poses) |
| 53 | |
| 54 | return scores |
| 55 | |
| 56 | |
| 57 | batch_size = 8 |