-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsw2stl.py
173 lines (136 loc) · 6.71 KB
/
sw2stl.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
#%%
import argparse
import h5py
import numpy as np
import json
import wget
import re
import urllib.request
import trimesh
import os
from tqdm import tqdm
def get_sw_json(json_path):
"""Extract .sw URL from the .json file"""
with open(json_path) as json_file:
sw_json = json.load(json_file)
sw_url = sw_json['spacewalk']['url']
return sw_url
def download_sw_file(local_filename, sw_url):
"""Download the .sw file from the SpaceWalk URL"""
urllib.request.urlretrieve(sw_url, local_filename)
def create_sphere(center, radius):
"""Create a sphere at the given center with the given radius"""
sphere = trimesh.creation.icosphere(subdivisions=2, radius=radius)
sphere.apply_translation(center)
return sphere
def create_cylinder(start_point, end_point, radius):
"""Create a cylinder with the given start and end points and radius"""
height = np.linalg.norm(np.array(end_point) - np.array(start_point))
direction = np.array(end_point) - np.array(start_point)
direction /= height
cylinder = trimesh.creation.cylinder(radius=radius, height=height, sections=16)
# Calculate rotation matrix
z_axis = np.array([0, 0, 1])
v = np.cross(z_axis, direction)
c = np.dot(z_axis, direction)
k = 1 / (1 + c)
rotation = np.array([
[v[0] * v[0] * k + c, v[0] * v[1] * k - v[2], v[0] * v[2] * k + v[1]],
[v[1] * v[0] * k + v[2], v[1] * v[1] * k + c, v[1] * v[2] * k - v[0]],
[v[2] * v[0] * k - v[1], v[2] * v[1] * k + v[0], v[2] * v[2] * k + c]
])
# Convert rotation matrix to 4x4 transformation matrix
transformation = np.eye(4)
transformation[:3, :3] = rotation
cylinder.apply_transform(transformation)
cylinder.apply_translation((start_point + end_point) / 2.0)
return cylinder
def normalize_vertices(vertices):
"""Normalize vertices to fit within a unit cube centered at the origin"""
min_coords = np.min(vertices, axis=0)
max_coords = np.max(vertices, axis=0)
center = (min_coords + max_coords) / 2
vertices -= center # Center the vertices
scale = np.max(max_coords - min_coords) # Find the scale based on the largest dimension
vertices /= scale # Normalize to fit within a unit cube
return vertices, scale
def convert_to_stl_and_offset(hdf5_file_path, ensemble_group, output_stl, time_point):
"""Convert the spatial position data from the hdf5 file to an .stl file with spheres and cylinders"""
# Open the HDF5 file and extract the spatial position data (vertex coordinates)
with h5py.File(hdf5_file_path, 'r') as swbf:
dataset = swbf[ensemble_group]
data = dataset['spatial_position']
sp = data[time_point]
vertices = np.array(sp, dtype=np.float32)
# Process and refine vertices (remove NaNs)
vertices = vertices[~np.isnan(vertices).any(axis=1)]
# Normalize the vertices
vertices, scale = normalize_vertices(vertices)
# Set dynamic sphere and cylinder sizes based on the normalized scale
sphere_radius = 0.01 # Small radius for spheres
cylinder_radius = 0.005 # Small radius for cylinders
# Initialize an empty mesh
combined_mesh = None
for i in tqdm(range(len(vertices))):
sphere = create_sphere(vertices[i], sphere_radius)
if combined_mesh is None:
combined_mesh = sphere
else:
combined_mesh = trimesh.boolean.union([combined_mesh, sphere])
if i < len(vertices) - 1:
cylinder = create_cylinder(vertices[i], vertices[i + 1], cylinder_radius)
combined_mesh = trimesh.boolean.union([combined_mesh, cylinder])
# Ensure the combined mesh is watertight (no holes or self-intersections)
if not combined_mesh.is_watertight:
combined_mesh = combined_mesh.fill_holes()
# Save the final combined mesh as an .stl file
combined_mesh.export(output_stl, file_type='stl')
print(f"STL file saved as {output_stl}")
def main():
"""Define arguments, determine .json or .sw to proceed, and call the convert_to_stl_and_offset function"""
parser = argparse.ArgumentParser(description='Convert dataset to STL and apply offset.')
parser.add_argument('-f','--file', required=True, help='Path to the file')
parser.add_argument('-j','--json_file', required=False, help='Path to the json file')
parser.add_argument('-hd','--hdf5_file', required=False, help='Path to the hdf5 file')
parser.add_argument('-e','--ensemble_group', required=True, help='Title of the ensemble group (cap senstive)')
parser.add_argument('-t','--time_point', required=True, help='''Time point number.
Note: The time point should be in the format "t_#" where # is the time point number.
Also consider model 1 = t_0, therefore, may need to -1 the # before plugging in to ensure correct model''')
parser.add_argument('-o','--output_stl', default=None, help='Output STL file title (Optional)')
args = parser.parse_args()
hdf5_file_path = None
if '.json' in args.file:
args.json_file = args.file
sw_url = get_sw_json(args.json_file)
if 'dropbox' in sw_url:
if 'dl=0' in sw_url:
sw_url = sw_url.replace('dl=0', 'dl=1')
hdf5_file_path = wget.download(sw_url)
else:
pattern = r'/([^/]+)$'
match = re.search(pattern, sw_url)
if match:
extracted_part_url = match.group(1)
else:
raise ValueError("No match found in the URL")
extracted_part_url = extracted_part_url.replace('.sw', '')
hdf5_file_path = f'{extracted_part_url}.sw'
if hdf5_file_path in os.listdir():
print(f"{hdf5_file_path} already exists in the current directory")
else:
download_sw_file(hdf5_file_path, sw_url)
elif '.sw' in args.file:
args.hdf5_file = args.file
hdf5_file_path = args.hdf5_file
else:
raise ValueError("Please provide the path a file that is either a .json or .sw")
if args.output_stl is None:
args.output_stl = f'{args.ensemble_group}_{args.time_point}.stl'
convert_to_stl_and_offset(hdf5_file_path, args.ensemble_group, args.output_stl, args.time_point)
if __name__ == '__main__':
main()
# example parameters
# Mammoth: -j = ./mammoth-spacewalk-session.json, -e = replica01_chr10_Direct_Inversion, -t = t_141
# HCT116 6h_auxin chr21:28-30: -j = ./HCT116-spacewalk-session.json, -e = HCT116, -t = t_14
# HCT116 6h_auxin chr21:28-30: -j = ./HCT116-t-36-spacewalk-session.json, -e = HCT116, -t = t_36
# K562 chr21:28-30: -j = ./K562-t-1008-spacewalk-session.json, -e = K562, -t = t_1008