def generate_dataset(INPUT_DIR, OUTPUT_DIR, DATA_NUM, cube_size=64):
# read file
plydirs = glob.glob(INPUT_DIR + '*.ply')
random.shuffle(plydirs)
random.shuffle(plydirs)
print('ply file direction list length:', len(plydirs))
num_data = 0
for filename in plydirs:
set_points, _ = load_points(filename, cube_size=cube_size, min_num=20)
random.shuffle(set_points)
# only half of the points will be used
for i, points in enumerate(set_points):
points = points.astype("uint8")
points_dir = os.path.join(
OUTPUT_DIR,
filename.split('/')[-1].split('.')[0] + '_' + str(i) + 'n.h5')
with h5py.File(points_dir, 'w') as h:
# print(set_points[i].dtype)
h.create_dataset('data', data=points, shape=points.shape)
num_data += 1
if num_data % 100 == 0:
print(num_data)
if num_data > DATA_NUM:
break
return
def preprocess(input_file, scale, cube_size, min_num):
"""Scaling, Partition & Voxelization.
Input: .ply file and arguments for pre-process.
Output: partitioned cubes, cube positions, and number of points in each cube.
"""
print('===== Preprocess =====')
# scaling (optional)
start = time.time()
if scale == 1:
scaling_file = input_file
else:
pc = load_ply_data(input_file)
pc_down = np.round(pc.astype('float32') * scale)
pc_down = np.unique(pc_down, axis=0)# remove duplicated points
scaling_file = './downscaling.ply'
write_ply_data(scaling_file, pc_down)
print("Scaling: {}s".format(round(time.time()-start, 4)))
# partition.
start = time.time()
partitioned_points, cube_positions = load_points(scaling_file, cube_size, min_num)
print("Partition: {}s".format(round(time.time()-start, 4)))
# voxelization.
start = time.time()
cubes = points2voxels(partitioned_points, cube_size)
points_numbers = np.sum(cubes, axis=(1,2,3,4)).astype(np.uint16)
print("Voxelization: {}s".format(round(time.time()-start, 4)))
print('cubes shape: {}'.format(cubes.shape))
print('points numbers (sum/mean/max/min): {} {} {} {}'.format(
points_numbers.sum(), round(points_numbers.mean()), points_numbers.max(), points_numbers.min()))
return cubes, cube_positions, points_numbers