def projection_plot(pcpath='', noise=0.05, outlier=0.05, savefig=False): f_list = [ pcpath + '/' + i for i in os.listdir(pcpath) if os.path.splitext(i)[1] == '.ply' ] fig = plt.figure(figsize=(38, 20), dpi=600, facecolor='w') colunms = 8 for i, j in enumerate(f_list): for k in range(colunms): pc = PointCloud(j) pc.down_sample(number_of_downsample=10000) pc.add_noise(noise) pc.add_outlier(outlier) pts_size = 2.5 if i == 7: pts_size = 1 try: mfig = pc.half_by_plane(n=1024, grid_resolution=(200, 200), show_result=pts_size) except: try: mfig = pc.half_by_plane(n=1024, grid_resolution=(250, 250), show_result=pts_size) except: try: mfig = pc.half_by_plane(n=1024, grid_resolution=(300, 300), show_result=pts_size) except: mfig = pc.half_by_plane(n=1024, grid_resolution=(650, 650), show_result=pts_size) f = mlab.gcf() # this two line for mlab.screenshot to work f.scene._lift() if savefig: mlab.savefig(str(i) + str(k) + '.png') img = mlab.screenshot(figure=mfig) mlab.close() ax = fig.add_subplot(len(f_list), colunms, i * colunms + k + 1) ax.imshow(img) ax.set_axis_off() plt.subplots_adjust(wspace=0, hspace=0) if savefig: plt.savefig('projection.png') plt.show() plt.close()
def scene_seg_dataset(pc_path, save_path, samples=1000, max_nb_pc=5, show_result=False): """ default number of points of each point cloud is 1024 :param pc_path: :param save_path: :param max_nb_pc: :return: """ f_list = [ PointCloud(pc_path + '/' + i) for i in os.listdir(pc_path) if os.path.splitext(i)[1] == '.ply' ] for i in f_list: i.down_sample() nb_classes = len(f_list) scene_dataset = np.zeros((samples, max_nb_pc * 1024, 3)) scene_label = np.zeros((samples, max_nb_pc * 1024), dtype=np.int32) rate_180, rate_240, rate_300, rate_360 = [0, 0, 0, 0] for i in range(samples): print('generating the {}th scene sample'.format(i)) nb_pc = np.random.choice(max_nb_pc) + 1 nb_pc = max_nb_pc for j in range(nb_pc): k = np.random.choice(nb_classes) pc = f_list[k] pc.transform() pc.cut_by_plane() pc2 = PointCloud(pc.visible) try: pc2.half_by_plane(n=1024, grid_resolution=(190, 190)) rate_180 += 1 except: try: pc2.half_by_plane(n=1024, grid_resolution=(260, 260)) rate_240 += 1 except: try: pc2.half_by_plane(n=1024, grid_resolution=(330, 330)) rate_300 += 1 except: pc2.half_by_plane(n=1024, grid_resolution=(400, 400)) rate_360 += 1 scene_dataset[i, j * 1024:j * 1024 + 1024, :] = pc2.visible scene_label[i, j * 1024:j * 1024 + 1024] = k print('180 240 300 360:', rate_180, rate_240, rate_300, rate_360) if show_result: for i in range(1): scene_pc = scene_dataset[i, :, :] scene_pc = PointCloud(scene_pc) # scene_lb = scene_label[i, :] figure = mlab.figure(size=(1000, 1000), bgcolor=(1, 1, 1)) colors = (np.random.random((nb_classes, 4)) * 255).astype(np.int8) colors[:, -1] = 255 colors = colors[scene_lb, :] scalars = np.arange(np.shape(colors)[0]) pts = mlab.quiver3d(scene_pc.position[:, 0], scene_pc.position[:, 1], scene_pc.position[:, 2], scene_pc.position[:, 0] * 10**-9 + 1, scene_pc.position[:, 0] * 10**-9 + 1, scene_pc.position[:, 0] * 10**-9 + 1, scalars=scalars, scale_factor=1, mode='sphere', figure=figure) pts.glyph.color_mode = 'color_by_scalar' pts.module_manager.scalar_lut_manager.lut.table = colors mlab.show() hdf5_file = h5py.File(save_path, mode='a') hdf5_file.create_dataset('train_set', (samples, max_nb_pc * 1024, 3), np.float32) # be careful about the dtype hdf5_file.create_dataset('train_labels', (samples, max_nb_pc * 1024), np.uint8) hdf5_file["train_set"][...] = scene_dataset hdf5_file["train_labels"][...] = scene_label hdf5_file.close()