SamplingPC

Classes for point cloud sampling

Sampling

Subsample points (and aligned attributes) using random or FPS-based indices.

This transform expects a dictionary containing:

"coord": NumPy array of shape (N, 3) with point coordinates.
Optionally other per-point arrays (e.g., "norm", "color", "label") that have length N along the first dimension.
For FPS-based methods ("random_fps" and "fps"), a key "fps_index" is also required, containing a 1D array of precomputed farthest-point-sampling indices into "coord".

If the requested number of points n_pts is less than the current number of points N, it selects a subset of indices and applies the same index selection to all aligned per-point fields (except keys containing "origin"). If n_pts >= N, no subsampling is applied.

The sampling strategy is controlled by method:

"random":
Uniformly sample n_pts indices from [0, N) without replacement.
"random_fps":
Sample n_pts indices from fps_index without replacement, then map the index to get the final subset. (it is for better augmentation instead of fixed fps index each time despite not guarantee perfect uniform coverage)
"fps":
Directly use the first n_pts entries from data_dict["fps_index"].

Parameters:

Name	Type	Description	Default
`n_pts`	`int`	Target number of points after sampling. If `n_pts >= N`, no sampling is applied. Defaults to 1024.	`1024`
`method`	`str`	Sampling strategy. One of: `"random"`: uniform random sampling. `"random_fps"`: random subset of precomputed FPS indices. `"fps"`: first `n_pts` precomputed FPS indices. Defaults to `"fps"`.	`'fps'`

Source code in src\augmentation_class.py

@TRANSFORMS.register()
class Sampling:
    """Subsample points (and aligned attributes) using random or FPS-based indices.

    This transform expects a dictionary containing:

    * `"coord"`: NumPy array of shape (N, 3) with point coordinates.
    * Optionally other per-point arrays (e.g., `"norm"`, `"color"`, `"label"`) that have length N along the first dimension.
    * For FPS-based methods (`"random_fps"` and `"fps"`), a key `"fps_index"` is also required, containing a 1D array of
      precomputed farthest-point-sampling indices into `"coord"`.

    If the requested number of points `n_pts` is **less** than the current number of points `N`, it selects a subset of
    indices and applies the same index selection to all aligned per-point fields (except keys containing `"origin"`).
    If `n_pts >= N`, no subsampling is applied.

    The sampling strategy is controlled by `method`:

    * `"random"`:
      - Uniformly sample `n_pts` indices from `[0, N)` without replacement.
    * `"random_fps"`:
      - Sample `n_pts` indices from `fps_index` without replacement, then map the index to get the final subset.
      (it is for better augmentation instead of fixed fps index each time despite not guarantee perfect uniform coverage)
    * `"fps"`:
      - Directly use the first `n_pts` entries from `data_dict["fps_index"]`.

    Args:
        n_pts (int, optional):
            Target number of points after sampling. If `n_pts >= N`, no sampling is applied. Defaults to 1024.
        method (str, optional):
            Sampling strategy. One of:

            * `"random"`: uniform random sampling.
            * `"random_fps"`: random subset of precomputed FPS indices.
            * `"fps"`: first `n_pts` precomputed FPS indices.

            Defaults to `"fps"`.
    """
    def __init__(self, n_pts=1024, method="fps"):
        self.n_pts = n_pts
        self.method = method
        assert method in ["random", "random_fps", "fps"]

    def __call__(self, data_dict: dict) -> dict:
        """Subsample points and aligned fields according to the chosen method.

        Args:
            data_dict (dict): Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3).
                Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain
                `"origin"` will be permuted with the same shuffle indices.

        Returns:
            dict: The same dictionary with `"coord"` and aligned per-point attributes subsampled to at most `n_pts` points, if applied.
        """
        # print(data_dict.keys())
        if "coord" in data_dict.keys():
            N = len(data_dict["coord"])
            if self.n_pts < N :
                if self.method == "random":
                    idx = np.random.choice(N, self.n_pts, replace=False)
                elif self.method == "random_fps":
                    assert self.n_pts <= len(data_dict["fps_index"])
                    idx = np.random.choice(len(data_dict["fps_index"]), self.n_pts, replace=False)
                    idx = data_dict["fps_index"][idx]
                elif self.method == "fps":
                    assert self.n_pts <= len(data_dict["fps_index"])
                    idx = data_dict["fps_index"][:self.n_pts]
                else:
                    raise NotImplementedError(f"method {self.method} is not supported.")

                for key, val in data_dict.items():
                    if isinstance(val, (np.ndarray, Sequence)) and len(val) == N and "origin" not in key:
                        data_dict[key] = val[idx]

        return data_dict

