def saliancey2range(resolution_control=0.005):
    for j, i in enumerate(f_list):
        print(' point cloud is', i)
        pc = PointCloud(i)
        pc.down_sample(number_of_downsample=2048)
        for k in range(4):
            if k == 0:
                k = -0.5
            fig = pc.compute_key_points(percentage=0.1,
                                        show_result=False,
                                        resolution_control=resolution_control,
                                        rate=0.05 * k + 0.05,
                                        use_deficiency=False,
                                        show_saliency=True)
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            img = mlab.screenshot()
            mlab.savefig(filename=str(j) + str(k) + '_without.png')
            mlab.close()
            fig = pc.compute_key_points(percentage=0.1,
                                        show_result=False,
                                        resolution_control=resolution_control,
                                        rate=0.05 * k + 0.05,
                                        use_deficiency=True,
                                        show_saliency=True)
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            img = mlab.screenshot()
            mlab.savefig(filename=str(j) + str(k) + '_with.png')
            mlab.close()

        del pc
def pose_estimation(posefile='', real_single_h5='', model_filepath=''):
    scene_idx = 19  # 0-53 for scene objcet 11 3 19 2
    model_idx = 2  # 0-7 for 8 class of model objcet

    poseset = tables.open_file(posefile, mode='r')
    random_pose = poseset.root.random_pose[scene_idx, :]
    predict_pose = poseset.root.predict_pose[scene_idx, :]
    print('random_pose:',
          poseset.root.random_pose[[0, 1, 4, 5, 6, 9, 11, 19], :])
    print('predict_pose:',
          poseset.root.predict_pose[[0, 1, 4, 5, 6, 9, 11, 19], :])

    readh5 = h5py.File(real_single_h5)
    for i in [0, 1, 4, 5, 6, 9, 11, 19]:
        scene_pc = readh5['train_set'][i, :]  # n * 1024 * 3
        scene_pc = PointCloud(scene_pc)
        scene_pc.save(path='pointcloud/fourkind/' + str(i) + '.ply')
    print('scene_pc:', scene_pc)

    model_pc = [
        model_filepath + '/' + i for i in os.listdir(model_filepath)
        if os.path.splitext(i)[1] == '.ply'
    ][model_idx]
    model_pc = PointCloud(model_pc)
    model_pc.down_sample()

    light = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [1.0, 1.0, 0.0],
                      [0.0, 1.0, 0.0]])
    shade = light * 0.7
    light1 = tuple(light[0, :].tolist())
    shade1 = tuple(shade[0, :].tolist())
    light2 = tuple(light[1, :].tolist())
    shade2 = tuple(shade[1, :].tolist())
    light3 = tuple(light[2, :].tolist())
    shade3 = tuple(shade[2, :].tolist())
    light4 = tuple(light[3, :].tolist())
    shade4 = tuple(shade[3, :].tolist())

    colorset = [[light4, light2], [shade2, light2], [shade3, light3],
                [shade4, light4]]

    # initial pose :
    fig = show_pc.show_trans(scene_pc,
                             model_pc,
                             colorset=colorset,
                             scale=1,
                             returnfig=True)

    filename1 = 'poseestimation/real/before_alignment1.png'
    while (True):
        if os.path.exists(filename1):
            filename1 = filename1.split('.')[0][:-1] + str(
                int(filename1.split('.')[0][-1]) + 1) + '.png'
            continue
        break
    f = mlab.gcf()  # this two line for mlab.screenshot to work
    f.scene._lift()
    mlab.savefig(filename=filename1)
    print('before image saved')
    mlab.close()
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 rneighbor2range():
    for j, i in enumerate(f_list):
        print(' point cloud is', i)
        pc = PointCloud(i)
        pc.down_sample(number_of_downsample=2048)
        for k in range(4):
            fig = pc.generate_r_neighbor(range_rate=0.025 * k + 0.025,
                                         show_result=True)
            pc.keypoints = None
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            mlab.savefig(filename=str(j) + str(k) + 'r.png')
            mlab.close()
        del pc
