Function load_completion_data

utils/data_provider.py:10–58

(path, batch_size, encode, npoint=2048, split='split_pcl2pcl.txt')

import glob

def load_completion_data(path, batch_size, encode, npoint=2048, split='split_pcl2pcl.txt'):
    save_path_train = os.path.join(path,"partial")
    save_path_gt = os.path.join(path,"complete")
    f_lidar = glob.glob(os.path.join(save_path_train, '*.npy'))
    new_dataset = []
    gt_dataset = []
    test_dataset = []
    test_dataset_gt = []

    a = np.loadtxt('./dataset/3depn/'+split, str)

    b = []
    for i in a:
        if int(i[:8]) == int(encode):
            i = i[9:]
            b.append(i)
    for i in f_lidar:
        raw_lidar = np.load(i)
        file = i.split('/')[-1].split('.')[0][:-5]+".npy"
        gt_lidar = np.load(os.path.join(save_path_gt,file))
        if file[:-4] in b:
            test_dataset.append(raw_lidar)
            test_dataset_gt.append(gt_lidar)
        else:
            new_dataset.append(raw_lidar)
            gt_dataset.append(gt_lidar)
    new_dataset = np.array(new_dataset)
    gt_dataset = np.array(gt_dataset)
    test_dataset = np.array(test_dataset)
    test_dataset_gt = np.array(test_dataset_gt)
    batch_dataset = []
    batch_dataset_gt = []
    test_batch_dataset = []
    test_batch_dataset_gt = []
    i=0
    while i+batch_size<=new_dataset.shape[0]:
        batch_dataset.append(new_dataset[i:i+batch_size])
        batch_dataset_gt.append(gt_dataset[i:i+batch_size])
        i = i + batch_size
    i=0
    while i+batch_size<=test_dataset.shape[0]:
        test_batch_dataset.append(test_dataset[i:i+batch_size])
        test_batch_dataset_gt.append(test_dataset_gt[i:i+batch_size])
        i = i + batch_size
    batch_dataset = np.array(batch_dataset)
    batch_dataset_gt = np.array(batch_dataset_gt)
    test_batch_dataset = np.array(test_batch_dataset)
    test_batch_dataset_gt = np.array(test_batch_dataset_gt)
    return batch_dataset, batch_dataset_gt, test_batch_dataset, test_batch_dataset_gt


def rotate_point_cloud_and_gt(batch_data,batch_gt=None):