`call(data_dict)`

Subsample points and aligned fields according to the chosen method.

Parameters:

Name	Type	Description	Default
`data_dict`	`dict`	Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3). Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain `"origin"` will be permuted with the same shuffle indices.	required

Returns:

Name	Type	Description
`dict`	`dict`	The same dictionary with `"coord"` and aligned per-point attributes subsampled to at most `n_pts` points, if applied.

Source code in src\augmentation_class.py

def __call__(self, data_dict: dict) -> dict:
    """Subsample points and aligned fields according to the chosen method.

    Args:
        data_dict (dict): Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3).
            Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain
            `"origin"` will be permuted with the same shuffle indices.

    Returns:
        dict: The same dictionary with `"coord"` and aligned per-point attributes subsampled to at most `n_pts` points, if applied.
    """
    # print(data_dict.keys())
    if "coord" in data_dict.keys():
        N = len(data_dict["coord"])
        if self.n_pts < N :
            if self.method == "random":
                idx = np.random.choice(N, self.n_pts, replace=False)
            elif self.method == "random_fps":
                assert self.n_pts <= len(data_dict["fps_index"])
                idx = np.random.choice(len(data_dict["fps_index"]), self.n_pts, replace=False)
                idx = data_dict["fps_index"][idx]
            elif self.method == "fps":
                assert self.n_pts <= len(data_dict["fps_index"])
                idx = data_dict["fps_index"][:self.n_pts]
            else:
                raise NotImplementedError(f"method {self.method} is not supported.")

            for key, val in data_dict.items():
                if isinstance(val, (np.ndarray, Sequence)) and len(val) == N and "origin" not in key:
                    data_dict[key] = val[idx]

    return data_dict

FPS PC with 4096 points

Before After

SamplingDynamic

Dynamically choose the number of sampled points based on a scalar attribute.

This transform expects a dictionary containing:

"coord": NumPy array of shape (N, 3) with point coordinates.
Optionally other per-point arrays (e.g., "norm", "color", "label") that have length N along the first dimension.
"fps_index": 1D NumPy array of precomputed farthest-point-sampling indices into "coord".
A scalar entry key (default "area") used to determine how many points to sample.

The number of target points is computed as:

pts = int(data_dict[key] * pts_ratio)

Then, from "fps_index" it selects:

The first pts indices if len(fps_index) > pts, or
All of fps_index otherwise.

All per-point arrays of length N (except those whose key contains "origin") are then indexed with this subset.

Parameters:

Name	Type	Description	Default
`key`	`str`	Name of the scalar field in `data_dict` used to determine the dynamic number of points. Common choices might include `"area"`, `"volume"`, etc. Defaults to `"area"`.	`'area'`
`pts_ratio`	`float`	Multiplicative factor that maps the scalar value `data_dict[key]` to a target number of points: `pts = int(data_dict[key] * pts_ratio)` For example, if `data_dict["area"] = 1.8` and `pts_ratio` is `8192 / 1.8`, then `pts ≈ 8192`. Defaults to `8192 / 1.8`.	`8192 / 1.8`

Source code in src\augmentation_class.py

@TRANSFORMS.register()
class SamplingDynamic:
    """Dynamically choose the number of sampled points based on a scalar attribute.

    This transform expects a dictionary containing:

    * `"coord"`: NumPy array of shape (N, 3) with point coordinates.
    * Optionally other per-point arrays (e.g., `"norm"`, `"color"`, `"label"`) that have length N along the first dimension.
    * `"fps_index"`: 1D NumPy array of precomputed farthest-point-sampling indices into `"coord"`.
    * A scalar entry `key` (default `"area"`) used to determine how many points to sample.

    The number of target points is computed as:

        pts = int(data_dict[key] * pts_ratio)

    Then, from `"fps_index"` it selects:

    * The first `pts` indices if `len(fps_index) > pts`, or
    * All of `fps_index` otherwise.

    All per-point arrays of length `N` (except those whose key contains `"origin"`) are then indexed with this subset.

    Args:
        key (str, optional):
            Name of the scalar field in `data_dict` used to determine the dynamic number of points. Common choices might
            include `"area"`, `"volume"`, etc.
            Defaults to `"area"`.
        pts_ratio (float, optional):
            Multiplicative factor that maps the scalar value `data_dict[key]` to a target number of points:

                pts = int(data_dict[key] * pts_ratio)

            For example, if `data_dict["area"] = 1.8` and `pts_ratio` is `8192 / 1.8`, then `pts ≈ 8192`.
            Defaults to `8192 / 1.8`.
    """
    def __init__(self, key="area", pts_ratio=8192/1.8):
        self.pts_ratio = pts_ratio
        self.key = key

    def __call__(self, data_dict: dict) -> dict:
        """Apply dynamic FPS-based subsampling to `"coord"` and aligned fields.

        Args:
            data_dict (dict): Input dictionary containing at least `"coord"`, `self.key`, and `"fps_index"`.

        Returns:
            dict: The same dictionary with `"coord"` and aligned per-point attributes subsampled according to the
            dynamically chosen number of points.
        """
        if "coord" in data_dict.keys():
            N = len(data_dict["coord"])
            assert self.key in data_dict
            assert "fps_index" in data_dict

            pts = int(data_dict[self.key] * self.pts_ratio)
            if len(data_dict["fps_index"]) > pts:
                idx = data_dict["fps_index"][:pts]
            else:
                idx = data_dict["fps_index"]

            for key, val in data_dict.items():
                if isinstance(val, (np.ndarray, Sequence)) and len(val) == N and "origin" not in key:
                    data_dict[key] = val[idx]

        return data_dict

