Method forward

pointersect/models/pointersect.py:222–543

Args: points: (b, m, k, 3) or (bm, k, 3), m: number of rays, n: number of neighbor points, xyz coord of points additional_features: (b, m, n, dim_point_feature) neighbor_num: (b, m) number of valid neighbors found by pr_cuda. dtype: lo

(
            self,
            points: torch.Tensor,
            additional_features: torch.Tensor = None,
            neighbor_num: torch.Tensor = None,
            valid_mask: torch.Tensor = None,
            printout: bool = False,
            max_chunk_size: int = -1,
            check_finite: bool = False,
    )

                raise NotImplementedError

    def forward(
            self,
            points: torch.Tensor,
            additional_features: torch.Tensor = None,
            neighbor_num: torch.Tensor = None,
            valid_mask: torch.Tensor = None,
            printout: bool = False,
            max_chunk_size: int = -1,
            check_finite: bool = False,
    ):
        """
        Args:
            points: (b, m, k, 3) or (bm, k, 3),  m: number of rays, n: number of neighbor points, xyz coord of points
            additional_features:  (b, m, n, dim_point_feature)
            neighbor_num:
                (b, m)  number of valid neighbors found by pr_cuda.  dtype: long
                pr_cuda returns index (b, m, k), among the k points, only neighbor_num[b, m] is valid,
                others are padding with dummy index 0 (mapped to (10^12, 10^12, 10^12))
            valid_mask:
                (b, m, k) or (bm, k) whether the point is valid to use

        Returns:
            ts: (b, m)
            surface_normals: (b, m, 3)
        """

        if max_chunk_size < 0:
            max_chunk_size = np.inf

        *b_shape, n, dim = points.shape
        device = points.device

        # Note that when we only find neighbors within a fixed distance of ray and the number < n,
        # we use a far away dummy point (1e12,1e12,1e12) to the fill the neighbor position
        # so pr points may contain a far away one

        points = points.reshape(-1, points.size(-2), points.size(-1))  # (b*m, n, 3)

        if valid_mask is not None:
            valid_mask = valid_mask.reshape(-1, points.size(-2))  # (bm, n)

        # pr use neighbor_num to know which points are invalid
        if self.use_pr:
            assert neighbor_num is not None, "need to include neighbor num when use pr"

        assert torch.isfinite(points).all()
        # The position of points can actually be infinite. Since the ray may shoot to opposite direction to all points,
        # the t_s will all be negative and thus become inf.

        # get positional encoding
        xs = self.pos_embedder(points)  # (b*m, n, dim)

        if check_finite and (not torch.isfinite(xs).all() or torch.isnan(xs).any()):
            print('positional embedding is wrong!')

        # concat additional feature
        if additional_features is not None:
            xs = torch.cat(