preprocessing.py

"""
preprocessing.py - Point-cloud loading, denoising, and downsampling.
"""

import numpy as np
import open3d as o3d


def load_point_cloud(filepath: str) -> np.ndarray:
    """Load a .ply or .pcd file and return an (N, 3) XYZ float array.

    Open3D handles both formats transparently.  Only XYZ coordinates are
    returned; colour/normal information is discarded because the segmentation
    pipeline relies purely on geometry.

    Args:
        filepath: Absolute or relative path to the point-cloud file.

    Returns:
        Float64 array of shape (N, 3) containing X, Y, Z columns.

    Raises:
        RuntimeError: When Open3D fails to read the file or the cloud is empty.
    """
    pcd = o3d.io.read_point_cloud(filepath)
    if not pcd.has_points():
        raise RuntimeError(f"File loaded but contains no points: {filepath}")

    points = np.asarray(pcd.points, dtype=np.float64)
    if points.shape[0] == 0:
        raise RuntimeError(f"Point cloud is empty after loading: {filepath}")

    print(f"  Loaded {points.shape[0]:,} points from '{filepath}'")
    return points


def remove_statistical_outliers(
    points: np.ndarray,
    nb_neighbors: int = 20,
    std_ratio: float = 2.0,
) -> np.ndarray:
    """Remove noise using statistical outlier removal (SOR).

    Each point's mean distance to its *nb_neighbors* nearest neighbours is
    computed.  Points whose mean distance exceeds the global mean by more than
    *std_ratio* standard deviations are classified as outliers and dropped.

    Args:
        points:       Input array of shape (N, 3).
        nb_neighbors: Number of neighbours used to estimate local density.
        std_ratio:    Threshold multiplier; lower values are more aggressive.

    Returns:
        Cleaned array of shape (M, 3) where M ≤ N.
    """
    if points.shape[0] == 0:
        print("  Warning: received empty point cloud – skipping outlier removal.")
        return points

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)

    cl, ind = pcd.remove_statistical_outlier(
        nb_neighbors=nb_neighbors,
        std_ratio=std_ratio,
    )
    cleaned = np.asarray(cl.points, dtype=np.float64)
    removed = points.shape[0] - cleaned.shape[0]
    print(
        f"  Statistical outlier removal: {removed:,} points removed "
        f"({cleaned.shape[0]:,} remaining)"
    )
    return cleaned


def voxel_downsample(points: np.ndarray, voxel_size: float = 0.05) -> np.ndarray:
    """Reduce point density by averaging all points within each voxel cell.

    Downsampling speeds up DBSCAN clustering while preserving the overall
    geometry.  A voxel_size of 0.05 m (5 cm) works well for room-scale scans.

    Args:
        points:     Input array of shape (N, 3).
        voxel_size: Edge length of each voxel cube in metres.

    Returns:
        Downsampled array of shape (M, 3) where M ≤ N.
    """
    if points.shape[0] == 0:
        print("  Warning: received empty point cloud – skipping voxel downsampling.")
        return points

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)

    ds_pcd = pcd.voxel_down_sample(voxel_size=voxel_size)
    downsampled = np.asarray(ds_pcd.points, dtype=np.float64)
    ratio = 100.0 * (1.0 - downsampled.shape[0] / points.shape[0])
    print(
        f"  Voxel downsampling (size={voxel_size}m): "
        f"{points.shape[0]:,} → {downsampled.shape[0]:,} points "
        f"({ratio:.1f}% reduction)"
    )
    return downsampled


def preprocess(
    filepath: str,
    voxel_size: float = 0.05,
    nb_neighbors: int = 20,
    std_ratio: float = 2.0,
) -> np.ndarray:
    """Convenience wrapper: load → denoise → downsample.

    Args:
        filepath:     Path to input .ply or .pcd file.
        voxel_size:   Voxel edge length for downsampling (metres).
        nb_neighbors: SOR neighbourhood size.
        std_ratio:    SOR standard deviation multiplier.

    Returns:
        Preprocessed point array of shape (M, 3).
    """
    points = load_point_cloud(filepath)
    points = remove_statistical_outliers(points, nb_neighbors=nb_neighbors, std_ratio=std_ratio)
    points = voxel_downsample(points, voxel_size=voxel_size)
    return points