-
Notifications
You must be signed in to change notification settings - Fork 3
/
vis_transport.py
214 lines (177 loc) · 8.78 KB
/
vis_transport.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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
#!/usr/bin/env python3
# Copyright 2004-present Facebook. All Rights Reserved.
import argparse
import json
import logging
import os
import random
import time
import torch
import naisr
import naisr.workspace as ws
from utils import cond_mkdir
import pandas as pd
from naisr import loss_pointcloud_sdf
import naisr_meshing
from naisr import *
from visualizer import plotter_evolution, plotter_evolution_comp
def get_covariates_difference(attributes, start_covariates, diff_names='all'):
differences = {}
for name, value in attributes.items():
if diff_names == 'all' or name in diff_names:
differences[name] = attributes[name] - start_covariates[name]
else:
differences[name] = 0
return differences
def get_transported_covariates(inferred_attributes, difference_covarites):
transported_covarites = {}
for name, value in inferred_attributes.items():
transported_covarites[name] = inferred_attributes[name] + difference_covarites[name].to(inferred_attributes[name].device)
return transported_covarites
def transport_covariates(attributes, inferred_attributes, start_covariates, diff_names='all'):
difference_covarites = get_covariates_difference(attributes, start_covariates, diff_names=diff_names)
transported_covarites = get_transported_covariates(inferred_attributes, difference_covarites)
return transported_covarites
if __name__ == "__main__":
arg_parser = argparse.ArgumentParser(
description="Use a trained DeepSDF decoder to reconstruct a shape given SDF "
+ "samples."
)
arg_parser.add_argument(
"--experiment",
"-e",
dest="experiment_directory",
default='examples/pediatric_airway/naivf_deepnaigsr.json',
required=False,
help="The experiment directory which includes specifications and saved model "
+ "files to use for reconstruction",
)
arg_parser.add_argument(
"--checkpoint",
"-c",
dest="checkpoint",
default="epoch_3000",
help="The checkpoint weights to use. This can be a number indicated an epoch "
+ "or 'latest' for the latest weights (this is the default)",
)
arg_parser.add_argument(
"--iters",
dest="iterations",
default=800,
help="The number of iterations of latent code optimization to perform.",
)
arg_parser.add_argument(
"--skip",
dest="skip",
action="store_true",
help="Skip meshes which have already been reconstructed.",
)
naisr.add_common_args(arg_parser)
args = arg_parser.parse_args()
naisr.configure_logging(args)
def empirical_stat(latent_vecs, indices):
lat_mat = torch.zeros(0).to(device)
for ind in indices:
lat_mat = torch.cat([lat_mat, latent_vecs[ind]], 0)
mean = torch.mean(lat_mat, 0)
var = torch.var(lat_mat, 0)
return mean, var
specs_filename = args.experiment_directory
if not os.path.isfile(specs_filename):
raise Exception(
'The experiment directory does not include specifications file "specs.json"'
)
specs = json.load(open(specs_filename))
device = specs['Device']
root_path = os.path.join(specs['LoggingRoot'], specs['ExperimentName'])
latent_size = specs["CodeLength"]
decoder = eval(specs['Network'])(
template_attributes=specs['TemplateAttributes'],
in_features=specs['InFeatures'],
hidden_features=specs['HiddenFeatures'],
hidden_layers=specs['HidenLayers'],
out_features=specs['OutFeatures'],
device=specs['Device'],
backbone=specs['Backbone'],
outermost_linear=False,
pos_enc=specs['PosEnc'],
latent_size=specs["CodeLength"])
#decoder = torch.nn.DataParallel(decoder)
saved_model_state = torch.load(
os.path.join(
root_path, ws.model_params_subdir, args.checkpoint + ".pth"
), map_location=torch.device(device)
)
saved_model_epoch = saved_model_state["epoch"]
decoder.load_state_dict(saved_model_state["model_state_dict"])
decoder = decoder.to(device)
logging.info(decoder)
err_sum = 0.0
repeat = 1
save_latvec_only = False
rerun = 0
transport_dir = os.path.join(root_path, ws.transport_subdir)
if not os.path.isdir(transport_dir):
os.makedirs(transport_dir)
transport_meshes_dir = os.path.join(transport_dir, ws.transport_meshes_subdir)
if not os.path.isdir(transport_meshes_dir):
os.makedirs(transport_meshes_dir)
transport_codes_dir = os.path.join(transport_dir, ws.transport_codes_subdir)
if not os.path.isdir(transport_codes_dir):
os.makedirs(transport_codes_dir)
#cases = naisr.get_youngest_ids(specs["Split"], split='test')
list_patient_scans = naisr.get_patients_for_transport(specs["DataSource"], specs["Split"], split='test_multiple')
training_cases = naisr.get_ids(specs["Split"], split='train')
import pandas as pd
df_data = pd.read_csv(specs["DataSource"], header=0)
list_metrics = []
for transport_name in ['all', 'age', 'weight']:
transport_covariate_meshes_dir = os.path.join(transport_meshes_dir, transport_name)
if not os.path.isdir(transport_covariate_meshes_dir):
os.makedirs(transport_covariate_meshes_dir)
for current_patient in list_patient_scans:
list_color = []
list_pred_shapepaths = []
list_gt_shapepaths = []
list_text = []
test_idx = current_patient['youngest_scan']
transport_covariate_meshes_dir_subj = os.path.join(transport_covariate_meshes_dir, str(test_idx))
start_arr_samples, start_attributes, start_gt = naisr.get_data_for_id(test_idx,
df_data,
training_cases,
specs["Attributes"])
path_3dshape = os.path.join(transport_covariate_meshes_dir_subj, "surface.stl")
list_pred_shapepaths.append(path_3dshape)
list_gt_shapepaths.append(start_gt['pvgt_path'][0])
list_text.append(start_gt['covariates'])
batch_lat = load_transport_vectors(transport_codes_dir, start_gt['id'][0], device)
infered_start_attributes = load_inferred_covariates(transport_codes_dir, start_gt['id'][0], device)
current_color = naisr_meshing.revert_points_to_template(decoder,
batch_lat,
infered_start_attributes,
transport_covariate_meshes_dir_subj,
device)
list_color.append(current_color)
# transporting to other covariates
other_scans= current_patient['other_scans']
for ith_scan_to_transp in other_scans:
transport_covariate_meshes_dir_subj = os.path.join(transport_covariate_meshes_dir, str(ith_scan_to_transp))
arr_samples, attributes, gt = naisr.get_data_for_id(ith_scan_to_transp,
df_data,
training_cases,
specs["Attributes"])
path_3dshape = os.path.join(transport_covariate_meshes_dir_subj, "surface.stl")
list_pred_shapepaths.append(path_3dshape)
list_gt_shapepaths.append(gt['pvgt_path'][0])
list_text.append(gt['covariates'])
new_covariates = transport_covariates(attributes, infered_start_attributes, start_attributes, diff_names='all')
current_color = naisr_meshing.revert_points_to_template(decoder,
batch_lat,
new_covariates,
transport_covariate_meshes_dir_subj,
device)
list_color.append(current_color)
savepath = os.path.join(transport_covariate_meshes_dir, str(test_idx))
print(list_gt_shapepaths)
plotter_evolution_comp(list_pred_shapepaths, list_gt_shapepaths, savepath, list_text=list_text, print_on_figure=False,)
plotter_evolution(list_pred_shapepaths, savepath, list_text=list_text, list_colors=list_color)