`call(data_dict)`

Apply dynamic FPS-based subsampling to "coord" and aligned fields.

Parameters:

Name	Type	Description	Default
`data_dict`	`dict`	Input dictionary containing at least `"coord"`, `self.key`, and `"fps_index"`.	required

Returns:

Name	Type	Description
`dict`	`dict`	The same dictionary with `"coord"` and aligned per-point attributes subsampled according to the
	`dict`	dynamically chosen number of points.

Source code in src\augmentation_class.py

def __call__(self, data_dict: dict) -> dict:
    """Apply dynamic FPS-based subsampling to `"coord"` and aligned fields.

    Args:
        data_dict (dict): Input dictionary containing at least `"coord"`, `self.key`, and `"fps_index"`.

    Returns:
        dict: The same dictionary with `"coord"` and aligned per-point attributes subsampled according to the
        dynamically chosen number of points.
    """
    if "coord" in data_dict.keys():
        N = len(data_dict["coord"])
        assert self.key in data_dict
        assert "fps_index" in data_dict

        pts = int(data_dict[self.key] * self.pts_ratio)
        if len(data_dict["fps_index"]) > pts:
            idx = data_dict["fps_index"][:pts]
        else:
            idx = data_dict["fps_index"]

        for key, val in data_dict.items():
            if isinstance(val, (np.ndarray, Sequence)) and len(val) == N and "origin" not in key:
                data_dict[key] = val[idx]

    return data_dict

FPS PC with surface area adjusted points

Before After

SphereCrop

Crop a point cloud to a fixed number of points using a spherical region.

This transform expects a dictionary containing:

"coord": NumPy array of shape (N, 3) with point coordinates.
Optionally other per-point arrays (e.g., "norm", "color", "label") that have length N along the first dimension.

The target number of points can be controlled either by an absolute cap (point_max) or a relative rate (sample_rate):

If sample_rate is not None, the effective maximum is point_max_eff = int(sample_rate * N).
Otherwise, the fixed point_max value is used.

The center of the crop is chosen according to mode:

"random": use a randomly selected point as center.
"center": use the point at index N // 2 as center (e.g., middle in the current ordering).

If N <= point_max_eff, no cropping is applied.

Parameters:

Name	Type	Description	Default
`point_max`	`int`	Maximum number of points to keep if`sample_rate` is `None`. Defaults to 80,000.	`80000`
`sample_rate`	`float \| None`	If provided, the effective maximum number of points is computed as: `point_max_eff = int(sample_rate * N)` where `N` is the current number of points. This allows dataset-dependent cropping. If `None`, the fixed `point_max` is used instead. Defaults to `None`.	`None`
`mode`	`str`	Strategy to select the crop center. Must be `"random"` or `"center"`. `"random"`: center is a random point from `"coord"`. `"center"`: center is `coord[N // 2]`. Defaults to `"random"`.	`'random'`

Source code in src\augmentation_class.py

