File size: 2,885 Bytes
91daf98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import os
import argparse
import ntpath
from tqdm import tqdm
import multiprocessing
from pathlib import Path
from glob import glob
import trimesh
from trimesh.sample import sample_surface
from plyfile import PlyData, PlyElement
import numpy as np
def write_ply(points, filename, text=False):
""" input: Nx3, write points to filename as PLY format. """
points = [(points[i,0], points[i,1], points[i,2]) for i in range(points.shape[0])]
vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4')])
el = PlyElement.describe(vertex, 'vertex', comments=['vertices'])
with open(filename, mode='wb') as f:
PlyData([el], text=text).write(f)
def find_files(folder, extension):
return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
class SamplePoints:
"""
Perform sampleing of points.
"""
def __init__(self):
"""
Constructor.
"""
parser = self.get_parser()
self.options = parser.parse_args()
def get_parser(self):
"""
Get parser of tool.
:return: parser
"""
parser = argparse.ArgumentParser(description='Scale a set of meshes stored as OFF files.')
parser.add_argument('--in_dir', type=str, help='Path to input directory.')
parser.add_argument('--out_dir', type=str, help='Path to output directory; files within are overwritten!')
parser.add_argument("--single-file", action='store_true', default=False)
return parser
def run_parallel(self, project_folder):
out_folder = os.path.join(project_folder, self.options.out_dir)
if not os.path.exists(out_folder):
os.makedirs(out_folder)
files = find_files(project_folder, 'final.stl')
for filepath in files:
N_POINTS = 2000
try:
out_mesh = trimesh.load(str(filepath))
out_pc, _ = sample_surface(out_mesh, N_POINTS)
save_path = os.path.join(out_folder, ntpath.basename(filepath)[:-4]+'_pcd.ply')
write_ply(out_pc, save_path)
except Exception as ex:
return project_folder
return
def run(self):
"""
Run simplification.
"""
if self.options.single_file:
self.run_parallel(self.options.in_dir)
else:
project_folders = sorted(glob(self.options.in_dir+'/*/'))
num_cpus = multiprocessing.cpu_count()
convert_iter = multiprocessing.Pool(num_cpus).imap(self.run_parallel, project_folders)
for _ in tqdm(convert_iter, total=len(project_folders)):
pass
if __name__ == '__main__':
app = SamplePoints()
app.run()
|