|
- # Copyright (c) Nanjing University, Vision Lab.
- # Last update:
- # 2020.11.26
-
- import os
- import numpy as np
- import sys
- import time
- import random
- # sys.path.append('.')
- from dataprocess.inout_points import load_ply_data, write_ply_data, \
- load_points, save_points, \
- points2voxels, voxels2points, \
- select_voxels
-
- def preprocess(input_file, scale, cube_size, min_num): #min_num=64
- """Scaling, Partition & Voxelization.
- Input: .ply file and arguments for pre-process.
- Output: partitioned cubes, cube positions, and number of points in each cube.
- """
- prefix=input_file.split('/')[-1].split('_')[0]+str(random.randint(1,100)) #redandblack2
- 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 = prefix+'downscaling.ply'
- write_ply_data(scaling_file, pc_down) #写入xyz到ply
- 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)))
- if scale != 1:
- os.system("rm "+scaling_file)
-
- # voxelization.
- start = time.time()
- cubes = points2voxels(partitioned_points, cube_size)
- points_numbers = np.sum(cubes, axis=(1,2,3,4)).astype(np.uint16) #每个cube里有多少个点
- print("Voxelization: {}s".format(round(time.time()-start, 4)))
-
- print('cubes shape: {}'.format(cubes.shape)) #(189, 64, 64, 64, 1)
- print('points numbers (sum/mean/max/min): {} {} {} {}'.format(
- points_numbers.sum(), round(points_numbers.mean()), points_numbers.max(), points_numbers.min())) #757478 4008.0 13906 68
-
- return cubes, cube_positions, points_numbers
-
- def postprocess(output_file, cubes, points_numbers, cube_positions, scale, cube_size, rho, fixed_thres=None):
- """Classify voxels to occupied or free, then extract points and write to file.
- Input: deocded cubes, cube positions, points numbers, cube size and rho=ouput numbers/input numbers.
- """
- prefix=output_file.split('/')[-1].split('_')[0]+str(random.randint(1,100))
- print('===== Post process =====')
- # Classify.
- start = time.time()
- output = select_voxels(cubes, points_numbers, rho, fixed_thres=fixed_thres) #output是针对每个cube的mask,1为occupied
-
- # Extract points.
- #points = voxels2points(output.numpy())
- points = voxels2points(output)
- print("Classify and extract points: {}s".format(round(time.time()-start, 4)))
-
- # scaling (optional)
- start = time.time()
- if scale == 1:
- save_points(points, cube_positions, output_file, cube_size)
- else:
- scaling_output_file = prefix+'downsampling_rec.ply'
- save_points(points, cube_positions, scaling_output_file, cube_size)
- pc = load_ply_data(scaling_output_file)
- pc_up = pc.astype('float32') * float(1/scale)
- write_ply_data(output_file, pc_up)
- os.system("rm "+scaling_output_file)
- print("Write point cloud to {}: {}s".format(output_file, round(time.time()-start, 4)))
-
- return
-
- if __name__=='__main__':
- import argparse
- parser = argparse.ArgumentParser(
- formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- parser.add_argument(
- "--input", type=str, default="testdata/8iVFB/longdress_vox10_1300.ply", dest="input")
- parser.add_argument(
- "--output", type=str, default="rec.ply", dest="output")
- parser.add_argument(
- "--scale", type=float, default=1.0, dest="scale")
- parser.add_argument(
- "--cube_size", type=int, default=64, dest="cube_size")
- parser.add_argument(
- "--min_num", type=int, default=20, dest="min_num")
- args = parser.parse_args()
- print(args)
-
- cubes, cube_positions, points_numbers = preprocess(args.input,
- args.scale,
- args.cube_size,
- args.min_num)
- print("===================\n", cubes.shape, np.sum(cubes), '\n', cube_positions.shape, '\n', points_numbers.shape)
- postprocess(args.output, cubes, points_numbers, cube_positions, args.scale, args.cube_size, rho=1.0)
- os.system("myutils/pc_error_d" \
- + ' -a ' + args.input + ' -b ' + args.output + " -r 1023")
- os.system("rm "+args.output)
|