| import torch | |
| from einops import rearrange, repeat | |
| from ....geometry.projection import get_world_rays | |
| from ....geometry.projection import sample_image_grid | |
| from ....test.export_ply import save_point_cloud_to_ply | |
| import torch.nn.functional as F | |
| import MinkowskiEngine as ME | |
| def project_features_to_me(intrinsics, extrinsics, out, depth, voxel_resolution, b, v): | |
| device = out.device | |
| # 0. 获取基础维度信息 | |
| h, w = depth.shape[2:] | |
| _, c, _, _ = out.shape | |
| # 1. 准备投影参数 | |
| intrinsics = rearrange(intrinsics, "b v i j -> b v () () () i j") | |
| extrinsics = rearrange(extrinsics, "b v i j -> b v () () () i j") | |
| depths = rearrange(depth, "b v h w -> b v (h w) () ()") | |
| # 2. 获取世界坐标 | |
| uv_grid = sample_image_grid((h, w), device)[0] | |
| uv_grid = repeat(uv_grid, "h w c -> 1 v (h w) () () c", v=v) | |
| origins, directions = get_world_rays(uv_grid, extrinsics, intrinsics) | |
| world_coords = origins + directions * depths[..., None] | |
| world_coords = world_coords.squeeze(3).squeeze(3) # [B, V, N, 3] | |
| # 3. 准备特征数据 | |
| features = rearrange(out, "(b v) c h w -> b v c h w", b=b, v=v) | |
| features = rearrange(features, "b v c h w -> b v h w c") | |
| features = rearrange(features, "b v h w c -> b v (h w) c") # [B, V, N, C] | |
| # 4. 合并所有点云数据 | |
| all_points = rearrange(world_coords, "b v n c -> (b v n) c") # [B*V*N, 3] | |
| feats_flat = features.reshape(-1, c) # [B*V*N, C] | |
| # 5. 分离量化操作 | |
| with torch.no_grad(): | |
| # 计算量化坐标 - 使用四舍五入而不是向下取整 | |
| quantized_coords = torch.round(all_points / voxel_resolution).long() | |
| # 创建坐标矩阵：批次索引 + 量化坐标 | |
| batch_indices = torch.arange(b, device=device).repeat_interleave(v * h * w).unsqueeze(1) | |
| combined_coords = torch.cat([batch_indices, quantized_coords], dim=1) | |
| # 获取唯一体素ID和映射索引 | |
| unique_coords, inverse_indices, counts = torch.unique( | |
| combined_coords, | |
| dim=0, | |
| return_inverse=True, | |
| return_counts=True | |
| ) | |
| # 6. 创建聚合特征和坐标 | |
| num_voxels = unique_coords.shape[0] | |
| # 7. 向量化聚合 - 特征平均 | |
| aggregated_feats = torch.zeros(num_voxels, c, device=device) | |
| aggregated_feats.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, c), feats_flat) | |
| aggregated_feats = aggregated_feats / counts.view(-1, 1).float() # 平均特征 | |
| # 8. 向量化聚合 - 坐标平均 | |
| aggregated_points = torch.zeros(num_voxels, 3, device=device) | |
| aggregated_points.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, 3), all_points) | |
| aggregated_points = aggregated_points / counts.view(-1, 1).float() | |
| # 9. 创建稀疏张量 | |
| # 使用正确的坐标格式：批次索引 + 量化坐标 | |
| sparse_tensor = ME.SparseTensor( | |
| features=aggregated_feats, | |
| coordinates=unique_coords.int(), | |
| tensor_stride=1, | |
| device=device | |
| ) | |
| # 10. 返回结果 | |
| return sparse_tensor, aggregated_points, counts | |
| # def project_features_to_me(intrinsics, extrinsics, out, depth, voxel_resolution, b, v): | |
| # device = out.device | |
| # # 0. 获取基础维度信息 | |
| # b, v = intrinsics.shape[:2] | |
| # h, w = depth.shape[2:] | |
| # _, c, _, _ = out.shape | |
| # # 1. 准备投影参数 | |
| # intrinsics = rearrange(intrinsics, "b v i j -> b v () () () i j") | |
| # extrinsics = rearrange(extrinsics, "b v i j -> b v () () () i j") | |
| # depths = rearrange(depth, "b v h w -> b v (h w) () ()") | |
| # # 2. 获取世界坐标 | |
| # uv_grid = sample_image_grid((h, w), device)[0] | |
| # uv_grid = repeat(uv_grid, "h w c -> 1 v (h w) () () c", v=v) | |
| # origins, directions = get_world_rays(uv_grid, extrinsics, intrinsics) | |
| # world_coords = origins + directions * depths[..., None] | |
| # world_coords = world_coords.squeeze(3).squeeze(3) # [B, V, N, 3] | |
| # # 3. 准备特征数据 | |
| # features = rearrange(out, "(b v) c h w -> b v c h w", b=b, v=v) | |
| # features = rearrange(features, "b v c h w -> b v h w c") | |
| # features = rearrange(features, "b v h w c -> b v (h w) c") # [B, V, N, C] | |
| # # 4. 合并所有点云数据 | |
| # all_points = rearrange(world_coords, "b v n c -> (b v n) c") # [B*V*N, 3] | |
| # feats_flat = features.reshape(-1, c) # [B*V*N, C] | |
| # # 5. 分离量化操作 | |
| # with torch.no_grad(): | |
| # # 计算量化坐标 | |
| # quantized_coords = (all_points / voxel_resolution).floor().long() | |
| # # 创建坐标矩阵：批处理索引 + 量化坐标 | |
| # batch_indices = torch.arange(b * v, device=device).repeat_interleave(h * w).unsqueeze(1) | |
| # combined_coords = torch.cat([batch_indices, quantized_coords], dim=1) | |
| # # 获取唯一体素ID和映射索引 | |
| # unique_coords, inverse_indices, counts = torch.unique( | |
| # combined_coords, | |
| # dim=0, | |
| # return_inverse=True, | |
| # return_counts=True | |
| # ) | |
| # # 6. 创建聚合特征和坐标 | |
| # num_voxels = unique_coords.shape[0] | |
| # # 7. 向量化聚合 - 特征求和 | |
| # aggregated_feats = torch.zeros(num_voxels, c, device=device) | |
| # aggregated_feats.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, c), feats_flat) | |
| # # 8. 向量化聚合 - 坐标平均 | |
| # aggregated_points = torch.zeros(num_voxels, 3, device=device) | |
| # aggregated_points.scatter_add_(0, inverse_indices.unsqueeze(1).expand(-1, 3), all_points) | |
| # aggregated_points = aggregated_points / counts.view(-1, 1).float() | |
| # # 9. 创建稀疏张量 | |
| # # 关键修改：使用相同的索引顺序 | |
| # # 确保 sparse_tensor 和 aggregated_points 对应相同的体素 | |
| # sparse_tensor = ME.SparseTensor( | |
| # features=aggregated_feats, | |
| # coordinates=torch.cat([ | |
| # torch.arange(num_voxels, device=device).unsqueeze(1), # 添加批次索引 | |
| # unique_coords[:, 1:].float() # 量化坐标 | |
| # ], dim=1).int(), | |
| # tensor_stride=1, | |
| # device=device | |
| # ) | |
| # # 10. 返回结果 | |
| # # 现在 sparse_tensor 和 aggregated_points 有相同的点数 (num_voxels) | |
| # return sparse_tensor, aggregated_points, counts | |