-
Notifications
You must be signed in to change notification settings - Fork 4
/
generate_vessel_graph.py
128 lines (110 loc) · 5.91 KB
/
generate_vessel_graph.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
import argparse
import csv
from vessel_graph_generation.forest import Forest
from vessel_graph_generation.greenhouse import Greenhouse
from vessel_graph_generation.utilities import prepare_output_dir, read_config
import vessel_graph_generation.tree2img as tree2img
import numpy as np
import nibabel as nib
import os
import yaml
from multiprocessing import cpu_count
import concurrent.futures
import warnings
from rich.console import Group, Console
from rich.live import Live
from rich.progress import Progress, TimeElapsedColumn
from utils.visualizer import DynamicDisplay
group = Group()
def main(config):
# Initialize greenhouse
greenhouse = Greenhouse(config['Greenhouse'])
# Prepare output directory
out_dir = prepare_output_dir(config['output'])
with open(os.path.join(out_dir, 'config.yml'), 'w') as f:
yaml.dump(config, f)
# Initialize forest
arterial_forest = Forest(config['Forest'], greenhouse.d, greenhouse.r, greenhouse.simspace, nerve_center=greenhouse.nerve_center, nerve_radius=greenhouse.nerve_radius)
venous_forest = Forest(config['Forest'], greenhouse.d, greenhouse.r, greenhouse.simspace, arterial=False, nerve_center=greenhouse.nerve_center, nerve_radius=greenhouse.nerve_radius)#, template_forest=arterial_forest)
greenhouse.set_forests(arterial_forest, venous_forest)
# Grow vessel network
greenhouse.develop_forest()
if config["output"]["save_stats"]:
greenhouse.save_stats(out_dir)
volume_dimension = [int(d) for d in greenhouse.simspace.shape*config['output']['image_scale_factor']]
art_edges = [{
"node1": current_node.position,
"node2": current_node.get_proximal_node().position,
"radius": current_node.radius
} for tree in arterial_forest.get_trees() for current_node in tree.get_tree_iterator(exclude_root=True, only_active=False)]
if venous_forest is not None:
ven_edges = [{
"node1": current_node.position,
"node2": current_node.get_proximal_node().position,
"radius": current_node.radius
} for tree in venous_forest.get_trees() for current_node in tree.get_tree_iterator(exclude_root=True, only_active=False)]
# Save vessel graph as csv file
if config['output']['save_trees']:
name = out_dir.split("/")[-1]
filepath = os.path.join(out_dir, name+'.csv')
with open(filepath, 'w+') as file:
writer = csv.writer(file)
writer.writerow(["node1", "node2", "radius"])
for row in art_edges+ven_edges:
writer.writerow([row["node1"], row["node2"], row["radius"]])
radius_list=[]
if config["output"].get("save_3D_volumes"):
art_mat, _ = tree2img.voxelize_forest(art_edges, volume_dimension, radius_list)
ven_mat, _ = tree2img.voxelize_forest(ven_edges, volume_dimension, radius_list)
art_ven_mat_gray = np.maximum(art_mat, ven_mat).astype(np.uint8)
if config["output"]["save_3D_volumes"] == "npy":
np.save(f'{out_dir}/art_ven_img_gray.npy', art_ven_mat_gray)
elif config["output"]["save_3D_volumes"] == "nifti":
nifti = nib.Nifti1Image(art_ven_mat_gray, np.eye(4))
nib.save(nifti, f"{out_dir}/art_ven_img_gray.nii.gz")
if config["output"]["save_2D_image"]:
radius_list=[]
image_res = [*volume_dimension]
del image_res[config["output"]["proj_axis"]]
art_mat,_ = tree2img.rasterize_forest(art_edges, image_res, MIP_axis=config["output"]["proj_axis"], radius_list=radius_list)
ven_mat,_ = tree2img.rasterize_forest(ven_edges, image_res, MIP_axis=config["output"]["proj_axis"], radius_list=radius_list)
art_ven_mat_gray = np.maximum(art_mat, ven_mat).astype(np.uint8)
tree2img.save_2d_img(art_ven_mat_gray, out_dir, "art_ven_img_gray")
if config["output"]["save_stats"]:
tree2img.plot_vessel_radii(out_dir, radius_list)
if __name__ == '__main__':
# Parse input arguments
parser = argparse.ArgumentParser(description='')
parser.add_argument('--config_file', type=str, required=True)
parser.add_argument('--num_samples', type=int, default=1)
parser.add_argument('--debug', action="store_true")
parser.add_argument('--threads', help="Number of parallel threads. By default all available threads but one are used.", type=int, default=-1)
args = parser.parse_args()
if args.debug:
warnings.filterwarnings('error')
# Read config file
assert os.path.isfile(args.config_file), f"Error: Your provided config path {args.config_file} does not exist!"
config = read_config(args.config_file)
assert config['output'].get('save_3D_volumes') in [None, 'npy', 'nifti'], f"Your provided option {config['output'].get('save_3D_volumes')} for 'save_3D_volumes' does not exist. Choose one of 'null', 'npy' or 'nifti'."
if args.threads == -1:
# If no argument is provided, use all available threads but one
cpus = cpu_count()
threads = min(cpus-1 if cpus>1 else 1,args.num_samples)
else:
threads=args.threads
with Live(group, console=Console(force_terminal=True), refresh_per_second=10):
progress = Progress(*Progress.get_default_columns(), TimeElapsedColumn())
progress.add_task(f"Generating {args.num_samples} vessel graphs:", total=args.num_samples)
with DynamicDisplay(group, progress):
if threads>1:
# Multi processing
with concurrent.futures.ProcessPoolExecutor(max_workers=threads) as executor:
future_dict = {executor.submit(main, config): i for i in range(args.num_samples)}
for future in concurrent.futures.as_completed(future_dict):
i = future_dict[future]
progress.advance(task_id=0)
else:
# Single processing
for i in range(args.num_samples):
main(config)
progress.advance(task_id=0)