def spatially_regular_gen(self, item): # Choose the cloud with the lowest probability cloud_idx = int(np.argmin(self.min_possibility[self.mode])) # choose the point with the minimum of possibility in the cloud as query point point_ind = np.argmin(self.possibility[self.mode][cloud_idx]) # Get all points within the cloud from tree structure points = np.array(self.input_trees[self.mode][cloud_idx].data, copy=False) # Center point of input region center_point = points[point_ind, :].reshape(1, -1) # Add noise to the center point noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) pick_point = center_point + noise.astype(center_point.dtype) # Check if the number of points in the selected cloud is less than the predefined num_points if len(points) < cfg.num_points: # Query all points within the cloud queried_idx = self.input_trees[self.mode][cloud_idx].query( pick_point, k=len(points))[1][0] else: # Query the predefined number of points queried_idx = self.input_trees[self.mode][cloud_idx].query( pick_point, k=cfg.num_points)[1][0] # Shuffle index queried_idx = DP.shuffle_idx(queried_idx) # Get corresponding points and colors based on the index queried_pc_xyz = points[queried_idx] queried_pc_xyz = queried_pc_xyz - pick_point queried_pc_colors = self.input_colors[ self.mode][cloud_idx][queried_idx] queried_pc_labels = self.input_labels[ self.mode][cloud_idx][queried_idx] # Update the possibility of the selected points dists = np.sum(np.square( (points[queried_idx] - pick_point).astype(np.float32)), axis=1) delta = np.square(1 - dists / np.max(dists)) self.possibility[self.mode][cloud_idx][queried_idx] += delta self.min_possibility[self.mode][cloud_idx] = float( np.min(self.possibility[self.mode][cloud_idx])) # up_sampled with replacement if len(points) < cfg.num_points: queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) return queried_pc_xyz.astype(np.float32), queried_pc_colors.astype( np.float32), queried_pc_labels, queried_idx.astype( np.int32), np.array([cloud_idx], dtype=np.int32)
def spatially_regular_gen(): # Generator loop for i in range(num_per_epoch): # num_per_epoch # Choose a random cloud cloud_idx = int(np.argmin(self.min_possibility[split])) # choose the point with the minimum of possibility as query point point_ind = np.argmin(self.possibility[split][cloud_idx]) # Get points from tree structure points = np.array(self.input_trees[split][cloud_idx].data, copy=False) # Center point of input region center_point = points[point_ind, :].reshape(1, -1) # Add noise to the center point noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) pick_point = center_point + noise.astype(center_point.dtype) if len(points) < cfg.num_points: queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] else: queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] queried_idx = DP.shuffle_idx(queried_idx) # Collect points and colors queried_pc_xyz = points[queried_idx] queried_pc_xyz = queried_pc_xyz - pick_point queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) delta = np.square(1 - dists / np.max(dists)) self.possibility[split][cloud_idx][queried_idx] += delta self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) if len(points) < cfg.num_points: queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) if True: yield (queried_pc_xyz.astype(np.float32), queried_pc_colors.astype(np.float32), queried_pc_labels, queried_idx.astype(np.int32), np.array([cloud_idx], dtype=np.int32))