-
Notifications
You must be signed in to change notification settings - Fork 26
/
predict.py
70 lines (59 loc) · 2.63 KB
/
predict.py
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
import argparse
from parse_config import ConfigParser
import deepmvlm
from utils3d import Utils3D
import os
def process_one_file(config, file_name):
print('Processing ', file_name)
name_lm_vtk = os.path.splitext(file_name)[0] + '_landmarks.vtk'
name_lm_txt = os.path.splitext(file_name)[0] + '_landmarks.txt'
dm = deepmvlm.DeepMVLM(config)
landmarks = dm.predict_one_file(file_name)
dm.write_landmarks_as_vtk_points(landmarks, name_lm_vtk)
dm.write_landmarks_as_text(landmarks, name_lm_txt)
dm.visualise_mesh_and_landmarks(file_name, landmarks)
def process_file_list(config, file_name):
print('Processing filelist ', file_name)
names = []
with open(file_name) as f:
for line in f:
line = (line.strip("/n")).strip("\n")
if len(line) > 4:
names.append(line)
print('Processing ', len(names), ' meshes')
dm = deepmvlm.DeepMVLM(config)
for file_name in names:
print('Processing ', file_name)
name_lm_txt = os.path.splitext(file_name)[0] + '_landmarks.txt'
landmarks = dm.predict_one_file(file_name)
dm.write_landmarks_as_text(landmarks, name_lm_txt)
def process_files_in_dir(config, dir_name):
print('Processing files in ', dir_name)
names = Utils3D.get_mesh_files_in_dir(dir_name)
print('Processing ', len(names), ' meshes')
dm = deepmvlm.DeepMVLM(config)
for file_name in names:
print('Processing ', file_name)
name_lm_txt = os.path.splitext(file_name)[0] + '_landmarks.txt'
landmarks = dm.predict_one_file(file_name)
dm.write_landmarks_as_text(landmarks, name_lm_txt)
def main(config):
name = str(config.name)
if name.lower().endswith(('.obj', '.wrl', '.vtk', '.vtp', '.ply', '.stl')) and os.path.isfile(name):
process_one_file(config, name)
elif name.lower().endswith('.txt') and os.path.isfile(name):
process_file_list(config, name)
elif os.path.isdir(name):
process_files_in_dir(config, name)
else:
print('Cannot process (not a mesh file, a filelist (.txt) or a directory)', name)
if __name__ == '__main__':
args = argparse.ArgumentParser(description='Deep-MVLM')
args.add_argument('-c', '--config', default=None, type=str,
help='config file path (default: None)')
args.add_argument('-d', '--device', default=None, type=str,
help='indices of GPUs to enable (default: all)')
args.add_argument('-n', '--name', default=None, type=str,
help='name of file, filelist (.txt) or directory to be processed')
global_config = ConfigParser(args)
main(global_config)