Class PointSample

embodiedscan/datasets/transforms/points.py:85–222

Point sample. Sampling data to a certain number. Required Keys: - points - pts_instance_mask (optional) - pts_semantic_mask (optional) Modified Keys: - points - pts_instance_mask (optional) - pts_semantic_mask (optional) Args: num_points (int): N

@TRANSFORMS.register_module()
class PointSample(BaseTransform):
    """Point sample.

    Sampling data to a certain number.

    Required Keys:

    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

    Modified Keys:

    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

    Args:
        num_points (int): Number of points to be sampled.
        sample_range (float, optional): The range where to sample points.
            If not None, the points with depth larger than `sample_range` are
            prior to be sampled. Defaults to None.
        replace (bool): Whether the sampling is with or without replacement.
            Defaults to False.
    """

    def __init__(self,
                 num_points: int,
                 sample_range: Optional[float] = None,
                 replace: bool = False) -> None:
        self.num_points = num_points
        self.sample_range = sample_range
        self.replace = replace

    def transform(self, input_dict: dict) -> dict:
        """Transform function to sample points to in indoor scenes.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
            dict: Results after sampling, 'points', 'pts_instance_mask'
            and 'pts_semantic_mask' keys are updated in the result dict.
        """
        points = input_dict['points']

        # if the depth map is all-zero
        if len(points) == 0:
            return input_dict

        points, choices = self._points_random_sampling(points,
                                                       self.num_points,
                                                       self.sample_range,
                                                       self.replace,
                                                       return_choices=True)
        input_dict['points'] = points

        pts_instance_mask = input_dict.get('pts_instance_mask', None)