Given n points (xyz) and m rays, return the neighboring points to each ray. Args: points: (*, n, 3) ray_origins: (*, m, 3) ray_directions: (*, m, 3) k: k nearest neighbors t_min: min t to co
(
points: torch.Tensor,
ray_origins: torch.Tensor,
ray_directions: torch.Tensor,
k: int,
t_min: float = 0.,
t_max: float = 1.e10,
t_init: torch.Tensor = None,
printout: bool = False,
)
| 786 | |
| 787 | |
| 788 | def get_k_neighbor_points( |
| 789 | points: torch.Tensor, |
| 790 | ray_origins: torch.Tensor, |
| 791 | ray_directions: torch.Tensor, |
| 792 | k: int, |
| 793 | t_min: float = 0., |
| 794 | t_max: float = 1.e10, |
| 795 | t_init: torch.Tensor = None, |
| 796 | printout: bool = False, |
| 797 | ) -> T.Dict[str, T.Any]: |
| 798 | """ |
| 799 | Given n points (xyz) and m rays, return the neighboring points to each ray. |
| 800 | |
| 801 | Args: |
| 802 | points: |
| 803 | (*, n, 3) |
| 804 | ray_origins: |
| 805 | (*, m, 3) |
| 806 | ray_directions: |
| 807 | (*, m, 3) |
| 808 | k: |
| 809 | k nearest neighbors |
| 810 | t_min: |
| 811 | min t to consider |
| 812 | t_max: |
| 813 | max t to consider |
| 814 | |
| 815 | Returns: |
| 816 | sorted_dists: |
| 817 | (*, m, min(k, n)) the distance of the k nearest points to each ray (inf if not within t range) |
| 818 | sorted_idxs: |
| 819 | (*, m, min(k,n)) the index of points of the k nearest points |
| 820 | sorted_ts: |
| 821 | (*, m, min(k, n)) projection length of each point on ray from the ray_origins (can be negative) |
| 822 | """ |
| 823 | device = points.device |
| 824 | |
| 825 | if printout: |
| 826 | if torch.cuda.is_available(): |
| 827 | # torch.cuda.synchronize(device) |
| 828 | print(f' before compute ts: {torch.cuda.memory_allocated(device=device) / 2 ** 30} GB') |
| 829 | # torch.cuda.synchronize(device) |
| 830 | |
| 831 | dist_dict = compute_point_ray_distance( |
| 832 | points=points, |
| 833 | ray_origins=ray_origins, |
| 834 | ray_directions=ray_directions, |
| 835 | ) |
| 836 | dists = dist_dict['dists'] # (*, m, n) (ray, point) |
| 837 | ts = dist_dict['ts'] # (*, m, n) |
| 838 | |
| 839 | if printout: |
| 840 | if torch.cuda.is_available(): |
| 841 | # torch.cuda.synchronize(device) |
| 842 | print(f' after compute ts: {torch.cuda.memory_allocated(device=device) / 2 ** 30} GB') |
| 843 | for key in dist_dict: |
| 844 | print(f" {key}: {dist_dict[key].shape} {np.prod(dist_dict[key].shape) * 4 / 2 ** 30:.2f} GB") |
| 845 | # torch.cuda.synchronize(device) |
no test coverage detected