def knn_plot(pc_path=''):

    f_list = [
        base_path + '/' + i for i in os.listdir(base_path)
        if os.path.splitext(i)[1] == '.ply'
    ]
    for j, i in enumerate(f_list):
        if j < 4:
            pc = PointCloud(i)
            pc.down_sample(number_of_downsample=4096)
            pc.add_noise(factor=0.04)
            pc.add_outlier(factor=0.04)
            fig = pc.compute_key_points(
                percentage=0.02,
                resolution_control=1 / 15,
                rate=0.05,
                use_deficiency=False,
                show_result=True)  # get the key points id

            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            mlab.savefig(filename=str(j) + '_0.png')
            mlab.close()

            colorset = np.random.random((100, 3))
            fig = pc.generate_k_neighbor(k=32,
                                         show_result=True,
                                         colorset=colorset)

            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            mlab.savefig(filename=str(j) + '_1.png')
            mlab.close()

            fig = pc.generate_k_neighbor(k=64,
                                         show_result=True,
                                         colorset=colorset)
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            mlab.savefig(filename=str(j) + '_2.png')
            mlab.close()

            fig = pc.generate_k_neighbor(k=128,
                                         show_result=True,
                                         colorset=colorset)
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            mlab.savefig(filename=str(j) + '_3.png')
            mlab.close()
def txt2normalply(txt_path, write_path='/ply/'):
    """

    :param txt_path:
    :param write_path:
    :return: nothing, write file into write_path
    """
    for i, j, k in os.walk(txt_path):
        if i == txt_path:
            for m, l in enumerate(k):
                a = np.loadtxt(i + '/' + l)
                PC = PointCloud(a)
                PC.down_sample(number_of_downsample=1024)

                pc = o3d.PointCloud()
                pc.points = o3d.Vector3dVector(PC.position)
                o3d.estimate_normals(
                    pc, o3d.KDTreeSearchParamHybrid(radius=10, max_nn=10))

                o3d.write_point_cloud(i + write_path + str(m) + '.ply', pc)
def key_points_plot(flist):
    for i in flist:
        Pc = PointCloud(i)
        Pc.down_sample(4096)
        fig = Pc.compute_key_points(percentage=0.1,
                                    resolution_control=None,
                                    show_result=True)

        f = mlab.gcf()  # this two line for mlab.screenshot to work
        f.scene._lift()
        img = mlab.screenshot()
        mlab.savefig(filename=str(i) + 'key_points.png')
        mlab.close()

        fig = Pc.compute_key_points(percentage=0.1,
                                    resolution_control=0.01,
                                    show_result=True)

        f = mlab.gcf()  # this two line for mlab.screenshot to work
        f.scene._lift()
        img = mlab.screenshot()
        mlab.savefig(filename=str(i) + 'key_points_with_resolution_ctrl.png')
        mlab.close()
def noise_outliers(pointclous, noise=0.05, outliers=0.05):

    fig = plt.figure(figsize=(38, 20), dpi=600, facecolor='w')
    for j, i in enumerate(pointclous):
        pc = PointCloud(i)
        pc.down_sample(number_of_downsample=1024)

        for k in range(4):
            if k == 3:
                k = 4
            pc.add_noise(factor=k * noise)
            pc.add_outlier(factor=k * outliers)
            m_fig = mlab.figure(bgcolor=(1, 1, 1))

            mlab.points3d(pc.position[:, 0],
                          pc.position[:, 1],
                          pc.position[:, 2],
                          pc.position[:, 2] * 10**-2 + 1,
                          colormap='Spectral',
                          scale_factor=2,
                          figure=m_fig)

            # mlab.gcf().scene.parallel_projection = True  # parallel projection
            f = mlab.gcf()  # this two line for mlab.screenshot to work
            f.scene._lift()
            # mlab.show()  # for testing
            img = mlab.screenshot(figure=m_fig)
            mlab.close()
            if k == 4:
                k = 3
            ax = fig.add_subplot(4, 8, (j + 1) + k * 8)
            ax.imshow(img)
            ax.set_axis_off()

    plt.subplots_adjust(wspace=0, hspace=0)
    plt.show()
