Spaces:
Running on Zero
Running on Zero
File size: 13,843 Bytes
be88765 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 | 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 |