点云规则格网化，且保存原始的点云索引

点云规则格网化，且保存原始的点云索引

点云深度学习Voxelize规则，参考PTV2：https://github.com/Gofinge/PointTransformerV2

1总执行文件

import numpy as np
import torch
from pcr.utils.registry import Registry
TRANSFORMS = Registry("transforms")
@TRANSFORMS.register_module()
class Voxelize(object):def __init__(self,voxel_size=0.05,hash_type="fnv",mode='train',keys=("coord", "normal", "color", "label"),return_inverse=False,return_discrete_coord=False,return_min_coord=False):self.voxel_size = voxel_sizeself.hash = self.fnv_hash_vec if hash_type == "fnv" else self.ravel_hash_vecassert mode in ["train", "test"]self.mode = modeself.keys = keysself.return_inverse = return_inverseself.return_discrete_coord = return_discrete_coordself.return_min_coord = return_min_coorddef __call__(self, data_dict):assert "coord" in data_dict.keys()discrete_coord = np.floor(data_dict["coord"] / np.array(self.voxel_size)).astype(np.int)min_coord = discrete_coord.min(0) * np.array(self.voxel_size)discrete_coord -= discrete_coord.min(0)key = self.hash(discrete_coord)idx_sort = np.argsort(key)key_sort = key[idx_sort]_, inverse, count = np.unique(key_sort, return_inverse=True, return_counts=True)if self.mode == 'train':  # train modeidx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(), count.size) % countidx_unique = idx_sort[idx_select]if self.return_discrete_coord:data_dict["discrete_coord"] = discrete_coord[idx_unique]if self.return_inverse:data_dict["mask"] = np.zeros_like(inverse)data_dict["mask"][idx_unique] = 1data_dict["inverse"] = np.zeros_like(inverse)data_dict["inverse"][idx_sort] = inversedata_dict["length"] = np.array(inverse.shape)if self.return_min_coord:data_dict["min_coord"] = min_coord.reshape([1, 3])for key in self.keys:data_dict[key] = data_dict[key][idx_unique]# print('data_dict["discrete_coord"].shape',data_dict["discrete_coord"].shape,' ',data_dict[key].shape)return data_dictelif self.mode == 'test':  # test modedata_part_list = []for i in range(count.max()):temp=np.insert(count, 0, 0)temp2=temp[0: -1]temp3= np.cumsum(temp2)temp4 = np.cumsum(temp2)+i % countidx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + i % countidx_part = idx_sort[idx_select]data_part = dict(index=idx_part)# TODO to be more robustfor key in self.keys:data_part[key] = data_dict[key][idx_part]if self.return_discrete_coord:data_part["discrete_coord"] = discrete_coord[idx_part]if self.return_inverse:data_part["inverse"] = np.zeros_like(inverse)data_part["inverse"][idx_sort] = inversedata_part["length"] = np.array(inverse.shape)if self.return_min_coord:data_part["min_coord"] = min_coord.reshape([1, 3])data_part_list.append(data_part)return data_part_listelse:raise NotImplementedError@staticmethoddef ravel_hash_vec(arr):"""Ravel the coordinates after subtracting the min coordinates."""assert arr.ndim == 2arr = arr.copy()arr -= arr.min(0)arr = arr.astype(np.uint64, copy=False)arr_max = arr.max(0).astype(np.uint64) + 1keys = np.zeros(arr.shape[0], dtype=np.uint64)# Fortran style indexingfor j in range(arr.shape[1] - 1):keys += arr[:, j]keys *= arr_max[j + 1]keys += arr[:, -1]return keys@staticmethoddef fnv_hash_vec(arr):"""FNV64-1A"""assert arr.ndim == 2# Floor first for negative coordinatesarr = arr.copy()arr = arr.astype(np.uint64, copy=False)hashed_arr = np.uint64(14695981039346656037) * np.ones(arr.shape[0], dtype=np.uint64)for j in range(arr.shape[1]):hashed_arr *= np.uint64(1099511628211)hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])return hashed_arr
class Compose(object):def __init__(self, cfg=None):self.cfg = cfg if cfg is not None else []self.transforms = []for t_cfg in self.cfg:self.transforms.append(TRANSFORMS.build(t_cfg))def __call__(self, data_dict):for t in self.transforms:data_dict = t(data_dict)return data_dictdata2 = torch.load('/media/1.pth')
Voxelize()
transform = Compose([dict(type="Voxelize", voxel_size=0.5, hash_type='fnv', mode='test',keys=("coord", "color", "semantic_gt"), return_discrete_coord=True)])
data2['coord']=np.zeros((8,3))
data2['coord'][:,0]=[9,7,1.01,1.02,3,4.01,4.02,4.03]
data2['coord'][:,1]=[9,7,1.01,1.02,3,4.01,4.02,4.03]
data2['coord'][:,2]=[9,7,1.01,1.02,3,4.01,4.02,4.03]
data2_voxelize = transform(data2)
# coord_p, idx_uni = np.random.rand(data["coord"].shape[0]) * 1e-3, np.array([])
# print(idx_uni.size)
for i in range(3):print(data2_voxelize[i]['coord'])
s=1

输入

data2[‘coord’]=np.zeros((8,3))
data2[‘coord’][:,0]=[9,7,1.01,1.02,3,4.01,4.02,4.03]
data2[‘coord’][:,1]=[9,7,1.01,1.02,3,4.01,4.02,4.03]
data2[‘coord’][:,2]=[9,7,1.01,1.02,3,4.01,4.02,4.03]

输出

[[9. 9. 9. ]
[7. 7. 7. ]
[4.01 4.01 4.01]
[3. 3. 3. ]
[1.01 1.01 1.01]]

[[9. 9. 9. ]
[7. 7. 7. ]
[4.02 4.02 4.02]
[3. 3. 3. ]
[1.02 1.02 1.02]]

[[9. 9. 9. ]
[7. 7. 7. ]
[4.03 4.03 4.03]
[3. 3. 3. ]
[1.01 1.01 1.01]]