@TRANSFORMS.register()
class SphereCrop:
    """Crop a point cloud to a fixed number of points using a spherical region.

    This transform expects a dictionary containing:

    * `"coord"`: NumPy array of shape (N, 3) with point coordinates.
    * Optionally other per-point arrays (e.g., `"norm"`, `"color"`, `"label"`) that have length N along the first dimension.

    The target number of points can be controlled either by an absolute cap (`point_max`) or a relative rate (`sample_rate`):

    * If `sample_rate` is not ``None``, the effective maximum is `point_max_eff = int(sample_rate * N)`.
    * Otherwise, the fixed `point_max` value is used.

    The center of the crop is chosen according to `mode`:

    * `"random"`: use a randomly selected point as center.
    * `"center"`: use the point at index `N // 2` as center (e.g., middle in
      the current ordering).

    If `N <= point_max_eff`, no cropping is applied.

    Args:
        point_max (int, optional):
            Maximum number of points to keep if`sample_rate` is ``None``.
            Defaults to 80,000.
        sample_rate (float | None, optional):
            If provided, the effective maximum number of points is computed as:

                point_max_eff = int(sample_rate * N)

            where `N` is the current number of points. This allows dataset-dependent cropping. If ``None``, the fixed `point_max`
            is used instead.
            Defaults to ``None``.
        mode (str, optional):
            Strategy to select the crop center. Must be `"random"` or `"center"`.

            * `"random"`: center is a random point from `"coord"`.
            * `"center"`: center is `coord[N // 2]`.

            Defaults to `"random"`.
    """
    def __init__(self, point_max: int = 80000, sample_rate: float = None, mode: str = "random"):
        self.point_max = point_max
        self.sample_rate = sample_rate
        assert mode in ["random", "center"]
        self.mode = mode

    def __call__(self, data_dict: dict) -> dict:
        """Apply spherical cropping to the `"coord"` entry in `data_dict`.

        Args:
            data_dict (dict): Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3).
                Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain
                `"origin"` will be permuted with the same shuffle indices.

        Returns:
            dict: The same dictionary with `"coord"` and aligned per-point attributes cropped to at most `point_max_eff`
            points, if applied.
        """
        if "coord" in data_dict.keys():
            n = len(data_dict["coord"])
            if self.sample_rate is not None:
                point_max = int(self.sample_rate * n)
            else:
                point_max = self.point_max

            # mode is "random" or "center"
            if n > point_max:
                if self.mode == "random":
                    center = data_dict["coord"][np.random.randint(n)]
                elif self.mode == "center":
                    center = data_dict["coord"][n // 2]
                else:
                    raise NotImplementedError

                idx = np.argsort(np.sum(np.square(data_dict["coord"] - center), 1))[: point_max]
                for key, value in data_dict.items():
                    if isinstance(value, (np.ndarray, Sequence)) and len(value) == n and "origin" not in key:
                        data_dict[key] = value[idx]

        return data_dict

`call(data_dict)`

Apply spherical cropping to the "coord" entry in data_dict.

Parameters:

Name	Type	Description	Default
`data_dict`	`dict`	Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3). Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain `"origin"` will be permuted with the same shuffle indices.	required

Returns:

Name	Type	Description
`dict`	`dict`	The same dictionary with `"coord"` and aligned per-point attributes cropped to at most `point_max_eff`
	`dict`	points, if applied.

Source code in src\augmentation_class.py

