| 59 | |
| 60 | |
| 61 | class SHEncoder(nn.Module): |
| 62 | def __init__(self, input_dim=3, degree=4): |
| 63 | super().__init__() |
| 64 | |
| 65 | self.input_dim = input_dim # coord dims, must be 3 |
| 66 | self.degree = degree # 0 ~ 4 |
| 67 | self.output_dim = degree ** 2 |
| 68 | |
| 69 | assert self.input_dim == 3, "SH encoder only support input dim == 3" |
| 70 | assert self.degree > 0 and self.degree <= 8, "SH encoder only supports degree in [1, 8]" |
| 71 | |
| 72 | def __repr__(self): |
| 73 | return f"SHEncoder: input_dim={self.input_dim} degree={self.degree}" |
| 74 | |
| 75 | def forward(self, inputs, size=1): |
| 76 | # inputs: [..., input_dim], normalized real world positions in [-size, size] |
| 77 | # return: [..., degree^2] |
| 78 | |
| 79 | inputs = inputs / size # [-1, 1] |
| 80 | |
| 81 | prefix_shape = list(inputs.shape[:-1]) |
| 82 | inputs = inputs.reshape(-1, self.input_dim) |
| 83 | |
| 84 | outputs = sh_encode(inputs, self.degree, inputs.requires_grad) |
| 85 | outputs = outputs.reshape(prefix_shape + [self.output_dim]) |
| 86 | |
| 87 | return outputs |