-
Notifications
You must be signed in to change notification settings - Fork 41
/
run_network.py
executable file
·240 lines (202 loc) · 8.82 KB
/
run_network.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
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
#!/usr/bin/env python3
import os
import sys
import argparse
import numpy as np
import re
from glob import glob
import time
import importlib
import json
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'datasets'))
from physics_data_helper import numpy_from_bgeo, write_bgeo_from_numpy
from create_physics_scenes import obj_surface_to_particles, obj_volume_to_particles
import open3d as o3d
def write_particles(path_without_ext, pos, vel=None, options=None):
"""Writes the particles as point cloud ply.
Optionally writes particles as bgeo which also supports velocities.
"""
arrs = {'pos': pos}
if not vel is None:
arrs['vel'] = vel
np.savez(path_without_ext + '.npz', **arrs)
if options and options.write_ply:
pcd = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pos))
o3d.io.write_point_cloud(path_without_ext + '.ply', pcd)
if options and options.write_bgeo:
write_bgeo_from_numpy(path_without_ext + '.bgeo', pos, vel)
def run_sim_tf(trainscript_module, weights_path, scene, num_steps, output_dir,
options):
# init the network
model = trainscript_module.create_model()
model.init()
model.load_weights(weights_path, by_name=True)
# prepare static particles
walls = []
for x in scene['walls']:
points, normals = obj_surface_to_particles(x['path'])
if 'invert_normals' in x and x['invert_normals']:
normals = -normals
points += np.asarray([x['translation']], dtype=np.float32)
walls.append((points, normals))
box = np.concatenate([x[0] for x in walls], axis=0)
box_normals = np.concatenate([x[1] for x in walls], axis=0)
# export static particles
write_particles(os.path.join(output_dir, 'box'), box, box_normals, options)
# compute lowest point for removing out of bounds particles
min_y = np.min(box[:, 1]) - 0.05 * (np.max(box[:, 1]) - np.min(box[:, 1]))
# prepare fluids
fluids = []
for x in scene['fluids']:
points = obj_volume_to_particles(x['path'])[0]
points += np.asarray([x['translation']], dtype=np.float32)
velocities = np.empty_like(points)
velocities[:, 0] = x['velocity'][0]
velocities[:, 1] = x['velocity'][1]
velocities[:, 2] = x['velocity'][2]
range_ = range(x['start'], x['stop'], x['step'])
fluids.append((points, velocities, range_))
pos = np.empty(shape=(0, 3), dtype=np.float32)
vel = np.empty_like(pos)
for step in range(num_steps):
# add from fluids to pos vel arrays
for points, velocities, range_ in fluids:
if step in range_: # check if we have to add the fluid at this point in time
pos = np.concatenate([pos, points], axis=0)
vel = np.concatenate([vel, velocities], axis=0)
if pos.shape[0]:
fluid_output_path = os.path.join(output_dir,
'fluid_{0:04d}'.format(step))
if isinstance(pos, np.ndarray):
write_particles(fluid_output_path, pos, vel, options)
else:
write_particles(fluid_output_path, pos.numpy(), vel.numpy(),
options)
inputs = (pos, vel, None, box, box_normals)
pos, vel = model(inputs)
# remove out of bounds particles
if step % 10 == 0:
print(step, 'num particles', pos.shape[0])
mask = pos[:, 1] > min_y
if np.count_nonzero(mask) < pos.shape[0]:
pos = pos[mask]
vel = vel[mask]
def run_sim_torch(trainscript_module, weights_path, scene, num_steps,
output_dir, options):
import torch
device = torch.device(options.device)
# init the network
model = trainscript_module.create_model()
weights = torch.load(weights_path)
model.load_state_dict(weights)
model.to(device)
model.requires_grad_(False)
# prepare static particles
walls = []
for x in scene['walls']:
points, normals = obj_surface_to_particles(x['path'])
if 'invert_normals' in x and x['invert_normals']:
normals = -normals
points += np.asarray([x['translation']], dtype=np.float32)
walls.append((points, normals))
box = np.concatenate([x[0] for x in walls], axis=0)
box_normals = np.concatenate([x[1] for x in walls], axis=0)
# export static particles
write_particles(os.path.join(output_dir, 'box'), box, box_normals, options)
# compute lowest point for removing out of bounds particles
min_y = np.min(box[:, 1]) - 0.05 * (np.max(box[:, 1]) - np.min(box[:, 1]))
box = torch.from_numpy(box).to(device)
box_normals = torch.from_numpy(box_normals).to(device)
# prepare fluids
fluids = []
for x in scene['fluids']:
points = obj_volume_to_particles(x['path'])[0]
points += np.asarray([x['translation']], dtype=np.float32)
velocities = np.empty_like(points)
velocities[:, 0] = x['velocity'][0]
velocities[:, 1] = x['velocity'][1]
velocities[:, 2] = x['velocity'][2]
range_ = range(x['start'], x['stop'], x['step'])
fluids.append(
(points.astype(np.float32), velocities.astype(np.float32), range_))
pos = np.empty(shape=(0, 3), dtype=np.float32)
vel = np.empty_like(pos)
for step in range(num_steps):
# add from fluids to pos vel arrays
for points, velocities, range_ in fluids:
if step in range_: # check if we have to add the fluid at this point in time
pos = np.concatenate([pos, points], axis=0)
vel = np.concatenate([vel, velocities], axis=0)
if pos.shape[0]:
fluid_output_path = os.path.join(output_dir,
'fluid_{0:04d}'.format(step))
if isinstance(pos, np.ndarray):
write_particles(fluid_output_path, pos, vel, options)
else:
write_particles(fluid_output_path, pos.numpy(), vel.numpy(),
options)
inputs = (torch.from_numpy(pos).to(device),
torch.from_numpy(vel).to(device), None, box, box_normals)
pos, vel = model(inputs)
pos = pos.cpu().numpy()
vel = vel.cpu().numpy()
# remove out of bounds particles
if step % 10 == 0:
print(step, 'num particles', pos.shape[0])
mask = pos[:, 1] > min_y
if np.count_nonzero(mask) < pos.shape[0]:
pos = pos[mask]
vel = vel[mask]
def main():
parser = argparse.ArgumentParser(
description=
"Runs a fluid network on the given scene and saves the particle positions as npz sequence",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("trainscript",
type=str,
help="The python training script.")
parser.add_argument(
"--weights",
type=str,
required=True,
help=
"The path to the .h5 network weights file for tensorflow ot the .pt weights file for torch."
)
parser.add_argument("--num_steps",
type=int,
default=250,
help="The number of simulation steps. Default is 250.")
parser.add_argument("--scene",
type=str,
required=True,
help="A json file which describes the scene.")
parser.add_argument("--output",
type=str,
required=True,
help="The output directory for the particle data.")
parser.add_argument("--write-ply",
action='store_true',
help="Export particle data also as .ply sequence")
parser.add_argument("--write-bgeo",
action='store_true',
help="Export particle data also as .bgeo sequence")
parser.add_argument("--device",
type=str,
default='cuda',
help="The device to use. Applies only for torch.")
args = parser.parse_args()
print(args)
module_name = os.path.splitext(os.path.basename(args.trainscript))[0]
sys.path.append('.')
trainscript_module = importlib.import_module(module_name)
with open(args.scene, 'r') as f:
scene = json.load(f)
os.makedirs(args.output)
if args.weights.endswith('.h5'):
return run_sim_tf(trainscript_module, args.weights, scene,
args.num_steps, args.output, args)
elif args.weights.endswith('.pt'):
return run_sim_torch(trainscript_module, args.weights, scene,
args.num_steps, args.output, args)
if __name__ == '__main__':
sys.exit(main())