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 create_fluid_data(output_dir, seed, options):
"""Creates a random scene for a specific seed and runs the simulator"""
np.random.seed(seed)
bounding_boxes = sorted(
glob(os.path.join(SCRIPT_DIR, 'models', 'Box*.obj')))
# override bounding boxes
if options.default_box:
bounding_boxes = [os.path.join(SCRIPT_DIR, 'models', 'Box.obj')]
fluid_shapes = sorted(
glob(os.path.join(SCRIPT_DIR, 'models', 'Fluid*.obj')))
rigid_shapes = sorted(
glob(os.path.join(SCRIPT_DIR, 'models', 'RigidBody*.obj')))
num_objects = np.random.choice([1, 2, 3])
# override the number of objects to generate
if options.num_objects > 0: num_objects = options.num_objects
# num_objects = random.choice([1])
print('num_objects', num_objects)
# create fluids and place them randomly
def create_fluid_object():
fluid_obj = np.random.choice(fluid_shapes)
fluid = obj_volume_to_particles(fluid_obj,
scale=np.random.uniform(0.5, 1.5))[0]
R = random_rotation_matrix(1.0)
fluid = fluid @ R
fluid_rasterized = rasterize_points(fluid, 2.01 * PARTICLE_RADIUS,
PARTICLE_RADIUS)
selected_pos = find_valid_fluid_start_positions(
bb_rasterized, fluid_rasterized)
fluid_pos = selected_pos - fluid_rasterized[0] * fluid_rasterized[1]
fluid += fluid_pos
fluid_vel = np.zeros_like(fluid)
max_vel = MAX_FLUID_START_VELOCITY_XZ
fluid_vel[:, 0] = np.random.uniform(-max_vel, max_vel)
fluid_vel[:, 2] = np.random.uniform(-max_vel, max_vel)
max_vel = MAX_FLUID_START_VELOCITY_Y
fluid_vel[:, 1] = np.random.uniform(-max_vel, max_vel)
density = np.random.uniform(500, 2000)
viscosity = np.random.exponential(scale=1 / 20) + 0.01
if options.uniform_viscosity:
viscosity = np.random.uniform(0.01, 0.3)
elif options.log10_uniform_viscosity:
viscosity = 0.01 * 10**np.random.uniform(0.0, 1.5)
if options.default_density: density = 1000
if options.default_viscosity: viscosity = 0.01
return {
'type': 'fluid',
'positions': fluid,
'velocities': fluid_vel,
'density': density,
'viscosity': viscosity,
}
scene_is_valid = False
for create_scene_i in range(100):
if scene_is_valid:
break
# select random bounding box
bb_obj = np.random.choice(bounding_boxes)
# convert bounding box to particles
bb, bb_normals = obj_surface_to_particles(bb_obj)
bb_vol = obj_volume_to_particles(bb_obj)[0]
# rasterize free volume
bb_rasterized = rasterize_points(np.concatenate([bb_vol, bb], axis=0),
2.01 * PARTICLE_RADIUS,
PARTICLE_RADIUS)
bb_rasterized = bb_rasterized[0], bb_rasterized[1], binary_erosion(
bb_rasterized[2], structure=np.ones((3, 3, 3)), iterations=3)
objects = []
create_fn_list = [create_fluid_object]
for object_i in range(num_objects):
create_fn = np.random.choice(create_fn_list)
create_success = False
for i in range(10):
if create_success:
break
try:
obj = create_fn()
objects.append(obj)
create_success = True
print('create object success')
except:
print('create object failed')
pass
scene_is_valid = True
def get_total_number_of_fluid_particles():
num_particles = 0
for obj in objects:
if obj['type'] == 'fluid':
num_particles += obj['positions'].shape[0]
return num_particles
def get_smallest_fluid_object():
num_particles = 100000000
obj_idx = -1
for idx, obj in enumerate(objects):
if obj['type'] == 'fluid':
if obj['positions'].shape[0] < num_particles:
obj_idx = idx
num_particles = min(obj['positions'].shape[0],
num_particles)
return obj_idx, num_particles
total_number_of_fluid_particles = get_total_number_of_fluid_particles()
if options.const_fluid_particles:
if options.const_fluid_particles > total_number_of_fluid_particles:
scene_is_valid = False
else:
while get_total_number_of_fluid_particles(
) != options.const_fluid_particles:
difference = get_total_number_of_fluid_particles(
) - options.const_fluid_particles
obj_idx, num_particles = get_smallest_fluid_object()
if num_particles < difference:
del objects[obj_idx]
else:
objects[obj_idx]['positions'] = objects[obj_idx][
'positions'][:-difference]
objects[obj_idx]['velocities'] = objects[obj_idx][
'velocities'][:-difference]
if options.max_fluid_particles:
if options.max_fluid_particles < total_number_of_fluid_particles:
scene_is_valid = False
sim_directory = os.path.join(output_dir, 'sim_{0:04d}'.format(seed))
os.makedirs(sim_directory, exist_ok=False)
# generate scene json file
scene = {
'Configuration': default_configuration,
'Simulation': default_simulation,
# 'Fluid': default_fluid,
'RigidBodies': [],
'FluidModels': [],
}
rigid_body_next_id = 1
# bounding box
box_output_path = os.path.join(sim_directory, 'box.bgeo')
write_bgeo_from_numpy(box_output_path, bb, bb_normals)
box_obj_output_path = os.path.join(sim_directory, 'box.obj')
copyfile(bb_obj, box_obj_output_path)
rigid_body = deepcopy(default_rigidbody)
rigid_body['id'] = rigid_body_next_id
rigid_body_next_id += 1
rigid_body['geometryFile'] = os.path.basename(
os.path.abspath(box_obj_output_path))
rigid_body['resolutionSDF'] = [64, 64, 64]
rigid_body["collisionObjectType"] = 5
scene['RigidBodies'].append(rigid_body)
fluid_count = 0
for obj in objects:
fluid_id = 'fluid{0}'.format(fluid_count)
fluid_count += 1
fluid = deepcopy(default_fluid)
fluid['viscosity'] = obj['viscosity']
fluid['density0'] = obj['density']
scene[fluid_id] = fluid
fluid_model = deepcopy(default_fluidmodel)
fluid_model['id'] = fluid_id
fluid_output_path = os.path.join(sim_directory, fluid_id + '.bgeo')
write_bgeo_from_numpy(fluid_output_path, obj['positions'],
obj['velocities'])
fluid_model['particleFile'] = os.path.basename(fluid_output_path)
scene['FluidModels'].append(fluid_model)
scene_output_path = os.path.join(sim_directory, 'scene.json')
with open(scene_output_path, 'w') as f:
json.dump(scene, f, indent=4)
run_simulator(os.path.abspath(scene_output_path), sim_directory)
def create_data(output_dir, seed, options):
"""Creates a random scene for a specific seed and runs the simulator"""
np.random.seed(seed)
script_dir = os.path.dirname(__file__)
# convert bounding box to particles
boundary_box = os.path.join(script_dir, 'models', 'Box2.obj')
bb, bb_normals = obj_surface_to_particles(boundary_box)
bb_vol = obj_volume_to_particles(boundary_box)[0]
fluid_shapes = sorted(
glob(os.path.join(script_dir, 'models', 'Fluid*.obj')))
# obstacle_shapes = sorted(glob(os.path.join(script_dir, 'models', 'Rigid*.obj')))
obstacle_shapes = [os.path.join(script_dir, 'models', 'Rigid_001.obj')]
num_objects = seed % 3 + 1
print('num_objects', num_objects)
# randomly placed fluid object's position can be invalid
scene_is_valid = False
for create_scene_i in range(100):
if scene_is_valid:
break
# rasterize free volume
bb_rasterized = rasterize_points(np.concatenate([bb_vol, bb], axis=0),
2.01 * PARTICLE_RADIUS,
PARTICLE_RADIUS)
bb_rasterized = bb_rasterized[0], bb_rasterized[1], binary_erosion(
bb_rasterized[2], structure=np.ones((3, 3, 3)), iterations=3)
# create obstacle
if options.obstacle:
create_success = False
for i in range(10):
if create_success:
break
try:
obstacle_obj = np.random.choice(obstacle_shapes)
R = random_rotation_matrix(1.0)
obstacle_info = create_rigid_object(
obstacle_obj, R, bb_rasterized)
create_success = True
print('create obstacle success')
except:
print('create obstacle failed')
pass
objects = []
create_fn_list = [create_fluid_object]
for object_i in range(num_objects):
create_fn = np.random.choice(create_fn_list)
create_success = False
for i in range(10):
if create_success:
break
try:
obj = create_fn(fluid_shapes, bb_rasterized)
objects.append(obj)
create_success = True
print('create object success')
except:
print('create object failed')
pass
scene_is_valid = True
# ignore all scene options
# total_number_of_fluid_particles = get_total_number_of_fluid_particles()
#
# if options.const_fluid_particles:
# if options.const_fluid_particles > total_number_of_fluid_particles:
# scene_is_valid = False
# else:
# while get_total_number_of_fluid_particles(
# ) != options.const_fluid_particles:
# difference = get_total_number_of_fluid_particles(
# ) - options.const_fluid_particles
# obj_idx, num_particles = get_smallest_fluid_object()
# if num_particles < difference:
# del objects[obj_idx]
# else:
# objects[obj_idx]['positions'] = objects[obj_idx][
# 'positions'][:-difference]
# objects[obj_idx]['velocities'] = objects[obj_idx][
# 'velocities'][:-difference]
#
# if options.max_fluid_particles:
# if options.max_fluid_particles < total_number_of_fluid_particles:
# scene_is_valid = False
sim_directory = os.path.join(output_dir, 'sim_{0:04d}'.format(seed))
os.makedirs(sim_directory, exist_ok=False)
# generate scene json file
scene = {
'Configuration': default_configuration,
'Simulation': default_simulation,
# 'Fluid': default_fluid,
'RigidBodies': [],
'FluidModels': [],
}
rigid_body_next_id = 1
# bounding box
box_output_path = os.path.join(sim_directory, 'box.bgeo')
write_bgeo_from_numpy(box_output_path, bb, bb_normals)
box_obj_output_path = os.path.join(sim_directory, 'box.obj')
copyfile(boundary_box, box_obj_output_path)
rigid_body = deepcopy(default_rigidbody)
rigid_body['id'] = rigid_body_next_id
rigid_body_next_id += 1
rigid_body['geometryFile'] = os.path.basename(
os.path.abspath(box_obj_output_path))
rigid_body['resolutionSDF'] = [64, 64, 64]
rigid_body["collisionObjectType"] = 5
scene['RigidBodies'].append(rigid_body)
# obstacle
if options.obstacle:
obstacle_output_path = os.path.join(sim_directory, 'obstacle.bgeo')
obs, obs_normals = obj_surface_to_particles(
obstacle_obj, num_points=default_obstacle_size)
obs = obs @ R
obs_normals = obs_normals @ R
write_bgeo_from_numpy(obstacle_output_path, obs, obs_normals)
obstacle_obj_output_path = os.path.join(sim_directory, 'obstacle.obj')
copyfile(obstacle_obj, obstacle_obj_output_path)
obstacle = deepcopy(default_rigidbody)
obstacle['id'] = rigid_body_next_id
rigid_body_next_id += 1
obstacle['translation'] = obstacle_info['positions'].tolist()
vec, angle = rotation_matrix_to_axis_angle(R)
obstacle['rotationAxis'] = vec.tolist()
obstacle['rotationAngle'] = angle
obstacle['geometryFile'] = os.path.basename(
os.path.abspath(obstacle_obj_output_path))
obstacle['resolutionSDF'] = [64, 64, 64]
obstacle["collisionObjectType"] = 5
scene['RigidBodies'].append(obstacle)
fluid_count = 0
for obj in objects:
fluid_id = 'fluid{0}'.format(fluid_count)
fluid_count += 1
fluid = deepcopy(default_fluid)
fluid['viscosity'] = obj['viscosity']
fluid['density0'] = obj['density']
scene[fluid_id] = fluid
fluid_model = deepcopy(default_fluidmodel)
fluid_model['id'] = fluid_id
fluid_output_path = os.path.join(sim_directory, fluid_id + '.bgeo')
write_bgeo_from_numpy(fluid_output_path, obj['positions'],
obj['velocities'])
fluid_model['particleFile'] = os.path.basename(fluid_output_path)
scene['FluidModels'].append(fluid_model)
scene_output_path = os.path.join(sim_directory, 'scene.json')
with open(scene_output_path, 'w') as f:
json.dump(scene, f, indent=4)
run_simulator(os.path.abspath(scene_output_path), sim_directory)