Find the k nearest points by each ray and rectify the coordinate system as if the rays are in the direction of (0,0,1) and originated from (0,0,0). additional feature inputs may include camera view direction or local frame vectors and these vectors need to be rectified as well (appl
(
points_w: torch.Tensor, # (b, n, 3)
ray_origins_w: torch.Tensor, # (b, m, 3)
ray_directions_w: torch.Tensor, # (b, m, 3)
k: int,
t_min: float = 1.0e-8,
t_max: float = 1.0e10,
t_init: torch.Tensor = None,
other_maps: T.List[torch.Tensor] = None,
rotate_other_maps: T.List[bool] = None,
translate_other_maps: T.List[bool] = None,
max_chunk_size: int = int(1e8),
pr_params: T.Dict[str, T.Any] = None,
randomize_translate: bool = False,
printout: bool = False,
cached_info: T.Union[T.Dict[str, torch.Tensor], None] = None,
valid_mask: torch.Tensor = None, # (b, n, 1)
enable_timing: bool = False,
)
| 12 | |
| 13 | |
| 14 | def find_neighbors_and_rectify( |
| 15 | points_w: torch.Tensor, # (b, n, 3) |
| 16 | ray_origins_w: torch.Tensor, # (b, m, 3) |
| 17 | ray_directions_w: torch.Tensor, # (b, m, 3) |
| 18 | k: int, |
| 19 | t_min: float = 1.0e-8, |
| 20 | t_max: float = 1.0e10, |
| 21 | t_init: torch.Tensor = None, |
| 22 | other_maps: T.List[torch.Tensor] = None, |
| 23 | rotate_other_maps: T.List[bool] = None, |
| 24 | translate_other_maps: T.List[bool] = None, |
| 25 | max_chunk_size: int = int(1e8), |
| 26 | pr_params: T.Dict[str, T.Any] = None, |
| 27 | randomize_translate: bool = False, |
| 28 | printout: bool = False, |
| 29 | cached_info: T.Union[T.Dict[str, torch.Tensor], None] = None, |
| 30 | valid_mask: torch.Tensor = None, # (b, n, 1) |
| 31 | enable_timing: bool = False, |
| 32 | ) -> T.Dict[str, T.Any]: |
| 33 | """ |
| 34 | Find the k nearest points by each ray and rectify the coordinate system as if the rays are |
| 35 | in the direction of (0,0,1) and originated from (0,0,0). |
| 36 | additional feature inputs may include camera view direction or local frame vectors |
| 37 | and these vectors need to be rectified as well (apply rotation only) |
| 38 | |
| 39 | Args: |
| 40 | points_w: (b, n, 3) |
| 41 | ray_origins_w: |
| 42 | ray_directions_w: |
| 43 | k: |
| 44 | t_min: |
| 45 | t_max: |
| 46 | other_maps: |
| 47 | a list of (b, n, d) to associated with each point_w |
| 48 | rotate_other_maps: |
| 49 | a list of bool to indicate whether the feature needs to be rotated, |
| 50 | i.e, multiplied by Rs_w2n |
| 51 | translate_other_maps: |
| 52 | a list of bool to indicate whether the feature needs to be translated |
| 53 | i.e, added by translation_w2n |
| 54 | cached_info: |
| 55 | a dictionary containing the grid cell to point index so pr does not |
| 56 | need to construct it again. |
| 57 | valid_mask: |
| 58 | (b, n, 1) whether the point (including point at inf) should be considered |
| 59 | in the neighbor search. |
| 60 | """ |
| 61 | |
| 62 | if max_chunk_size < 0: |
| 63 | max_chunk_size = int(1e13) |
| 64 | |
| 65 | if other_maps is None: |
| 66 | other_maps = [] |
| 67 | |
| 68 | b, n, _ = points_w.shape |
| 69 | m = ray_origins_w.size(-2) |
| 70 | timing_info = dict() |
| 71 |