(
self,
output: torch.Tensor,
rescale_colors: bool = False,
)
| 205 | return model_fn |
| 206 | |
| 207 | def split_model_output( |
| 208 | self, |
| 209 | output: torch.Tensor, |
| 210 | rescale_colors: bool = False, |
| 211 | ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: |
| 212 | assert ( |
| 213 | len(self.aux_channels) + 3 == output.shape[1] |
| 214 | ), "there must be three spatial channels before aux" |
| 215 | pos, joined_aux = output[:, :3], output[:, 3:] |
| 216 | |
| 217 | aux = {} |
| 218 | for i, name in enumerate(self.aux_channels): |
| 219 | v = joined_aux[:, i] |
| 220 | if name in {"R", "G", "B", "A"}: |
| 221 | v = v.clamp(0, 255).round() |
| 222 | if rescale_colors: |
| 223 | v = v / 255.0 |
| 224 | aux[name] = v |
| 225 | return pos, aux |
| 226 | |
| 227 | def output_to_point_clouds(self, output: torch.Tensor) -> List[PointCloud]: |
| 228 | res = [] |
no outgoing calls
no test coverage detected