| 14 | 13: (219., 219., 141.), 14: (255., 127., 14.), 15: (158., 218., 229.), 16: (44., 160., 44.), 17: (112., 128., 144.), 18: (227., 119., 194.), 19: (82., 84., 163.)} |
| 15 | |
| 16 | class Voxelize(object): |
| 17 | def __init__(self, |
| 18 | voxel_size=0.05, |
| 19 | hash_type="fnv", |
| 20 | mode='train', |
| 21 | keys=("coord", "normal", "color", "label"), |
| 22 | return_discrete_coord=False, |
| 23 | return_min_coord=False): |
| 24 | self.voxel_size = voxel_size |
| 25 | self.hash = self.fnv_hash_vec if hash_type == "fnv" else self.ravel_hash_vec |
| 26 | assert mode in ["train", "test"] |
| 27 | self.mode = mode |
| 28 | self.keys = keys |
| 29 | self.return_discrete_coord = return_discrete_coord |
| 30 | self.return_min_coord = return_min_coord |
| 31 | |
| 32 | def __call__(self, data_dict): |
| 33 | assert "coord" in data_dict.keys() |
| 34 | discrete_coord = np.floor(data_dict["coord"] / np.array(self.voxel_size)).astype(int) |
| 35 | min_coord = discrete_coord.min(0) * np.array(self.voxel_size) |
| 36 | discrete_coord -= discrete_coord.min(0) |
| 37 | key = self.hash(discrete_coord) |
| 38 | idx_sort = np.argsort(key) |
| 39 | key_sort = key[idx_sort] |
| 40 | _, inverse, count = np.unique(key_sort, return_inverse=True, return_counts=True) |
| 41 | if self.mode == 'train': # train mode |
| 42 | # idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(), count.size) % count |
| 43 | idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) |
| 44 | idx_unique = idx_sort[idx_select] |
| 45 | if self.return_discrete_coord: |
| 46 | data_dict["discrete_coord"] = discrete_coord[idx_unique] |
| 47 | if self.return_min_coord: |
| 48 | data_dict["min_coord"] = min_coord.reshape([1, 3]) |
| 49 | for key in self.keys: |
| 50 | data_dict[key] = data_dict[key][idx_unique] |
| 51 | return data_dict |
| 52 | |
| 53 | elif self.mode == 'test': # test mode |
| 54 | data_part_list = [] |
| 55 | for i in range(count.max()): |
| 56 | idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + i % count |
| 57 | idx_part = idx_sort[idx_select] |
| 58 | data_part = dict(index=idx_part) |
| 59 | for key in data_dict.keys(): |
| 60 | if key in self.keys: |
| 61 | data_part[key] = data_dict[key][idx_part] |
| 62 | else: |
| 63 | data_part[key] = data_dict[key] |
| 64 | if self.return_discrete_coord: |
| 65 | data_part["discrete_coord"] = discrete_coord[idx_part] |
| 66 | if self.return_min_coord: |
| 67 | data_part["min_coord"] = min_coord.reshape([1, 3]) |
| 68 | data_part_list.append(data_part) |
| 69 | return data_part_list |
| 70 | else: |
| 71 | raise NotImplementedError |
| 72 | |
| 73 | @staticmethod |