def augment_data(base_path='',
                 pc_path='',
                 add_noise=0.04,
                 add_outlier=0.04,
                 n=5000,
                 not_project=False,
                 show_result=False):

    pc = PointCloud(pc_path)
    pc.down_sample()

    if add_noise is not None:
        pc.add_noise(factor=add_noise)
    if add_outlier is not None:
        pc.add_outlier(factor=add_outlier)

    if not_project:
        for i in range(n):
            if i % 10 == 0:
                print('saving number', i + 1,
                      'th lab random sample point clouds')
            temp = deepcopy(pc)
            temp.down_sample(number_of_downsample=1024)
            np.savetxt(base_path + '/random_sample' + str(i) + '.txt',
                       temp.position,
                       delimiter=' ')
    else:
        for i in range(n):
            if i % 10 == 0:
                print('saving number', i + 1, 'th lab_project point clouds')

            #pc.cut_by_plane()  # todo manually
            #pc2 = PointCloud(pc.visible)
            pc2 = pc
            try:
                pc2.half_by_plane(n=1024, grid_resolution=(200, 200))
            except:
                try:
                    pc2.half_by_plane(n=1024, grid_resolution=(250, 250))
                except:
                    try:
                        pc2.half_by_plane(n=1024, grid_resolution=(300, 300))
                    except:
                        pc2.half_by_plane(n=1024, grid_resolution=(650, 650))

            np.savetxt(base_path + '/lab_project' + str(i) + '.txt',
                       pc2.visible,
                       delimiter=' ')  # pc.visible will variant

    if show_result:
        dir_list = [
            base_path + '/' + i for i in os.listdir(base_path)
            if os.path.isdir(i)
        ]
        fig = mlab.figure(size=(1000, 1000), bgcolor=(1, 1, 1))
        for i in dir_list:
            color = np.random.random((1, 3))
            pc = np.loadtxt(i + '/lab_project1')
            mlab.points3d(pc[:, 0],
                          pc[:, 1],
                          pc[:, 2],
                          pc[:, 2] * 10**-9,
                          color=color,
                          figure=fig)
