Spaces:
Running on Zero
Running on Zero
| """ | |
| General utility functions for dataset processing. | |
| This module contains common functions that can be used across different dataset loaders, | |
| such as point cloud processing, normal estimation, and downsampling. | |
| """ | |
| import numpy as np | |
| import open3d as o3d | |
| from typing import Optional, Tuple, Union, List | |
| import logging | |
| import torch | |
| import os | |
| # Import necessary modules from other utils files | |
| from .io_utils import get_dataset_name | |
| logger = logging.getLogger(__name__) | |
| def random_downsample_points(points: np.ndarray, max_points: int, seed: Optional[int] = None) -> np.ndarray: | |
| """ | |
| Randomly downsample a point cloud to a maximum number of points. | |
| Args: | |
| points: Input point cloud (N, 3) or (N, 4) with intensity | |
| max_points: Maximum number of points to keep | |
| seed: Random seed for reproducibility (optional) | |
| Returns: | |
| Downsampled point cloud | |
| """ | |
| if len(points) <= max_points: | |
| return points | |
| if seed is not None: | |
| np.random.seed(seed) | |
| indices = np.random.choice(len(points), max_points, replace=False) | |
| downsampled_points = points[indices] | |
| # logger.debug(f"Randomly downsampled {len(points)} to {len(downsampled_points)} points") | |
| return downsampled_points | |
| def farthest_point_sampling(points: np.ndarray, num_samples: int, seed: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Perform farthest point sampling (FPS) on a point cloud using PyTorch3D when available. | |
| Note: this is still super slow when the input point cloud is too large. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| num_samples: Number of points to sample | |
| seed: Random seed for reproducibility (optional) | |
| Returns: | |
| Tuple of (sampled_points, sampled_indices) | |
| - sampled_points: The sampled points (num_samples, 3) | |
| - sampled_indices: The indices of sampled points in original array (num_samples,) | |
| """ | |
| if len(points) == 0: | |
| return np.array([]), np.array([], dtype=int) | |
| if num_samples >= len(points): | |
| return points.copy(), np.arange(len(points)) | |
| if seed is not None: | |
| np.random.seed(seed) | |
| if PYTORCH3D_AVAILABLE: | |
| torch.manual_seed(seed) | |
| if PYTORCH3D_AVAILABLE: | |
| # Use PyTorch3D's efficient implementation | |
| try: | |
| # Determine device (use GPU if available for large point clouds) | |
| device = torch.device('cuda' if torch.cuda.is_available() and len(points) > 10000 else 'cpu') | |
| # Convert to torch tensor | |
| points_tensor = torch.from_numpy(points).float().unsqueeze(0).to(device) # Add batch dimension | |
| # Sample farthest points | |
| sampled_points_tensor, sampled_indices_tensor = sample_farthest_points( | |
| points_tensor, K=num_samples, random_start_point=True | |
| ) | |
| # Convert back to numpy | |
| sampled_points = sampled_points_tensor.squeeze(0).cpu().numpy() # Remove batch dimension and move to CPU | |
| sampled_indices = sampled_indices_tensor.squeeze(0).cpu().numpy().astype(int) | |
| device_str = f" on {device}" if device.type == 'cuda' else "" | |
| logger.info(f"FPS (PyTorch3D{device_str}) downsampled {len(points)} to {len(sampled_points)} points") | |
| return sampled_points, sampled_indices | |
| except Exception as e: | |
| logger.warning(f"PyTorch3D FPS failed ({e}), falling back to NumPy implementation") | |
| # Fall through to NumPy implementation | |
| # Fallback NumPy implementation | |
| logger.debug("Using NumPy FPS implementation (PyTorch3D not available)") | |
| num_points = len(points) | |
| sampled_indices = [] | |
| distances = np.full(num_points, np.inf) | |
| # Start with a random point | |
| current_idx = np.random.randint(0, num_points) | |
| sampled_indices.append(current_idx) | |
| for _ in range(1, num_samples): | |
| # Update distances to the current point | |
| current_point = points[current_idx] | |
| current_distances = np.linalg.norm(points - current_point, axis=1) | |
| distances = np.minimum(distances, current_distances) | |
| # Select the farthest point | |
| current_idx = np.argmax(distances) | |
| sampled_indices.append(current_idx) | |
| sampled_indices = np.array(sampled_indices) | |
| sampled_points = points[sampled_indices] | |
| logger.debug(f"FPS (NumPy) downsampled {len(points)} to {len(sampled_points)} points") | |
| return sampled_points, sampled_indices | |
| def farthest_point_sampling_with_normals(points: np.ndarray, | |
| normals: Optional[np.ndarray] = None, | |
| num_samples: int = 1000, | |
| seed: Optional[int] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]: | |
| """ | |
| Perform farthest point sampling on a point cloud with optional normals. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| normals: Optional normal vectors (N, 3) | |
| num_samples: Number of points to sample | |
| seed: Random seed for reproducibility (optional) | |
| Returns: | |
| Tuple of (sampled_points, sampled_normals) | |
| - sampled_points: The sampled points (num_samples, 3) | |
| - sampled_normals: The sampled normals (num_samples, 3) or None if normals is None | |
| """ | |
| sampled_points, sampled_indices = farthest_point_sampling(points, num_samples, seed) | |
| if normals is not None and len(normals) > 0: | |
| sampled_normals = normals[sampled_indices] | |
| return sampled_points, sampled_normals | |
| else: | |
| return sampled_points, None | |
| def downsample_points(points: np.ndarray, | |
| normals: Optional[np.ndarray] = None, | |
| method: str = "voxel", | |
| voxel_size: float = 0.1, | |
| num_points: Optional[int] = None, | |
| seed: Optional[int] = None, | |
| use_torch: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray]]: | |
| """ | |
| Unified function for downsampling point clouds using different methods. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| normals: Optional normal vectors (N, 3) | |
| method: Downsampling method ("voxel", "fps", or "random") | |
| voxel_size: Size of voxels for voxel downsampling | |
| num_points: Target number of points (required for FPS and random methods) | |
| seed: Random seed for reproducibility (optional) | |
| use_torch: If True and method is "voxel", use torch-based voxel downsampling for speedup (default: False) | |
| Returns: | |
| Tuple of (downsampled_points, downsampled_normals) | |
| - downsampled_points: The downsampled points | |
| - downsampled_normals: The downsampled normals (or None if normals is None) | |
| """ | |
| if len(points) == 0: | |
| return points, normals | |
| if method == "voxel": | |
| if use_torch: | |
| # Use torch-based voxel downsampling for speedup | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| points_tensor = torch.from_numpy(points).float().to(device) | |
| # Get downsampled indices using torch | |
| downsampled_indices_tensor = voxel_down_sample_torch(points_tensor, voxel_size) | |
| downsampled_indices = downsampled_indices_tensor.cpu().numpy() | |
| # Apply indices to points and normals | |
| downsampled_points = points[downsampled_indices] | |
| downsampled_normals = normals[downsampled_indices] if normals is not None else None | |
| else: | |
| # Use numpy-based voxel downsampling (original implementation) | |
| downsampled_indices = _get_voxel_downsample_indices(points, voxel_size) | |
| downsampled_points = points[downsampled_indices] | |
| downsampled_normals = normals[downsampled_indices] if normals is not None else None | |
| elif method == "fps": | |
| if num_points is None: | |
| raise ValueError("num_points must be specified for FPS downsampling") | |
| downsampled_points, downsampled_normals = farthest_point_sampling_with_normals( | |
| points, normals, num_points, seed | |
| ) | |
| elif method == "random": | |
| if num_points is None: | |
| raise ValueError("num_points must be specified for random downsampling") | |
| if len(points) <= num_points: | |
| downsampled_points = points.copy() | |
| downsampled_normals = normals.copy() if normals is not None else None | |
| else: | |
| if seed is not None: | |
| np.random.seed(seed) | |
| indices = np.random.choice(len(points), num_points, replace=False) | |
| downsampled_points = points[indices] | |
| downsampled_normals = normals[indices] if normals is not None else None | |
| else: | |
| raise ValueError(f"Unknown downsampling method: {method}. Choose from 'voxel', 'fps', or 'random'") | |
| method_str = f"{method} method" | |
| if method == "voxel" and use_torch: | |
| method_str = f"{method} method (torch-based)" | |
| logger.debug(f"Downsampled {len(points)} to {len(downsampled_points)} points using {method_str}") | |
| return downsampled_points, downsampled_normals | |
| def _get_voxel_downsample_indices(points: np.ndarray, voxel_size: float) -> np.ndarray: | |
| """ | |
| Get indices of points after voxel downsampling. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| voxel_size: Size of voxels for downsampling | |
| Returns: | |
| Indices of downsampled points | |
| """ | |
| if len(points) == 0: | |
| return np.array([], dtype=int) | |
| # Convert to voxel coordinates | |
| voxel_coords = np.floor(points / voxel_size).astype(int) | |
| # Create unique voxel keys and track which points belong to each voxel | |
| voxel_to_points = {} | |
| for i, voxel_coord in enumerate(voxel_coords): | |
| voxel_key = tuple(voxel_coord) | |
| if voxel_key not in voxel_to_points: | |
| voxel_to_points[voxel_key] = [] | |
| voxel_to_points[voxel_key].append(i) | |
| # For each voxel, find the point closest to the voxel center | |
| downsampled_indices = [] | |
| for voxel_key, point_indices in voxel_to_points.items(): | |
| if len(point_indices) == 1: | |
| # Only one point in this voxel, keep it | |
| downsampled_indices.append(point_indices[0]) | |
| else: | |
| # Multiple points in this voxel, find the one closest to voxel center | |
| voxel_center = (np.array(voxel_key) + 0.5) * voxel_size | |
| voxel_points = points[point_indices] | |
| # Calculate distances to voxel center | |
| distances = np.linalg.norm(voxel_points - voxel_center, axis=1) | |
| # Find the point with minimum distance | |
| closest_idx = np.argmin(distances) | |
| downsampled_indices.append(point_indices[closest_idx]) | |
| return np.array(downsampled_indices) | |
| def voxel_down_sample_torch(points: torch.tensor, voxel_size: float): | |
| """ | |
| voxel based downsampling. Returns the indices of the points which are closest to the voxel centers. | |
| Args: | |
| points (torch.Tensor): [N,3] point coordinates | |
| voxel_size (float): grid resolution | |
| Returns: | |
| indices (torch.Tensor): [M] indices of the original point cloud, downsampled point cloud would be `points[indices]` | |
| Reference: Louis Wiesmann | |
| """ | |
| _quantization = 1000 # if change to 1, then it would take the first (smallest) index lie in the voxel | |
| offset = torch.floor(points.min(dim=0)[0] / voxel_size).long() | |
| grid = torch.floor(points / voxel_size) | |
| center = (grid + 0.5) * voxel_size | |
| dist = ((points - center) ** 2).sum(dim=1) ** 0.5 | |
| dist = ( | |
| dist / dist.max() * (_quantization - 1) | |
| ).long() # for speed up # [0-_quantization] | |
| grid = grid.long() - offset | |
| v_size = grid.max().ceil() | |
| grid_idx = grid[:, 0] + grid[:, 1] * v_size + grid[:, 2] * v_size * v_size | |
| unique, inverse = torch.unique(grid_idx, return_inverse=True) | |
| idx_d = torch.arange(inverse.size(0), dtype=inverse.dtype, device=inverse.device) | |
| offset = 10 ** len(str(idx_d.max().item())) | |
| idx_d = idx_d + dist.long() * offset | |
| idx = torch.empty( | |
| unique.shape, dtype=inverse.dtype, device=inverse.device | |
| ).scatter_reduce_( | |
| dim=0, index=inverse, src=idx_d, reduce="amin", include_self=False | |
| ) | |
| # https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_reduce_.html | |
| # This operation may behave nondeterministically when given tensors on | |
| # a CUDA device. consider to change a more stable implementation | |
| idx = idx % offset | |
| return idx | |
| def estimate_point_normals_open3d(points: np.ndarray, k_neighbors: int = 20, radius: float = 0.5) -> np.ndarray: | |
| """ | |
| Estimate point-wise normals using Open3D. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| k_neighbors: Number of neighbors for normal estimation (default: 20) | |
| radius: Search radius for normal estimation (default: 0.5m) | |
| Returns: | |
| Normal vectors (N, 3) with unit length | |
| """ | |
| if len(points) == 0: | |
| return np.array([]) | |
| # Create Open3D point cloud | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points) | |
| # Estimate normals | |
| # Use both k-neighbors and radius for robust estimation | |
| pcd.estimate_normals( | |
| search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=k_neighbors) | |
| ) | |
| # Orient normals consistently (optional but recommended) | |
| pcd.orient_normals_consistent_tangent_plane(k=k_neighbors) | |
| # Get normals as numpy array | |
| normals = np.asarray(pcd.normals) | |
| logger.debug(f"Estimated normals for {len(points)} points") | |
| return normals | |
| def transform_points(points: np.ndarray, transform: np.ndarray) -> np.ndarray: | |
| """ | |
| Transform points using 4x4 transformation matrix. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| transform: 4x4 transformation matrix | |
| Returns: | |
| Transformed point cloud (N, 3) | |
| """ | |
| # Convert to homogeneous coordinates | |
| points_homo = np.hstack([points, np.ones((points.shape[0], 1))]) | |
| # Apply transformation | |
| transformed_points = (transform @ points_homo.T).T | |
| return transformed_points[:, :3] | |
| def transform_normals(normals: np.ndarray, transform: np.ndarray) -> np.ndarray: | |
| """ | |
| Transform normals using 4x4 transformation matrix. | |
| Normals are transformed using the rotation part of the transformation matrix. | |
| Args: | |
| normals: Normal vectors (N, 3) | |
| transform: 4x4 transformation matrix | |
| Returns: | |
| Transformed normal vectors (N, 3) | |
| """ | |
| if len(normals) == 0: | |
| return normals | |
| # Extract rotation matrix (3x3 upper-left part) | |
| rotation_matrix = transform[:3, :3] | |
| # Transform normals using rotation matrix only | |
| # Note: We don't add homogeneous coordinate for normals | |
| transformed_normals = (rotation_matrix @ normals.T).T | |
| # Normalize the transformed normals | |
| norms = np.linalg.norm(transformed_normals, axis=1, keepdims=True) | |
| # Avoid division by zero | |
| norms = np.where(norms < 1e-8, 1.0, norms) | |
| transformed_normals = transformed_normals / norms | |
| return transformed_normals | |
| def voxel_downsample_points(points: np.ndarray, voxel_size: float = 0.1) -> np.ndarray: | |
| """ | |
| Downsample points using voxel grid, keeping the actual points closest to voxel centers. | |
| Args: | |
| points: Input point cloud (N, 3) | |
| voxel_size: Size of voxels for downsampling | |
| Returns: | |
| Downsampled point cloud with actual points (not voxel centers) | |
| Note: | |
| This function is kept for backward compatibility. Consider using downsample_points() for more options. | |
| """ | |
| downsampled_points, _ = downsample_points(points, method="voxel", voxel_size=voxel_size) | |
| logger.debug(f"Voxel downsampled {len(points)} to {len(downsampled_points)} points") | |
| return downsampled_points | |
| def voxel_downsample_points_open3d(points: np.ndarray, voxel_size: float = 0.1) -> np.ndarray: | |
| """ | |
| Downsample points using Open3D's voxel downsampling (returns voxel centers). | |
| Args: | |
| points: Input point cloud (N, 3) | |
| voxel_size: Size of voxels for downsampling | |
| Returns: | |
| Downsampled point cloud with voxel center points | |
| """ | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points) | |
| # Apply voxel downsampling (returns voxel centers) | |
| downsampled_pcd = pcd.voxel_down_sample(voxel_size=voxel_size) | |
| return np.asarray(downsampled_pcd.points) | |
| def get_pose_center(pose: np.ndarray) -> np.ndarray: | |
| """ | |
| Extract the center position from a 4x4 pose matrix. | |
| Args: | |
| pose: 4x4 pose matrix | |
| Returns: | |
| Center position (3,) as numpy array | |
| """ | |
| return pose[:3, 3] | |
| def filter_keyframes_by_motion(poses: List[np.ndarray], | |
| translation_threshold: float = 0.5, | |
| rotation_threshold_degrees: float = 5.0, | |
| min_frame_interval: int = 1) -> List[int]: | |
| """ | |
| Filter poses to extract keyframes based on translation and rotation thresholds. | |
| This function identifies keyframes by comparing the motion between consecutive poses. | |
| A frame is selected as a keyframe if either: | |
| 1. The translation distance from the last keyframe exceeds translation_threshold, OR | |
| 2. The rotation angle from the last keyframe exceeds rotation_threshold_degrees | |
| Args: | |
| poses: List of 4x4 pose matrices | |
| translation_threshold: Minimum translation distance in meters to trigger keyframe selection (default: 0.5m) | |
| rotation_threshold_degrees: Minimum rotation angle in degrees to trigger keyframe selection (default: 5.0°) | |
| min_frame_interval: Minimum number of frames between keyframes (default: 1, meaning consecutive frames allowed) | |
| Returns: | |
| List of keyframe indices (0-based) | |
| """ | |
| if poses is None or len(poses) == 0: | |
| return [] | |
| if len(poses) == 1: | |
| return [0] | |
| # Convert rotation threshold from degrees to radians | |
| rotation_threshold_rad = np.radians(rotation_threshold_degrees) | |
| keyframe_indices = [0] # Always include the first frame | |
| last_keyframe_pose = poses[0] | |
| last_keyframe_idx = 0 | |
| for i in range(1, len(poses)): | |
| current_pose = poses[i] | |
| # Check minimum frame interval | |
| if i - last_keyframe_idx < min_frame_interval: | |
| continue | |
| # Calculate translation distance | |
| translation_distance = np.linalg.norm( | |
| get_pose_center(current_pose) - get_pose_center(last_keyframe_pose) | |
| ) | |
| # Calculate rotation angle difference | |
| rotation_angle_rad = calculate_rotation_angle_between_poses(current_pose, last_keyframe_pose) | |
| # Check if either threshold is exceeded | |
| if (translation_distance >= translation_threshold or | |
| rotation_angle_rad >= rotation_threshold_rad): | |
| keyframe_indices.append(i) | |
| last_keyframe_pose = current_pose | |
| last_keyframe_idx = i | |
| # Always include the last frame if it's not already included | |
| if keyframe_indices[-1] != len(poses) - 1: | |
| keyframe_indices.append(len(poses) - 1) | |
| logger.info(f"Selected {len(keyframe_indices)} keyframes from {len(poses)} total frames " | |
| f"(translation_threshold: {translation_threshold}m, " | |
| f"rotation_threshold: {rotation_threshold_degrees:.1f}°)") | |
| return keyframe_indices | |
| def calculate_rotation_angle_between_poses(pose1: np.ndarray, pose2: np.ndarray) -> float: | |
| """ | |
| Calculate the rotation angle between two 4x4 pose matrices. | |
| Args: | |
| pose1: First 4x4 pose matrix | |
| pose2: Second 4x4 pose matrix | |
| Returns: | |
| Rotation angle in radians (0 to π) | |
| """ | |
| # Extract rotation matrices | |
| R1 = pose1[:3, :3] | |
| R2 = pose2[:3, :3] | |
| # Calculate relative rotation | |
| R_rel = R1.T @ R2 | |
| # Calculate rotation angle using trace formula | |
| # For rotation matrix R, the angle θ satisfies: trace(R) = 1 + 2*cos(θ) | |
| trace = np.trace(R_rel) | |
| # Clamp trace to valid range [-1, 3] to handle numerical errors | |
| trace = np.clip(trace, -1.0, 3.0) | |
| # Calculate angle: cos(θ) = (trace - 1) / 2 | |
| cos_angle = (trace - 1.0) / 2.0 | |
| cos_angle = np.clip(cos_angle, -1.0, 1.0) # Ensure valid range for arccos | |
| angle = np.arccos(cos_angle) | |
| return angle | |
| def filter_poses_and_data_by_keyframes(poses: List[np.ndarray], | |
| data_list: Optional[List] = None, | |
| translation_threshold: float = 0.5, | |
| rotation_threshold_degrees: float = 5.0, | |
| min_frame_interval: int = 1) -> Tuple[List[np.ndarray], List[int], Optional[List]]: | |
| """ | |
| Filter poses and associated data by keyframes based on motion thresholds. | |
| This is a convenience function that combines keyframe filtering with data filtering. | |
| Args: | |
| poses: List of 4x4 pose matrices | |
| data_list: Optional list of data (e.g., point clouds, timestamps) to filter alongside poses | |
| translation_threshold: Minimum translation distance in meters to trigger keyframe selection | |
| rotation_threshold_degrees: Minimum rotation angle in degrees to trigger keyframe selection | |
| min_frame_interval: Minimum number of frames between keyframes | |
| Returns: | |
| Tuple of (filtered_poses, keyframe_indices, filtered_data_list) | |
| - filtered_poses: List of keyframe poses | |
| - keyframe_indices: List of original indices of selected keyframes | |
| - filtered_data_list: Filtered data list (None if data_list was None) | |
| """ | |
| keyframe_indices = filter_keyframes_by_motion( | |
| poses, translation_threshold, rotation_threshold_degrees, min_frame_interval | |
| ) | |
| # Filter poses | |
| filtered_poses = [poses[i] for i in keyframe_indices] | |
| # Filter data if provided | |
| filtered_data_list = None | |
| if data_list is not None: | |
| if len(data_list) != len(poses): | |
| logger.warning(f"Data list length ({len(data_list)}) doesn't match poses length ({len(poses)})") | |
| else: | |
| filtered_data_list = [data_list[i] for i in keyframe_indices] | |
| return filtered_poses, keyframe_indices, filtered_data_list | |
| def calculate_point_cloud_overlap_ratio_fast(points1: np.ndarray, | |
| points2: np.ndarray, | |
| voxel_size: float = 2.0, | |
| max_points_per_cloud: int = 20000) -> float: | |
| """ | |
| Fast approximate overlap calculation using downsampling and voxelization. | |
| Args: | |
| points1: First point cloud (N, 3) | |
| points2: Second point cloud (M, 3) | |
| voxel_size: Size of voxels for discretization | |
| max_points_per_cloud: Maximum points to use per cloud for speed | |
| Returns: | |
| Overlap ratio (0.0 to 1.0) | |
| """ | |
| if len(points1) == 0 or len(points2) == 0: | |
| return 0.0 | |
| # Downsample if too many points | |
| if len(points1) > max_points_per_cloud: | |
| indices = np.random.choice(len(points1), max_points_per_cloud, replace=False) | |
| points1 = points1[indices] | |
| if len(points2) > max_points_per_cloud: | |
| indices = np.random.choice(len(points2), max_points_per_cloud, replace=False) | |
| points2 = points2[indices] | |
| # Voxelize both point clouds | |
| def voxelize_points(points): | |
| # Convert to voxel coordinates | |
| voxel_coords = np.floor(points / voxel_size).astype(int) | |
| # Create unique voxel keys | |
| voxel_keys = set(map(tuple, voxel_coords)) | |
| return voxel_keys | |
| voxels1 = voxelize_points(points1) | |
| voxels2 = voxelize_points(points2) | |
| # Calculate intersection and union | |
| intersection = len(voxels1.intersection(voxels2)) | |
| union = len(voxels1.union(voxels2)) | |
| if union == 0: | |
| return 0.0 | |
| overlap_ratio = intersection / union | |
| return overlap_ratio | |
| def sample_truncated_gaussian(min_val: int, max_val: int, mean: float, std: float, max_attempts: int = 100) -> int: | |
| """ | |
| Sample from a truncated Gaussian distribution bounded by min_val and max_val. | |
| Args: | |
| min_val: Minimum allowed value | |
| max_val: Maximum allowed value | |
| mean: Mean of the Gaussian distribution | |
| std: Standard deviation of the Gaussian distribution | |
| max_attempts: Maximum attempts to find a valid sample | |
| Returns: | |
| Sampled integer value within [min_val, max_val] | |
| """ | |
| for _ in range(max_attempts): | |
| # Sample from normal distribution | |
| sample = np.random.normal(mean, std) | |
| # Round to nearest integer | |
| sample_int = int(round(sample)) | |
| # Check if within bounds | |
| if min_val <= sample_int <= max_val: | |
| return sample_int | |
| # If we can't find a valid sample, fall back to uniform distribution | |
| logger.warning(f"Could not sample from truncated Gaussian with mean={mean}, std={std}, bounds=[{min_val}, {max_val}]. Using uniform distribution.") | |
| import random | |
| return random.randint(min_val, max_val) | |
| def deskewing( | |
| points: np.ndarray, | |
| ts: np.ndarray, | |
| pose: np.ndarray, | |
| ts_mid_pose: float = 0.5 | |
| ) -> np.ndarray: | |
| """ | |
| Deskew a batch of points at timestamp ts by a relative transformation matrix. | |
| PyTorch implementation using roma library for faster SLERP. | |
| Args: | |
| points: Point cloud with shape (N, 3) or (N, 4) if including intensity | |
| ts: Timestamps for each point, shape (N,) with values in range [0, 1] | |
| pose: Transformation matrix from current to last frame, shape (4, 4) | |
| ts_mid_pose: Timestamp corresponding to the pose (default: 0.5) | |
| Returns: | |
| Deskewed point cloud with same shape as input points | |
| """ | |
| if ts is None: | |
| return points # no deskewing | |
| try: | |
| import torch | |
| import roma | |
| except ImportError: | |
| raise ImportError("PyTorch and roma are required for fast deskewing. Install with: pip install torch roma") | |
| # Convert inputs to torch tensors | |
| points_torch = torch.from_numpy(points).float() | |
| ts_torch = torch.from_numpy(ts).float() | |
| pose_torch = torch.from_numpy(pose).float() | |
| # Ensure ts is 1D | |
| ts_torch = ts_torch.squeeze(-1) | |
| # Normalize the tensor to the range [0, 1] | |
| # NOTE: you need to figure out the begin and end of a frame because | |
| # sometimes there's only partial measurements, some part are blocked by some occlusions | |
| min_ts = torch.min(ts_torch) | |
| max_ts = torch.max(ts_torch) | |
| # Avoid division by zero | |
| if max_ts - min_ts > 1e-8: | |
| ts_torch = (ts_torch - min_ts) / (max_ts - min_ts) | |
| else: | |
| # If all timestamps are the same, set them to 0.5 | |
| ts_torch = torch.full_like(ts_torch, 0.5) | |
| # this is related to: https://github.com/PRBonn/kiss-icp/issues/299 | |
| ts_torch -= ts_mid_pose | |
| # Use roma for fast SLERP | |
| rotmat_slerp = roma.rotmat_slerp( | |
| torch.eye(3).to(points_torch), pose_torch[:3, :3].to(points_torch), ts_torch | |
| ) | |
| tran_lerp = ts_torch[:, None] * pose_torch[:3, 3].to(points_torch) | |
| points_deskewed_torch = points_torch.clone() | |
| points_deskewed_torch[:, :3] = (rotmat_slerp @ points_torch[:, :3].unsqueeze(-1)).squeeze(-1) + tran_lerp | |
| # Convert back to numpy | |
| points_deskewed = points_deskewed_torch.cpu().numpy() | |
| return points_deskewed | |
| def get_global_transformation_matrix(sequence_name: str) -> np.ndarray: | |
| """Get global transformation matrix based on sequence name.""" | |
| if sequence_name.startswith('7-scenes') or \ | |
| sequence_name.startswith('bundlefusion') or \ | |
| sequence_name.startswith('rgbd-scenes') or \ | |
| sequence_name.startswith('sun3d'): # for the test set | |
| # Transformation for 7-scenes, bundlefusion, rgbd-scenes | |
| return np.array([[0, 0, 1], | |
| [-1, 0, 0], | |
| [0, -1, 0]], dtype=np.float32) | |
| # elif sequence_name.startswith('sun3d'): | |
| # # Transformation for sun3d | |
| # return np.array([[0, 0, -1], | |
| # [1, 0, 0], | |
| # [0, 1, 0]], dtype=np.float32) | |
| elif sequence_name.startswith('analysis-by-synthesis'): | |
| # Identity transformation for analysis-by-synthesis | |
| return np.eye(3, dtype=np.float32) | |
| else: | |
| # Default to identity for other sequences | |
| return np.eye(3, dtype=np.float32) | |
| def save_num_points_to_folder(output_dir: str, | |
| dataset_name: str, | |
| sample_num_points: List[int], | |
| sample_dirs: List[Tuple[str, str]]): | |
| """Save num_points data to folder structure following precompute_num_points.py format.""" | |
| num_points_dir = os.path.join(output_dir, "num_points") | |
| os.makedirs(num_points_dir, exist_ok=True) | |
| # For folder structure, we need to map samples to train/val splits | |
| # We'll try to read the data_split files to determine which samples belong to which split | |
| data_split_dir = os.path.join(output_dir, 'data_split') | |
| # Create sample path to num_points mapping | |
| sample_to_num_points = {} | |
| all_num_points = [] | |
| def normalize_path(path: str) -> str: | |
| """Normalize path by removing dataset prefix and _processed suffix for matching.""" | |
| # Remove dataset name prefix if present | |
| if path.startswith(f"{dataset_name}/"): | |
| path = path[len(f"{dataset_name}/"):] | |
| # Remove _processed suffix if present | |
| if path.endswith("_processed"): | |
| path = path[:-len("_processed")] | |
| return path | |
| for i, (relative_path, _) in enumerate(sample_dirs): | |
| num_points = sample_num_points[i] if i < len(sample_num_points) else 0 | |
| normalized_path = normalize_path(relative_path) | |
| sample_to_num_points[normalized_path] = num_points | |
| all_num_points.append(num_points) | |
| # Create split-specific num_points lists | |
| split_num_points = {} | |
| # Read all split files and create corresponding num_points lists | |
| if os.path.exists(data_split_dir): | |
| for split_file_name in os.listdir(data_split_dir): | |
| if split_file_name.endswith('.txt'): | |
| split_name = split_file_name.replace('.txt', '') | |
| split_file_path = os.path.join(data_split_dir, split_file_name) | |
| with open(split_file_path, 'r') as f: | |
| samples = [line.strip() for line in f if line.strip()] | |
| # Create num_points list for this split | |
| split_points = [] | |
| for sample in samples: | |
| normalized_sample = normalize_path(sample) | |
| if normalized_sample in sample_to_num_points: | |
| split_points.append(sample_to_num_points[normalized_sample]) | |
| else: | |
| logger.warning(f"Sample {sample} (normalized: {normalized_sample}) not found in processed data for {split_name} split") | |
| split_points.append(0) | |
| if split_points: | |
| split_num_points[split_name] = split_points | |
| # Always include 'all' split | |
| split_num_points['all'] = all_num_points | |
| # Save all split files | |
| for split_name, num_points_data in split_num_points.items(): | |
| if num_points_data: | |
| split_file = os.path.join(num_points_dir, f'{split_name}.txt') | |
| with open(split_file, 'w') as f: | |
| for num_point in num_points_data: | |
| f.write(f"{num_point}\n") | |
| logger.info(f"Saved num_points for {split_name} split: {split_file} ({len(num_points_data)} samples)") | |
| logger.debug(f"{split_name} num_points stats: min={min(num_points_data)}, max={max(num_points_data)}, mean={np.mean(num_points_data):.1f}") | |
| def save_points_to_ply(points: np.ndarray, ply_path: str, normals: Optional[np.ndarray] = None) -> None: | |
| """Save points to a PLY file using Open3D.""" | |
| if o3d is None: | |
| raise ImportError("open3d is required. Install with: pip install open3d") | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64, copy=False)) | |
| if normals is not None: | |
| pcd.normals = o3d.utility.Vector3dVector(normals.astype(np.float64, copy=False)) | |
| os.makedirs(os.path.dirname(ply_path), exist_ok=True) | |
| o3d.io.write_point_cloud(ply_path, pcd, write_ascii=False) | |