:param unknown: (B, n, 3) tensor of the xyz positions of the unknown features :param known: (B, m, 3) tensor of the xyz positions of the known features :param unknow_feats: (B, C1, n) tensor of the features to be propigated to :param known_feats: (B, C2, m) tensor of
(
self, unknown: torch.Tensor, known: torch.Tensor, unknow_feats: torch.Tensor, known_feats: torch.Tensor
)
| 125 | self.mlp = pt_utils.SharedMLP(mlp, bn=bn) |
| 126 | |
| 127 | def forward( |
| 128 | self, unknown: torch.Tensor, known: torch.Tensor, unknow_feats: torch.Tensor, known_feats: torch.Tensor |
| 129 | ) -> torch.Tensor: |
| 130 | """ |
| 131 | :param unknown: (B, n, 3) tensor of the xyz positions of the unknown features |
| 132 | :param known: (B, m, 3) tensor of the xyz positions of the known features |
| 133 | :param unknow_feats: (B, C1, n) tensor of the features to be propigated to |
| 134 | :param known_feats: (B, C2, m) tensor of features to be propigated |
| 135 | :return: |
| 136 | new_features: (B, mlp[-1], n) tensor of the features of the unknown features |
| 137 | """ |
| 138 | if known is not None: |
| 139 | dist, idx = pointnet2_utils.three_nn(unknown, known) |
| 140 | dist_recip = 1.0 / (dist + 1e-8) |
| 141 | norm = torch.sum(dist_recip, dim=2, keepdim=True) |
| 142 | weight = dist_recip / norm |
| 143 | |
| 144 | interpolated_feats = pointnet2_utils.three_interpolate(known_feats, idx, weight) |
| 145 | else: |
| 146 | interpolated_feats = known_feats.expand(*known_feats.size()[0:2], unknown.size(1)) |
| 147 | |
| 148 | if unknow_feats is not None: |
| 149 | new_features = torch.cat([interpolated_feats, unknow_feats], dim=1) # (B, C2 + C1, n) |
| 150 | else: |
| 151 | new_features = interpolated_feats |
| 152 | |
| 153 | new_features = new_features.unsqueeze(-1) |
| 154 | new_features = self.mlp(new_features) |
| 155 | |
| 156 | return new_features.squeeze(-1) |
| 157 | |
| 158 | |
| 159 | if __name__ == "__main__": |
nothing calls this directly
no outgoing calls
no test coverage detected