| 7 | |
| 8 | |
| 9 | class Voxelization(nn.Module): |
| 10 | def __init__(self, resolution, normalize=True, eps=0): |
| 11 | super().__init__() |
| 12 | self.r = int(resolution) |
| 13 | self.normalize = normalize |
| 14 | self.eps = eps |
| 15 | |
| 16 | def forward(self, features, coords): |
| 17 | coords = coords.detach() |
| 18 | norm_coords = coords - coords.mean(2, keepdim=True) |
| 19 | if self.normalize: |
| 20 | norm_coords = norm_coords / (norm_coords.norm(dim=1, keepdim=True).max(dim=2, keepdim=True).values * 2.0 + self.eps) + 0.5 |
| 21 | else: |
| 22 | norm_coords = (norm_coords + 1) / 2.0 |
| 23 | norm_coords = torch.clamp(norm_coords * self.r, 0, self.r - 1) |
| 24 | vox_coords = torch.round(norm_coords).to(torch.int32) |
| 25 | return F.avg_voxelize(features, vox_coords, self.r), norm_coords |
| 26 | |
| 27 | def extra_repr(self): |
| 28 | return 'resolution={}{}'.format(self.r, ', normalized eps = {}'.format(self.eps) if self.normalize else '') |