utils.py

import numpy as np
from scipy.spatial import KDTree
import matplotlib.pyplot as plt
import open3d as o3d
from sklearn.neighbors import NearestNeighbors


score_thresh = 0.85
def keep_score(score):
    return score > score_thresh

# def keep_dist(new_pts,existing_pts):
#     median_existing = np.median(existing_pts, axis=1)

#     distances = np.linalg.norm(new_pts - median_existing[:, np.newaxis], axis=0)
#     kept_points = new_pts[:, distances <= dist_thresh]

#     return kept_points

def keep_dist_region_growing(new_pts, dist_thresh=0.1, min_density_thresh=500, max_density_thresh=7500):
    """
    Filters noisy points using region growing while preserving edges.
    
    Parameters:
    - new_pts: New points to filter (shape: 3 x N).
    - existing_pts: Existing points to compare against (shape: 3 x M).
    - normals: Normal vectors corresponding to new_pts (shape: 3 x N).
    - dist_thresh: Maximum distance for inclusion.
    - density_thresh: Minimum number of neighbors required for inclusion.
    
    Returns:
    - kept_points: Points from new_pts that meet the criteria.
    """
    # Detect edges in the point cloud
    
    # Build a KDTree for efficient neighbor search
    tree = KDTree(new_pts.T)
    
    # Initialize a mask for points to keep
    keep_mask = np.zeros(new_pts.shape[1], dtype=bool)
    
    # Iterate through all points in new_pts
    for i in range(new_pts.shape[1]):
        # Apply region growing logic for non-edge points
        neighbors = tree.query_ball_point(new_pts[:, i], dist_thresh)
        if min_density_thresh <= len(neighbors) <= max_density_thresh:
            keep_mask[i] = True
    
    # Return filtered points
    kept_points = new_pts[:, keep_mask]
    
    return kept_points

def guided_filter(new_pts, radius=0.0001, epsilon=0.1):
    kdtree = o3d.geometry.KDTreeFlann(new_pts)
    points_copy = np.array(new_pts)
    points = np.asarray(new_pts)
    num_points = new_pts.shape[1]

    for i in range(num_points):
        k, idx, _ = kdtree.search_radius_vector_3d(new_pts[:, i], radius)
        if k < 1000:
            continue

        neighbors = points[:, idx]
        mean = np.mean(neighbors.T, 0)
        cov = np.cov(neighbors)
        e = np.linalg.inv(cov + epsilon * np.eye(3))

        A = cov @ e
        b = mean - A @ mean

        points_copy[:, i] = A @ points[:, i] + b

    return points_copy

def keep_dist_median(new_pts, dist_thresh=0.1):
    
    # Compute the median point in each axis
    median_point = np.median(new_pts, axis=1, keepdims=True)

    # Compute Euclidean distances from the median point
    distances = np.linalg.norm(new_pts - median_point, axis=0)

    # Keep points within the threshold distance
    keep_mask = distances <= dist_thresh
    kept_points = new_pts[:, keep_mask]

    return kept_points

def statistical_outlier_removal(points_np, nb_neighbors=20, std_ratio=2.0):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points_np.T)  # shape: (3, N) → (N, 3)

    pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
    filtered_points = np.asarray(pcd.points).T  # shape back to (3, N)
    return filtered_points


def save_pointcloud_as_ply(points, filename):
    """
    Save a point cloud to a .ply file using open3d.

    Args:
        points (np.ndarray): Array of shape (3, N) or (N, 3)
        filename (str): Output filename (.ply)
    """
    if points.shape[0] == 3 and points.shape[1] != 3:
        points = points.T  # Convert from (3, N) to (N, 3)
    
    point_cloud = o3d.geometry.PointCloud()
    point_cloud.points = o3d.utility.Vector3dVector(points)
    o3d.io.write_point_cloud(filename, point_cloud)


def compute_precision_cm(points, min_neighbors=1000, radius_mm=30):
    points = points.T
    if points.shape[1] > 3:
        points = points[:, :3]
    
    print(f"Total points: {len(points)}, Searching for ≥{min_neighbors} neighbors in {radius_mm:.2f}mm radius")
    
    nbrs = NearestNeighbors(radius=radius_mm).fit(points)
    distances, indices = nbrs.radius_neighbors(points, return_distance=True)
    
    # Collect densities for debugging
    neighbor_counts = np.array([len(idx) for idx in indices])
    print(f"Max neighbors in radius: {neighbor_counts.max()}")
    
    valid_dists = []
    for d, idx in zip(distances, indices):
        if len(idx) > min_neighbors:
            valid_dists.append(np.mean(d[1:]))  # exclude self-distance
            
    if not valid_dists:
        print(f"No regions with >{min_neighbors} neighbors found")
        return None
        
    mean_dist_mm = np.mean(valid_dists) #* 1000
    return mean_dist_mm / 10.0  # mm → cm