|
import numpy as np |
|
import pyvista as pv |
|
import argparse |
|
import os |
|
import glob |
|
import trimesh |
|
|
|
|
|
def parse_command_line(): |
|
print('---'*10) |
|
print('Parsing Command Line Arguments') |
|
parser = argparse.ArgumentParser(description='Defacing protocol') |
|
parser.add_argument('-bp', metavar='base path', type=str, |
|
help="Absolute path of the base directory") |
|
parser.add_argument('-gp', metavar='ground truth path', type=str, |
|
help="Relative path of the ground truth model") |
|
parser.add_argument('-pp', metavar='prediction path', type=str, |
|
help="Relative path of the prediction model") |
|
argv = parser.parse_args() |
|
return argv |
|
|
|
|
|
def distanceVertex2Mesh(mesh, vertex): |
|
faces_as_array = mesh.faces.reshape((mesh.n_faces, 4))[:, 1:] |
|
mesh_box = trimesh.Trimesh(vertices=mesh.points, |
|
faces=faces_as_array) |
|
cp, cd, ci = trimesh.proximity.closest_point(mesh_box, vertex) |
|
return cd |
|
|
|
|
|
def main(): |
|
args = parse_command_line() |
|
base = args.bp |
|
gt_path = args.gp |
|
pred_path = args.pp |
|
output_dir = os.path.join(base, 'output') |
|
try: |
|
os.mkdir(output_dir) |
|
except: |
|
print(f'{output_dir} already exists') |
|
|
|
for i in glob.glob(os.path.join(base, gt_path) + '/*.vtk'): |
|
filename = os.path.basename(i).split('.')[0] |
|
|
|
|
|
|
|
output_sub_dir = os.path.join( |
|
base, 'output', filename) |
|
try: |
|
os.mkdir(output_sub_dir) |
|
except: |
|
print(f'{output_sub_dir} already exists') |
|
|
|
gt_mesh = pv.read(i) |
|
pred_mesh = pv.read(os.path.join( |
|
base, pred_path, filename + '.vtk')) |
|
pred_vertices = np.array(pred_mesh.points) |
|
cd = distanceVertex2Mesh(gt_mesh, pred_vertices) |
|
pred_mesh['dist'] = cd |
|
pred_mesh.save(os.path.join(output_sub_dir, filename + '.vtk')) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
