Source code for gunz_cm.preprocs.points.downsample

"""Downsampling operations for 3D point clouds."""
__author__ = "Yeremia Gunawan Adhisantoso"
__license__ = "Clear BSD"
__email__ = "adhisant@tnt.uni-hannover.de"

from pydantic import validate_call, ConfigDict
import numpy as np
import pandas as pd

from ...utils.validate import require_positive_int

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def downsample_points(
    points: np.ndarray,
    ds_ratio: int,
    def_coor: float = np.nan,
) -> np.ndarray:
    """Downsample the given points by a specified ratio.

    Notes
    -----
    - The function ensures that the `ds_ratio` is greater than 1.
    - Points with all NaN values are ignored during downsampling.
    - The resulting array is filled with `def_coor` for indices without valid points.

    Parameters
    ----------
    points : np.ndarray
        The array of points to be downsampled.
    ds_ratio : int
        The downsampling ratio. Must be greater than 1.
    def_coor : float, optional
        The default coordinate value for indices without valid points, by default np.nan.

    Returns
    -------
    np.ndarray
        The downsampled points array.

    """
    require_positive_int("ds_ratio", ds_ratio, min_value=2)

    num_points = points.shape[0]
    valid_points_mask = ~np.isnan(points).all(axis=1)
    valid_points_ids = np.arange(num_points)[valid_points_mask]

    points_df = pd.DataFrame({
        'ids': valid_points_ids,
        'x': points[valid_points_mask, 0],
        'y': points[valid_points_mask, 1],
        'z': points[valid_points_mask, 2],
    })
    points_df['lr_ids'] = points_df['ids'] // ds_ratio

    lr_points_df = (
        points_df[['lr_ids', 'x', 'y', 'z']]
        .groupby('lr_ids')
        .mean()
    )
    num_lr_points = int(np.ceil(num_points/ds_ratio))
    valid_lr_ids = np.array(lr_points_df.index)
    lr_points = np.full((num_lr_points, 3), def_coor)
    lr_points[valid_lr_ids] = lr_points_df.to_numpy()

    return lr_points