def save_data(save_path='',
              base_path='',
              n=5000,
              use_key_feature=True,
              train_data=True,
              nb_types=4,
              show_result=False,
              normalize=True,
              shuffle=True):
    """
    transform the txt point clouds into h5py dataset for simplicity. data augmentation of projection is implemented here
    :param save_path:
    :param n:
    :param train_data whether it is training data or it is test data. if its testdata, label is random.
    :param base_path:  path contains txt or ply point cloud data
    :param use_key_feature: if you want to use the local key features
    :param nb_types: number of classes of used object
    :return:
    """
    compute_time = []
    if train_data:
        pc_tile = np.empty(shape=(nb_types * n, 1024, 3))
        if use_key_feature:
            pc_key_feature = np.empty(shape=(nb_types * n, int(
                1024 * 0.1), 9))  # key feature space, 102=1024*0.1,
            # 9 for multi-scale eigen-value
            #pc_pl = tf.placeholder(tf.float32, shape=(1, 1024, 3))

        for k in range(
                nb_types
        ):  # number of types of  objects model, test data can ignore type label
            for i, j in enumerate(range(k * n, (k + 1) * n)):

                if i % 10 == 0:
                    print('reading number', i + 1,
                          'th lab' + str(k + 1) + ' point clouds')

                if use_key_feature:
                    pc = np.loadtxt(
                        base_path + '/lab' + str(k + 1) + '/lab_project' +
                        str(i) + '.txt'
                    )  # pc = tf.convert_to_tensor(pc, dtype=tf.float32)
                    pc = PointCloud(pc)
                    if normalize:
                        pc.normalize(
                        )  # partial point cloud should not normalize

                    expand = np.expand_dims(pc.position, axis=0)
                    pc_tile[j, :, :] = expand
                    # print('*****************************************')
                    # print('reading point cloud cost time:{}'.format(t1 - t0))

                    pc_key_eig = get_local_eig_np(
                        expand, useiss=False)  # 1 x nb_keypoints x 9

                    # print('*****************************************')
                    # print('get local cost time:{}'.format(t2 - t1))
                    #pc_key_feature[i, :, :] = np.squeeze(sess.run(pc_key_eig, feed_dict={pc_pl: pc}))
                    pc_key_feature[j, :, :] = np.squeeze(pc_key_eig)
                else:

                    pc_tile[j, :, :] = np.expand_dims(
                        np.loadtxt(base_path + '/lab' + str(k + 1) +
                                   '/lab_project' + str(i) + '.txt'),
                        axis=0)

                # print('-----------------------------------------')
                # print('one pc cost total:{}second'.format(te-ts))
                # print('----------------------------------------')

        pc_label = np.tile(np.arange(nb_types), (n, 1)).reshape((-1, ),
                                                                order='F')
        train_set_shape = (nb_types * n, 1024, 3)
        train_set_local_shape = (nb_types * n, 102, 9)
        train_label_shape = (nb_types * n, )
    else:
        ply_list = [
            base_path + '/' + i for i in os.listdir(base_path)
            if os.path.splitext(i)[1] == '.ply'
            or os.path.splitext(i)[1] == '.txt'
        ]
        n = len(ply_list)
        less_than_th = []
        for i, j in enumerate(ply_list):
            pc = PointCloud(j)
            if pc.nb_points < 1024:
                less_than_th.append(i)

        n = n - len(less_than_th)
        print("there are: ,", n, ' point clouds with # of points available')
        pc_label = np.arange(n)
        train_set_shape = (n, 1024, 3)
        train_set_local_shape = (n, 102, 9)
        train_label_shape = (n, )

        pc_tile = np.empty(shape=(n, 1024, 3))
        pc_key_feature = np.empty(shape=(n, int(
            1024 * 0.1), 9))  # key feature space, 102=1024*0.1,

        for i, j in enumerate(ply_list):
            if i not in less_than_th:
                print(j)
                start_time = time.clock()
                mypc = PointCloud(j)
                if mypc.nb_points > 1024:
                    mypc.down_sample(number_of_downsample=1024)
                    if normalize:
                        mypc.normalize()
                    expand = np.expand_dims(mypc.position, axis=0)
                    pc_tile[i, :, :] = expand
                    pc_key_eig = get_local_eig_np(expand, useiss=False)
                    end_time = time.clock()
                    compute_time.append([end_time - start_time])
                    if use_key_feature:
                        pc_key_feature[i, :, :] = np.squeeze(pc_key_eig)

    hdf5_file = h5py.File(save_path, mode='a')
    hdf5_file.create_dataset('train_set', train_set_shape,
                             np.float32)  # be careful about the dtype
    hdf5_file.create_dataset('train_labels', train_label_shape, np.uint8)
    hdf5_file.create_dataset('train_set_local', train_set_local_shape,
                             np.float32)
    if shuffle:

        idx = np.arange(np.shape(pc_tile)[0])
        np.random.shuffle(idx)
        pc_tile = pc_tile[idx, ...]
        pc_label = pc_label[idx, ...]
        pc_key_feature = pc_key_feature[idx, ...]

    hdf5_file["train_set"][...] = pc_tile
    hdf5_file["train_labels"][...] = pc_label
    hdf5_file["train_set_local"][...] = pc_key_feature
    hdf5_file.close()
    return compute_time