Args: arr: (b, n, *) bidx: batch index Returns: (n, dim)
(self, arr: torch.Tensor, bidx: int)
| 267 | setattr(self, 'included_point_at_inf', state_dict.get('included_point_at_inf', False)) |
| 268 | |
| 269 | def extract_valid_attr(self, arr: torch.Tensor, bidx: int) -> T.Union[torch.Tensor, None]: |
| 270 | """ |
| 271 | Args: |
| 272 | arr: |
| 273 | (b, n, *) |
| 274 | bidx: |
| 275 | batch index |
| 276 | Returns: |
| 277 | (n, dim) |
| 278 | """ |
| 279 | |
| 280 | if arr is None: |
| 281 | return None |
| 282 | |
| 283 | if self.valid_mask is None: |
| 284 | if not self.included_point_at_inf: |
| 285 | return arr[bidx] # (n, dim) |
| 286 | else: |
| 287 | return arr[bidx, 1:] # (n, dim) |
| 288 | else: |
| 289 | if self.valid_mask.dtype == torch.bool: |
| 290 | valid_mask = self.valid_mask |
| 291 | else: |
| 292 | valid_mask = self.valid_mask > 0.5 |
| 293 | if self.included_point_at_inf: |
| 294 | assert (valid_mask[bidx, 0] < 0.5).all() |
| 295 | arr = arr[bidx, valid_mask[bidx, :, 0]] # (n, dim) |
| 296 | return arr |
| 297 | |
| 298 | def extract_valid_point_cloud(self, bidx: int) -> 'PointCloud': |
| 299 | """ |
no outgoing calls
no test coverage detected