Spaces:
Running on Zero
Running on Zero
| import os | |
| import numpy as np | |
| import logging | |
| from typing import List, Tuple, Optional, Union | |
| import torch # For torch.manual_seed in _allocate_fps_points | |
| from pytorch3d.ops import sample_farthest_points | |
| import open3d as o3d | |
| logger = logging.getLogger(__name__) | |
| def calculate_voxel_coverage(points: np.ndarray, voxel_size: float) -> int: | |
| """ | |
| Calculate the number of unique voxels occupied by the point cloud. | |
| Args: | |
| points: Point cloud array (N, 3) | |
| voxel_size: Voxel size in meters | |
| Returns: | |
| Number of unique voxels occupied | |
| """ | |
| if len(points) == 0: | |
| return 0 | |
| # Convert points to voxel coordinates | |
| voxel_coords = np.floor(points / voxel_size).astype(int) | |
| # Find unique voxels | |
| unique_voxels = np.unique(voxel_coords, axis=0) | |
| return len(unique_voxels) | |
| def calculate_adaptive_sample_count_per_part( | |
| parts_points: List[np.ndarray], | |
| voxel_size: float, | |
| voxel_ratio: float, | |
| min_points_per_part: int, | |
| max_sample_points: int | |
| ) -> List[int]: | |
| """ | |
| Calculate the number of sample points for each part based on its occupied voxels. | |
| Args: | |
| parts_points: List of point cloud arrays for each part | |
| voxel_size: Voxel size in meters for spatial coverage calculation | |
| voxel_ratio: Ratio of occupied voxels to sample points | |
| min_points_per_part: Minimum number of points each part should have | |
| max_sample_points: Maximum number of sample points for voxel_adaptive method | |
| Returns: | |
| List of sample points to allocate for each part | |
| """ | |
| if not parts_points: | |
| return [] | |
| sample_counts_per_part = [] | |
| for i, points in enumerate(parts_points): | |
| if len(points) == 0: | |
| sample_counts_per_part.append(0) | |
| continue | |
| # Calculate voxel coverage for this part | |
| voxel_coverage = calculate_voxel_coverage(points, voxel_size) | |
| # Calculate adaptive sample count for this part | |
| adaptive_sample_count = int(voxel_coverage * voxel_ratio) | |
| # Apply min/max constraints for this part | |
| adaptive_sample_count = max(min_points_per_part, adaptive_sample_count) | |
| # Don't exceed available points in this part | |
| adaptive_sample_count = min(len(points), adaptive_sample_count) | |
| # Apply global max constraint | |
| adaptive_sample_count = min(max_sample_points, adaptive_sample_count) | |
| sample_counts_per_part.append(adaptive_sample_count) | |
| logger.debug(f"Part {i}: {voxel_coverage} occupied voxels -> {adaptive_sample_count} sample points " | |
| f"(ratio: {voxel_ratio}, voxel_size: {voxel_size}m)") | |
| total_sample_points = sum(sample_counts_per_part) | |
| logger.debug(f"Voxel-adaptive per-part sampling: total {total_sample_points} sample points across {len(parts_points)} parts") | |
| return sample_counts_per_part | |
| def allocate_fps_points( | |
| parts_data: Union[List[int], List[np.ndarray]], | |
| allocation_method: str, | |
| num_points: int, | |
| min_points_per_part: int, | |
| voxel_size: float, | |
| voxel_ratio: float, | |
| total_sample_points: Optional[List[int]] = None # for voxel_adaptive method | |
| ) -> np.ndarray: | |
| """ | |
| Allocate FPS points to parts proportionally with minimum constraints. | |
| Args: | |
| parts_data: Either list of point counts (for point_count method) or | |
| list of point arrays (for spatial_coverage and voxel_adaptive methods) | |
| allocation_method: Method for allocating FPS points ('point_count', 'spatial_coverage', or 'voxel_adaptive') | |
| num_points: Total number of sample points to allocate (for point_count and spatial_coverage methods) | |
| min_points_per_part: Minimum number of points each part should have after FPS | |
| voxel_size: Voxel size in meters for spatial coverage calculation (used in spatial_coverage and voxel_adaptive) | |
| voxel_ratio: Ratio of occupied voxels to sample points for voxel_adaptive method | |
| total_sample_points: Total number of sample points to allocate for voxel_adaptive (per-part counts) | |
| Returns: | |
| Array of target points to sample from each part | |
| """ | |
| if allocation_method == 'point_count': | |
| if isinstance(parts_data[0], np.ndarray): | |
| pts_per_part = np.array([len(part) for part in parts_data]) | |
| else: | |
| pts_per_part = np.array(parts_data) | |
| return _allocate_by_point_count_logic(pts_per_part, num_points, min_points_per_part) | |
| elif allocation_method == 'spatial_coverage': | |
| if isinstance(parts_data[0], np.ndarray): | |
| coverage_per_part = np.array([calculate_voxel_coverage(part, voxel_size) for part in parts_data]) | |
| pts_per_part = np.array([len(part) for part in parts_data]) | |
| else: | |
| raise ValueError("spatial_coverage allocation method requires point arrays, not point counts") | |
| return _allocate_by_spatial_coverage_logic(coverage_per_part, pts_per_part, num_points, min_points_per_part) | |
| elif allocation_method == 'voxel_adaptive': | |
| if isinstance(parts_data[0], np.ndarray): | |
| if total_sample_points is None: | |
| raise ValueError("voxel_adaptive allocation method requires total_sample_points parameter (per-part sample counts)") | |
| if not isinstance(total_sample_points, list): | |
| raise ValueError("voxel_adaptive allocation method requires total_sample_points to be a list of per-part counts") | |
| per_part_sample_counts = np.array(total_sample_points) | |
| pts_per_part = np.array([len(part) for part in parts_data]) | |
| target_per_part = np.minimum(per_part_sample_counts, pts_per_part) | |
| logger.debug(f"Voxel-adaptive per-part allocation (voxel_size={voxel_size}m, ratio={voxel_ratio}):") | |
| for i, (requested_count, available_count, final_count) in enumerate(zip(per_part_sample_counts, pts_per_part, target_per_part)): | |
| logger.debug(f" Part {i}: requested {requested_count}, available {available_count} -> {final_count} sampled") | |
| return target_per_part | |
| else: | |
| raise ValueError("voxel_adaptive allocation method requires point arrays, not point counts") | |
| else: | |
| raise ValueError(f"Unknown allocation method: {allocation_method}") | |
| def _allocate_by_point_count_logic(pts_per_part: np.ndarray, num_points: int, min_points_per_part: int) -> np.ndarray: | |
| """ | |
| Allocate FPS points based on point count (original method) - internal logic. | |
| """ | |
| n_parts = len(pts_per_part) | |
| min_per_part = np.minimum(min_points_per_part, pts_per_part) | |
| total_min_points = min_per_part.sum() | |
| if total_min_points > num_points: | |
| logger.warning(f"Minimum points per part ({min_points_per_part}) * {n_parts} parts = {total_min_points} " | |
| f"exceeds target ({num_points}). Scaling down.") | |
| scale_factor = num_points / total_min_points | |
| min_per_part = np.maximum(1, np.round(min_per_part * scale_factor).astype(int)) | |
| total_min_points = min_per_part.sum() | |
| target_per_part = min_per_part.copy() | |
| remaining_points = num_points - total_min_points | |
| if remaining_points > 0: | |
| remaining_capacity = pts_per_part - min_per_part | |
| total_capacity = remaining_capacity.sum() | |
| if total_capacity > 0: | |
| extra = np.round(remaining_capacity * remaining_points / total_capacity).astype(int) | |
| target_per_part = np.minimum(target_per_part + extra, pts_per_part) | |
| diff = num_points - target_per_part.sum() | |
| while diff != 0 and ((diff > 0 and (target_per_part < pts_per_part).sum() > 0) or | |
| (diff < 0 and (target_per_part > min_per_part).sum() > 0)): | |
| if diff > 0: | |
| valid_mask = target_per_part < pts_per_part | |
| else: | |
| valid_mask = target_per_part > min_per_part | |
| if valid_mask.sum() > 0: | |
| idx = np.random.choice(np.where(valid_mask)[0]) | |
| target_per_part[idx] += 1 if diff > 0 else -1 | |
| diff -= 1 if diff > 0 else -1 | |
| else: | |
| break | |
| return target_per_part | |
| def _allocate_by_spatial_coverage_logic(coverage_per_part: np.ndarray, pts_per_part: np.ndarray, num_points: int, min_points_per_part: int) -> np.ndarray: | |
| """ | |
| Allocate FPS points based on spatial coverage (voxel count) - internal logic. | |
| """ | |
| n_parts = len(coverage_per_part) | |
| min_per_part = np.minimum(min_points_per_part, pts_per_part) | |
| total_min_points = min_per_part.sum() | |
| if total_min_points > num_points: | |
| logger.warning(f"Minimum points per part ({min_points_per_part}) * {n_parts} parts = {total_min_points} " | |
| f"exceeds target ({num_points}). Scaling down.") | |
| scale_factor = num_points / total_min_points | |
| min_per_part = np.maximum(1, np.round(min_per_part * scale_factor).astype(int)) | |
| total_min_points = min_per_part.sum() | |
| target_per_part = min_per_part.copy() | |
| remaining_points = num_points - total_min_points | |
| if remaining_points > 0: | |
| remaining_capacity = pts_per_part - min_per_part | |
| total_coverage = coverage_per_part.sum() | |
| if total_coverage > 0: | |
| coverage_proportion = coverage_per_part / total_coverage | |
| extra_by_coverage = np.round(coverage_proportion * remaining_points).astype(int) | |
| extra_by_coverage = np.minimum(extra_by_coverage, remaining_capacity) | |
| target_per_part += extra_by_coverage | |
| diff = num_points - target_per_part.sum() | |
| iteration_count = 0 | |
| max_iterations = remaining_points | |
| while diff != 0 and iteration_count < max_iterations: | |
| if diff > 0 and (target_per_part < pts_per_part).sum() > 0: | |
| valid_mask = target_per_part < pts_per_part | |
| valid_indices = np.where(valid_mask)[0] | |
| if len(valid_indices) > 0: | |
| coverage_weights = coverage_per_part[valid_indices] | |
| if coverage_weights.sum() > 0: | |
| coverage_weights = coverage_weights / coverage_weights.sum() | |
| idx = valid_indices[np.random.choice(len(valid_indices), p=coverage_weights)] | |
| else: | |
| idx = np.random.choice(valid_indices) | |
| target_per_part[idx] += 1 | |
| diff -= 1 | |
| elif diff < 0 and (target_per_part > min_per_part).sum() > 0: | |
| valid_mask = target_per_part > min_per_part | |
| valid_indices = np.where(valid_mask)[0] | |
| if len(valid_indices) > 0: | |
| coverage_weights = coverage_per_part[valid_indices] | |
| if coverage_weights.sum() > 0: | |
| inv_weights = 1.0 / (coverage_weights + 1e-8) | |
| inv_weights = inv_weights / inv_weights.sum() | |
| idx = valid_indices[np.random.choice(len(valid_indices), p=inv_weights)] | |
| else: | |
| idx = np.random.choice(valid_indices) | |
| target_per_part[idx] -= 1 | |
| diff += 1 | |
| else: | |
| break | |
| iteration_count += 1 | |
| if iteration_count >= max_iterations: | |
| logger.warning(f"Fine-tuning reached maximum iterations ({max_iterations}), final diff: {diff}") | |
| return target_per_part | |
| def apply_batched_fps( | |
| batch_augmented_tensor: torch.Tensor, | |
| batch_lengths_tensor: torch.Tensor, | |
| batch_k_tensor: torch.Tensor, | |
| global_seed: int, | |
| device: torch.device | |
| ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: | |
| """ | |
| Apply batched FPS using PyTorch3D. | |
| Args: | |
| batch_augmented_tensor: Batched augmented point clouds (N, max_points, 3) | |
| batch_lengths_tensor: Lengths of point clouds in the batch (N,) | |
| batch_k_tensor: Number of points to sample for each point cloud (N,) | |
| global_seed: Global random seed for reproducible sampling | |
| device: Device to perform computation on | |
| Returns: | |
| Tuple of (list of sampled augmented parts, list of sampled indices) | |
| """ | |
| try: | |
| # Set the FPS-specific random seed for reproducible sampling | |
| torch.manual_seed(global_seed) | |
| # PyTorch3D FPS requires a contiguous tensor, ensure it | |
| batch_augmented_tensor = batch_augmented_tensor.contiguous() | |
| sampled_points, indices_tensor = sample_farthest_points( | |
| batch_augmented_tensor.to(device), | |
| lengths=batch_lengths_tensor.to(device), | |
| K=batch_k_tensor.to(device), | |
| random_start_point=True | |
| ) | |
| # Extract sampled points using indices (keep as tensors to avoid redundant conversions) | |
| sampled_augmented_parts = [] | |
| for i, (k_i, indices_i) in enumerate(zip(batch_k_tensor, indices_tensor)): | |
| # Get valid indices (not padded) - keep as tensor | |
| valid_indices = indices_i[:k_i] | |
| sampled_augmented_parts.append(batch_augmented_tensor[i][valid_indices]) | |
| return sampled_augmented_parts, indices_tensor | |
| except Exception as e: | |
| logger.error(f"PyTorch3D batched FPS failed: {e}") | |
| raise RuntimeError(f"PyTorch3D FPS failed: {e}") from e |