def __call__(self, data_dict: dict) -> dict:
    """Apply spherical cropping to the `"coord"` entry in `data_dict`.

    Args:
        data_dict (dict): Input dictionary that must contain a `"coord"` key with a NumPy array of shape (N, 3).
            Any other entry whose value is a NumPy array or `Sequence` of length N and whose key does not contain
            `"origin"` will be permuted with the same shuffle indices.

    Returns:
        dict: The same dictionary with `"coord"` and aligned per-point attributes cropped to at most `point_max_eff`
        points, if applied.
    """
    if "coord" in data_dict.keys():
        n = len(data_dict["coord"])
        if self.sample_rate is not None:
            point_max = int(self.sample_rate * n)
        else:
            point_max = self.point_max

        # mode is "random" or "center"
        if n > point_max:
            if self.mode == "random":
                center = data_dict["coord"][np.random.randint(n)]
            elif self.mode == "center":
                center = data_dict["coord"][n // 2]
            else:
                raise NotImplementedError

            idx = np.argsort(np.sum(np.square(data_dict["coord"] - center), 1))[: point_max]
            for key, value in data_dict.items():
                if isinstance(value, (np.ndarray, Sequence)) and len(value) == n and "origin" not in key:
                    data_dict[key] = value[idx]

    return data_dict

Sphere Crop PC

Before After

GridSample

Grid-based voxelization and optional per-voxel sampling/pooling.

This transform expects a dictionary containing:

"coord": NumPy array of shape (N, 3) with point coordinates.
Optionally per-point attributes (e.g. "norm", "color", "label") that have length N along the first dimension.

It first normalizes coordinates by subtracting the global minimum so that all coordinates are non-negative, then assigns each point to a 3D grid cell (voxel) either by:

a fixed cell size (grid_size), or
a fixed number of cells per axis (grid_number).

The per-point integer voxel index is stored in "grid_coord".

Optionally, it can:

Store per-point relative coordinates within each voxel in "relative_coord", normalized to approximately [-1, 1] per axis.
Perform voxel-level sampling/aggregation:
method="random":
- mode="train": pick one random point per voxel.
- mode="test": generate up to max_count “views”, each view containing one point per voxel (cycling through the points in each voxel). Returns a list of dictionaries.
method="mean":
- Average features per voxel and replace "coord", "norm", "color" accordingly.
- For "label":
- If integer dtype → take majority vote per voxel.
- Else → average values per voxel.

Parameters:

Name	Type	Description	Default
`grid_size`	`float`	Fixed grid cell size along each axis when `grid_number` is `None`. The same size is used for x, y, z. Coordinates are voxelized as: `grid_coord = floor((coord - coord_min) / grid_size)` Defaults to 0.02.	`0.02`
`grid_number`	`tuple[int] \| None`	If not `None`, use a fixed number of cells per axis instead of a fixed size. In that case, the effective grid size is computed as: `grid_size = coord_norm.max(axis=0) / grid_number` and voxel indices are clipped to `[0, grid_number-1]`. Defaults to `None`.	`None`
`sampling`	`bool`	Whether to perform per-voxel sampling or pooling. If False, the transform only computes `"grid_coord"` (and `"relative_coord"` if `return_relative=True`) and does not change the number of points. Defaults to True.	`True`
`method`	`str`	Sampling/pooling strategy when `sampling=True`. Supported values: `"random"`: per-voxel random selection. `"mean"`: per-voxel averaging of features (and label fusion). Defaults to `"random"`.	`'random'`
`return_relative`	`bool`	If True, adds `"relative_coord"` to `data_dict`, containing per-point offsets relative to the voxel center, normalized to approximately `[-1, 1]` per axis. Defaults to False.	`False`
`mode`	`str`	Behavior mode for `method="random"`. * `"train"`: returns a single dictionary, selecting one random point per voxel. * `"test"`: returns a list of dictionaries, each containing one point per voxel, cycling through all points within each voxel. Defaults to `"train"`.	`'train'`

Source code in src\augmentation_class.py

@TRANSFORMS.register()
class GridSample:
    """Grid-based voxelization and optional per-voxel sampling/pooling.

    This transform expects a dictionary containing:

    * `"coord"`: NumPy array of shape (N, 3) with point coordinates.
    * Optionally per-point attributes (e.g. `"norm"`, `"color"`, `"label"`) that have length N along the first dimension.

    It first normalizes coordinates by subtracting the global minimum so that all coordinates are non-negative,
    then assigns each point to a 3D grid cell (voxel) either by:

    * a fixed cell size (`grid_size`), or
    * a fixed number of cells per axis (`grid_number`).

    The per-point integer voxel index is stored in `"grid_coord"`.

    Optionally, it can:

    * Store per-point relative coordinates within each voxel in `"relative_coord"`, normalized to approximately `[-1, 1]` per axis.

    * Perform voxel-level sampling/aggregation:

      - `method="random"`:
        * `mode="train"`: pick one random point per voxel.
        * `mode="test"`: generate up to `max_count` “views”, each view containing one point per voxel
        (cycling through the points in each voxel). Returns a list of dictionaries.

      - `method="mean"`:
        * Average features per voxel and replace `"coord"`, `"norm"`, `"color"` accordingly.
        * For `"label"`:
          - If integer dtype → take majority vote per voxel.
          - Else → average values per voxel.

    Args:
        grid_size (float, optional):
            Fixed grid cell size along each axis when `grid_number` is ``None``. The
            same size is used for x, y, z. Coordinates are voxelized as:

                grid_coord = floor((coord - coord_min) / grid_size)

            Defaults to 0.02.
        grid_number (tuple[int] | None, optional):
            If not ``None``, use a fixed number of cells per axis instead of a fixed size. In that case, the effective
            grid size is computed as:

                grid_size = coord_norm.max(axis=0) / grid_number

            and voxel indices are clipped to `[0, grid_number-1]`.
            Defaults to ``None``.
        sampling (bool, optional):
            Whether to perform per-voxel sampling or pooling. If False, the transform only computes `"grid_coord"`
            (and `"relative_coord"` if `return_relative=True`) and does not change the number of points.
            Defaults to True.
        method (str, optional):
            Sampling/pooling strategy when `sampling=True`. Supported values:

            * `"random"`: per-voxel random selection.
            * `"mean"`: per-voxel averaging of features (and label fusion).

            Defaults to `"random"`.
        return_relative (bool, optional):
            If True, adds `"relative_coord"` to `data_dict`, containing per-point offsets relative to the voxel center,
            normalized to approximately `[-1, 1]` per axis.
            Defaults to False.
        mode (str, optional):
            Behavior mode for `method="random"`.
            * `"train"`: returns a single dictionary, selecting one random point per voxel.
            * `"test"`: returns a list of dictionaries, each containing one point per voxel, cycling through all points within each voxel.

            Defaults to `"train"`.
    """

    def __init__(self, grid_size: float = 0.02, grid_number: tuple[int] = None, sampling: bool = True,
                 method: str = "random", return_relative: bool = False, mode: str = "train"):
        self.grid_number = grid_number
        self.grid_size = grid_size
        self.sampling = sampling
        assert method in ["random", "mean"]
        self.method = method
        self.return_relative = return_relative
        assert mode in ["train", "test"]
        self.mode = mode

    @staticmethod
    def _ravel_3d(arr_grid):
        """Convert 3D integer grid coordinates to a 1D key via row-major flattening.

        Given integer grid coordinates of shape (N, 3), this computes a unique 1D key for each voxel coordinate
        such that points in the same voxel share the same key. This is used for grouping points by voxel
        (e.g. for random voxel-wise sampling).

        Args:
            arr_grid (np.ndarray): Integer grid coordinates of shape (N, 3).

        Returns:
            np.ndarray: 1D array of length N containing raveled voxel keys.
        """
        assert arr_grid.ndim == 2 and arr_grid.shape[-1] == 3
        max_grid = np.max(arr_grid, axis=0)
        keys = arr_grid[:, 0] + arr_grid[:, 1] * (max_grid[0] + 1) + arr_grid[:, 2] * (max_grid[0] + 1) * (
                max_grid[1] + 1)
        return keys

    def __call__(self, data_dict: dict) -> dict | list:
        """Apply grid voxelization and optional voxel-wise sampling/pooling.


        Args:
            data_dict (dict): Input dictionary that must contain `"coord"`.
                Optionally `"norm"`, `"color"`, `"label"`, etc.

        Returns:
            dict | list[dict]:
                * If `sampling=False` or `method="mean"` or `mode="train"`:
                  returns a single modified `data_dict`.
                * If `method="random"` and `mode="test"`: returns a list of dictionaries, each representing a
                  different per-voxel sampling pass.
        """
        if "coord" in data_dict.keys():
            coord = data_dict["coord"]
            coord_min = coord.min(axis=0)
            coord_norm = coord - coord_min
            # kill tiny negative noise due to FP
            coord_norm = np.clip(coord_norm, 0.0, None)  # everything < 0 → 0.0

            if self.grid_number is None:
                # fixed grid cell size
                grid_size = self.grid_size
                grid_coord = np.floor(coord_norm / grid_size).astype(int)
            else:
                # fixed number of cells per axis
                grid_number = np.array(self.grid_number, dtype=np.int32)
                grid_size = coord_norm.max(axis=0) / grid_number
                # avoid division by zero in degenerate dims
                grid_size[grid_size == 0] = 1e-6

                grid_coord = np.floor(coord_norm / grid_size).astype(np.int32)
                # floor() would give integers in [0, grid_number-1] except when coord_norm == max, you get exactly grid_number.
                grid_coord = np.clip(grid_coord, a_min=0, a_max=grid_number - 1)

            data_dict["grid_coord"] = grid_coord

            if self.return_relative:
                grid_center = (grid_coord + 0.5) * grid_size
                # normalized to -1, 1 range
                data_dict["relative_coord"] = (coord_norm - grid_center) * (2.0 / grid_size)

            if not self.sampling:
                return data_dict

            N = len(data_dict["coord"])
            if self.method == "random":
                key = self._ravel_3d(grid_coord)
                idx_sort = np.argsort(key)  # [min_key_index, ... max_key_index]
                key_sort = key[idx_sort]
                unique_keys, inverse, count = np.unique(key_sort, return_inverse=True, return_counts=True)

                if self.mode != "test":
                    # idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(),
                    #                                                                          count.size) % count

                    # TRAIN: pick 1 random point per voxel
                    offsets = np.cumsum(np.insert(count, 0, 0))[:-1]
                    # one random offset in [0, count_i) per voxel
                    rand_offsets = np.random.randint(0, count.max(), size=count.size) % count
                    idx_select = offsets + rand_offsets  # indices into sorted list

                    idx_unique = idx_sort[idx_select]
                    for key, value in data_dict.items():
                        if isinstance(value, (np.ndarray, Sequence)) and len(value) == N and "origin" not in key:
                            data_dict[key] = value[idx_unique]
                else:
                    data_part_list = []
                    max_count = int(count.max())
                    offsets = np.cumsum(np.insert(count, 0, 0))[:-1]

                    for i in range(max_count):
                        idx_select = offsets + (i % count)
                        idx_unique = idx_sort[idx_select]

                        data_part = dict(origin_index=idx_unique)
                        for key, value in data_dict.items():
                            if isinstance(value, (np.ndarray, Sequence)) and len(value) == N and "origin" not in key:
                                data_part[key] = value[idx_unique]
                            else:
                                data_part[key] = data_dict[key]
                        data_part_list.append(data_part)
                    return data_part_list


            elif self.method == "mean":
                unique_keys, inverse, count = np.unique(grid_coord, return_inverse=True, return_counts=True, axis=0)
                unique_sizes = len(unique_keys)

                # ----- combine features along feature dimension -----
                feat_parts = []
                feat_slices = {}   # record where each feature lives in combined vector
                start = 0
                for name in ["coord", "norm", "color"]:
                    if name in data_dict:
                        arr = np.array(data_dict[name])
                        dim = arr.shape[1]
                        feat_parts.append(arr)
                        feat_slices[name] = slice(start, start + dim)
                        start += dim
                if feat_parts:
                    combined_feat = np.concat(feat_parts, axis=1)

                    sum_feat = np.zeros((unique_sizes, combined_feat.shape[1]), dtype=combined_feat.dtype)
                    np.add.at(sum_feat, inverse, combined_feat)
                    avg_feat = sum_feat / count[:, np.newaxis]


                    data_dict["grid_coord"] = unique_keys
                    # ----- write back per feature -----
                    for key in feat_slices.keys():
                        data_dict[key] = avg_feat[:, feat_slices[key]]

                if "label" in data_dict and len(data_dict["label"]) == N:
                    labels = np.array(data_dict["label"])
                    if np.issubdtype(labels.dtype, np.integer):
                        # treat as classification
                        num_classes_tmp = int(labels.max()) + 1
                        hist = np.zeros((unique_sizes, num_classes_tmp), dtype=np.int32)
                        np.add.at(hist, (inverse, labels), 1)   # (inverse, labels) is the 2D index for the hist
                        data_dict["label"] = hist.argmax(axis=1)
                    else:
                        # treat as regression
                        if labels.ndim == 1:
                            labels = labels[:, np.newaxis]
                        sum_label = np.zeros((unique_sizes, labels.shape[1]), dtype=np.float32)
                        np.add.at(sum_label, inverse, labels)
                        avg_label = sum_label / count[:, None]
                        if labels.ndim == 1:
                            avg_label = avg_label[:, 0]
                        data_dict["label"] = avg_label

        return data_dict

`call(data_dict)`

Apply grid voxelization and optional voxel-wise sampling/pooling.

Parameters:

Name	Type	Description	Default
`data_dict`	`dict`	Input dictionary that must contain `"coord"`. Optionally `"norm"`, `"color"`, `"label"`, etc.	required

Returns:

Type	Description
`dict \| list`	dict \| list[dict]: * If `sampling=False` or `method="mean"` or `mode="train"`: returns a single modified `data_dict`. * If `method="random"` and `mode="test"`: returns a list of dictionaries, each representing a different per-voxel sampling pass.

Source code in src\augmentation_class.py

def __call__(self, data_dict: dict) -> dict | list:
    """Apply grid voxelization and optional voxel-wise sampling/pooling.


    Args:
        data_dict (dict): Input dictionary that must contain `"coord"`.
            Optionally `"norm"`, `"color"`, `"label"`, etc.

    Returns:
        dict | list[dict]:
            * If `sampling=False` or `method="mean"` or `mode="train"`:
              returns a single modified `data_dict`.
            * If `method="random"` and `mode="test"`: returns a list of dictionaries, each representing a
              different per-voxel sampling pass.
    """
    if "coord" in data_dict.keys():
        coord = data_dict["coord"]
        coord_min = coord.min(axis=0)
        coord_norm = coord - coord_min
        # kill tiny negative noise due to FP
        coord_norm = np.clip(coord_norm, 0.0, None)  # everything < 0 → 0.0

        if self.grid_number is None:
            # fixed grid cell size
            grid_size = self.grid_size
            grid_coord = np.floor(coord_norm / grid_size).astype(int)
        else:
            # fixed number of cells per axis
            grid_number = np.array(self.grid_number, dtype=np.int32)
            grid_size = coord_norm.max(axis=0) / grid_number
            # avoid division by zero in degenerate dims
            grid_size[grid_size == 0] = 1e-6

            grid_coord = np.floor(coord_norm / grid_size).astype(np.int32)
            # floor() would give integers in [0, grid_number-1] except when coord_norm == max, you get exactly grid_number.
            grid_coord = np.clip(grid_coord, a_min=0, a_max=grid_number - 1)

        data_dict["grid_coord"] = grid_coord

        if self.return_relative:
            grid_center = (grid_coord + 0.5) * grid_size
            # normalized to -1, 1 range
            data_dict["relative_coord"] = (coord_norm - grid_center) * (2.0 / grid_size)

        if not self.sampling:
            return data_dict

        N = len(data_dict["coord"])
        if self.method == "random":
            key = self._ravel_3d(grid_coord)
            idx_sort = np.argsort(key)  # [min_key_index, ... max_key_index]
            key_sort = key[idx_sort]
            unique_keys, inverse, count = np.unique(key_sort, return_inverse=True, return_counts=True)

            if self.mode != "test":
                # idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(),
                #                                                                          count.size) % count

                # TRAIN: pick 1 random point per voxel
                offsets = np.cumsum(np.insert(count, 0, 0))[:-1]
                # one random offset in [0, count_i) per voxel
                rand_offsets = np.random.randint(0, count.max(), size=count.size) % count
                idx_select = offsets + rand_offsets  # indices into sorted list

                idx_unique = idx_sort[idx_select]
                for key, value in data_dict.items():
                    if isinstance(value, (np.ndarray, Sequence)) and len(value) == N and "origin" not in key:
                        data_dict[key] = value[idx_unique]
            else:
                data_part_list = []
                max_count = int(count.max())
                offsets = np.cumsum(np.insert(count, 0, 0))[:-1]

                for i in range(max_count):
                    idx_select = offsets + (i % count)
                    idx_unique = idx_sort[idx_select]

                    data_part = dict(origin_index=idx_unique)
                    for key, value in data_dict.items():
                        if isinstance(value, (np.ndarray, Sequence)) and len(value) == N and "origin" not in key:
                            data_part[key] = value[idx_unique]
                        else:
                            data_part[key] = data_dict[key]
                    data_part_list.append(data_part)
                return data_part_list


        elif self.method == "mean":
            unique_keys, inverse, count = np.unique(grid_coord, return_inverse=True, return_counts=True, axis=0)
            unique_sizes = len(unique_keys)

            # ----- combine features along feature dimension -----
            feat_parts = []
            feat_slices = {}   # record where each feature lives in combined vector
            start = 0
            for name in ["coord", "norm", "color"]:
                if name in data_dict:
                    arr = np.array(data_dict[name])
                    dim = arr.shape[1]
                    feat_parts.append(arr)
                    feat_slices[name] = slice(start, start + dim)
                    start += dim
            if feat_parts:
                combined_feat = np.concat(feat_parts, axis=1)

                sum_feat = np.zeros((unique_sizes, combined_feat.shape[1]), dtype=combined_feat.dtype)
                np.add.at(sum_feat, inverse, combined_feat)
                avg_feat = sum_feat / count[:, np.newaxis]


                data_dict["grid_coord"] = unique_keys
                # ----- write back per feature -----
                for key in feat_slices.keys():
                    data_dict[key] = avg_feat[:, feat_slices[key]]

            if "label" in data_dict and len(data_dict["label"]) == N:
                labels = np.array(data_dict["label"])
                if np.issubdtype(labels.dtype, np.integer):
                    # treat as classification
                    num_classes_tmp = int(labels.max()) + 1
                    hist = np.zeros((unique_sizes, num_classes_tmp), dtype=np.int32)
                    np.add.at(hist, (inverse, labels), 1)   # (inverse, labels) is the 2D index for the hist
                    data_dict["label"] = hist.argmax(axis=1)
                else:
                    # treat as regression
                    if labels.ndim == 1:
                        labels = labels[:, np.newaxis]
                    sum_label = np.zeros((unique_sizes, labels.shape[1]), dtype=np.float32)
                    np.add.at(sum_label, inverse, labels)
                    avg_label = sum_label / count[:, None]
                    if labels.ndim == 1:
                        avg_label = avg_label[:, 0]
                    data_dict["label"] = avg_label

    return data_dict