import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt

pseudo = np.load('path/VoxFormer/mmdetection3d/VoxFormer_UQ/kitti/dataset/sequences_msnet3dmax1std_sweep1/00/voxels/000070.npy')


# 获取非零元素的索引
non_zero_indices = np.argwhere(pseudo > 0.1)

# 获取非零值并排序，计算归一化的颜色映射索引
non_zero_values = pseudo[non_zero_indices[:, 0], non_zero_indices[:, 1], non_zero_indices[:, 2]]
sorted_norm_indices = np.argsort(np.argsort(non_zero_values)).astype(float) / len(non_zero_values)

# # 使用归一化的索引映射到颜色
# cmap = plt.get_cmap('hot')
# colors = cmap(sorted_norm_indices)[:,:3]  # 获取RGB颜色

# 使用 'jet' 颜色映射
cmap = plt.get_cmap('jet')
colors = cmap(sorted_norm_indices)[:,:3]  # 获取RGB颜色

# 设置点云的点和颜色
point_cloud = o3d.geometry.PointCloud()
point_cloud.points = o3d.utility.Vector3dVector(non_zero_indices)
point_cloud.colors = o3d.utility.Vector3dVector(colors)

# 可视化点云数据
o3d.visualization.draw_geometries([point_cloud])
