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

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

点云深度学习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_size
        self.hash = self.fnv_hash_vec if hash_type == "fnv" else self.ravel_hash_vec
        assert mode in ["train", "test"]
        self.mode = mode
        self.keys = keys
        self.return_inverse = return_inverse
        self.return_discrete_coord = return_discrete_coord
        self.return_min_coord = return_min_coord

    def __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 mode
            idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(), count.size) % count
            idx_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] = 1
                data_dict["inverse"] = np.zeros_like(inverse)
                data_dict["inverse"][idx_sort] = inverse
                data_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_dict

        elif self.mode == 'test':  # test mode
            data_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 % count
                idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + i % count
                idx_part = idx_sort[idx_select]
                data_part = dict(index=idx_part)

                # TODO to be more robust
                for 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] = inverse
                    data_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_list
        else:
            raise NotImplementedError

    @staticmethod
    def ravel_hash_vec(arr):
        """
        Ravel the coordinates after subtracting the min coordinates.
        """
        assert arr.ndim == 2
        arr = arr.copy()
        arr -= arr.min(0)
        arr = arr.astype(np.uint64, copy=False)
        arr_max = arr.max(0).astype(np.uint64) + 1

        keys = np.zeros(arr.shape[0], dtype=np.uint64)
        # Fortran style indexing
        for j in range(arr.shape[1] - 1):
            keys += arr[:, j]
            keys *= arr_max[j + 1]
        keys += arr[:, -1]
        return keys

    @staticmethod
    def fnv_hash_vec(arr):
        """
        FNV64-1A
        """
        assert arr.ndim == 2
        # Floor first for negative coordinates
        arr = 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_dict

data2 = 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]]