def _is_moving(o_id: int, transforms: Transforms,
rigidbodies: Rigidbodies) -> bool:
"""
:param o_id: The ID of the object.
:param transforms: The Transforms output data.
:param rigidbodies: The Rigidbodies output data.
:return: True if the object is still moving.
"""
y: Optional[float] = None
sleeping: bool = False
for i in range(transforms.get_num()):
if transforms.get_id(i) == o_id:
y = transforms.get_position(i)[1]
break
assert y is not None, f"y value is none for {o_id}"
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == o_id:
sleeping = rigidbodies.get_sleeping(i)
break
# If the object isn't sleeping, it is still moving.
# If the object fell into the abyss, we don't count it as moving (to prevent an infinitely long simulation).
return not sleeping and y > -10
def _get_velocity(rigidbodies: Rigidbodies, o_id: int) -> float:
"""
:param rigidbodies: The rigidbody data.
:param o_id: The ID of the object.
:return: The velocity magnitude of the object.
"""
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == o_id:
return np.linalg.norm(rigidbodies.get_velocity(i))
def get_collisions(
resp: List[bytes]
) -> Tuple[List[Collision], List[EnvironmentCollision],
Optional[Rigidbodies]]:
"""
Parse collision and rigibody data from the output data.
:param resp: The response from the build.
:return: A list of collisions on this frame (can be empty), a list of environment collisions on this frame (can be empty), and Rigidbodies data (can be `None`).
"""
if len(resp) == 1:
return [], [], None
collisions: List[Collision] = []
environment_collisions: List[EnvironmentCollision] = []
rigidbodies: Optional[Rigidbodies] = None
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == 'coll':
collisions.append(Collision(r))
if r_id == 'rigi':
rigidbodies = Rigidbodies(r)
if r_id == 'enco':
environment_collisions.append(EnvironmentCollision(r))
return collisions, environment_collisions, rigidbodies
def _write_frame(self, frames_grp: h5py.Group, resp: List[bytes], frame_num: int) -> \
Tuple[h5py.Group, h5py.Group, dict, bool]:
frame, objs, tr, done = super()._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame_num)
num_objects = len(self.object_ids)
# Physics data.
velocities = np.empty(dtype=np.float32, shape=(num_objects, 3))
angular_velocities = np.empty(dtype=np.float32, shape=(num_objects, 3))
# Collision data.
collision_ids = np.empty(dtype=np.int32, shape=(0, 2))
collision_relative_velocities = np.empty(dtype=np.float32, shape=(0, 3))
collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3))
# Environment Collision data.
env_collision_ids = np.empty(dtype=np.int32, shape=(0, 1))
env_collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3))
sleeping = True
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "rigi":
ri = Rigidbodies(r)
ri_dict = dict()
for i in range(ri.get_num()):
ri_dict.update({ri.get_id(i): {"vel": ri.get_velocity(i),
"ang": ri.get_angular_velocity(i)}})
# Check if any objects are sleeping that aren't in the abyss.
if not ri.get_sleeping(i) and tr[ri.get_id(i)]["pos"][1] >= -1:
sleeping = False
# Add the Rigibodies data.
for o_id, i in zip(self.object_ids, range(num_objects)):
velocities[i] = ri_dict[o_id]["vel"]
angular_velocities[i] = ri_dict[o_id]["ang"]
elif r_id == "coll":
co = Collision(r)
collision_ids = np.append(collision_ids, [co.get_collider_id(), co.get_collidee_id()])
collision_relative_velocities = np.append(collision_relative_velocities, co.get_relative_velocity())
for i in range(co.get_num_contacts()):
collision_contacts = np.append(collision_contacts, (co.get_contact_normal(i),
co.get_contact_point(i)))
elif r_id == "enco":
en = EnvironmentCollision(r)
env_collision_ids = np.append(env_collision_ids, en.get_object_id())
for i in range(en.get_num_contacts()):
env_collision_contacts = np.append(env_collision_contacts, (en.get_contact_normal(i),
en.get_contact_point(i)))
objs.create_dataset("velocities", data=velocities.reshape(num_objects, 3), compression="gzip")
objs.create_dataset("angular_velocities", data=angular_velocities.reshape(num_objects, 3), compression="gzip")
collisions = frame.create_group("collisions")
collisions.create_dataset("object_ids", data=collision_ids.reshape((-1, 2)), compression="gzip")
collisions.create_dataset("relative_velocities", data=collision_relative_velocities.reshape((-1, 3)),
compression="gzip")
collisions.create_dataset("contacts", data=collision_contacts.reshape((-1, 2, 3)), compression="gzip")
env_collisions = frame.create_group("env_collisions")
env_collisions.create_dataset("object_ids", data=env_collision_ids, compression="gzip")
env_collisions.create_dataset("contacts", data=env_collision_contacts.reshape((-1, 2, 3)),
compression="gzip")
return frame, objs, tr, sleeping
def __init__(self, index: int, rigidbodies: Rigidbodies,
segmentation_colors: SegmentationColors, audio: ObjectInfo):
"""
:param index: The index of the object in `segmentation_colors`
:param rigidbodies: Rigidbodies output data.
:param segmentation_colors: Segmentation colors output data.
"""
self.object_id = segmentation_colors.get_object_id(index)
self.model_name = segmentation_colors.get_object_name(index)
self.segmentation_color = np.array(
segmentation_colors.get_object_color(index))
self.audio = audio
# Get the mass.
self.mass: float = -1
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == self.object_id:
self.mass = rigidbodies.get_mass(i)
break
assert self.mass >= 0, f"Mass not found: {self.object_id}"
def _get_frame_state(t: Transforms, r: Rigidbodies, num: int) -> dict:
"""
Returns a dictionary representing the current frame state.
:param t: The transforms data.
:param r: The rigidbodies data.
:param num: The current frame.
"""
objects = dict()
assert t.get_num() == r.get_num()
for i in range(t.get_num()):
objects.update({
t.get_id(i): {
"position": t.get_position(i),
"rotation": t.get_rotation(i),
"forward": t.get_forward(i),
"velocity": r.get_velocity(i),
"angular_velocity": r.get_angular_velocity(i),
"mass": r.get_mass(i)
}
})
return {"frame": num, "objects": objects}
def get_sound(self, collision: Union[Collision, EnvironmentCollision],
rigidbodies: Rigidbodies, id1: int, mat1: str, id2: int,
mat2: str, other_amp: float, target_amp: float,
resonance: float) -> Optional[Base64Sound]:
"""
Produce sound of two colliding objects as a byte array.
:param collision: TDW `Collision` or `EnvironmentCollision` output data.
:param rigidbodies: TDW `Rigidbodies` output data.
:param id1: The object ID for one of the colliding objects.
:param mat1: The material label for one of the colliding objects.
:param id2: The object ID for the other object.
:param mat2: The material label for the other object.
:param other_amp: Sound amplitude of object 2.
:param target_amp: Sound amplitude of object 1.
:param resonance: The resonances of the objects.
:return Sound data as a Base64Sound object.
"""
# The sound amplitude of object 2 relative to that of object 1.
amp2re1 = other_amp / target_amp
# Set the object modes.
if id2 not in self.object_modes:
self.object_modes.update({id2: {}})
if id1 not in self.object_modes[id2]:
self.object_modes[id2].update({
id1:
CollisionInfo(self._get_object_modes(mat2),
self._get_object_modes(mat1),
amp=target_amp * self.initial_amp)
})
obj_col = isinstance(collision, Collision)
# Unpack useful parameters.
# Compute normal velocity at impact.
vel = 0
if obj_col:
vel = collision.get_relative_velocity()
else:
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == id2:
vel = rigidbodies.get_velocity(i)
# If the y coordinate of the velocity is negative, it implies a scrape or roll along the floor.
if vel[1] < 0:
return None
break
vel = np.asarray(vel)
speed = np.square(vel)
speed = np.sum(speed)
speed = math.sqrt(speed)
nvel = vel / np.linalg.norm(vel)
num_contacts = collision.get_num_contacts()
nspd = []
for jc in range(0, num_contacts):
tmp = np.asarray(collision.get_contact_normal(jc))
tmp = tmp / np.linalg.norm(tmp)
tmp = np.arccos(np.clip(np.dot(tmp, nvel), -1.0, 1.0))
# Scale the speed by the angle (i.e. we want speed Normal to the surface).
tmp = speed * np.cos(tmp)
nspd.append(tmp)
normal_speed = np.mean(nspd)
# Get indices of objects in collisions.
id1_index = None
id2_index = None
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == id1:
id1_index = i
if rigidbodies.get_id(i) == id2:
id2_index = i
# Use default values for environment collisions.
if not obj_col:
m1 = 100
m2 = rigidbodies.get_mass(id2_index)
# Use the Rigidbody masses.
elif id1_index is not None and id2_index is not None:
m1 = rigidbodies.get_mass(id1_index)
m2 = rigidbodies.get_mass(id2_index)
# Fallback: Try to use default mass values if the ID's aren't in the Rigidbody data.
elif id1 in self.object_names and id2 in self.object_names and self.object_names[id1] in self.object_info and \
self.object_names[id2] in self.object_info:
m1 = self.object_info[self.object_names[id1]].mass
m2 = self.object_info[self.object_names[id2]].mass
# Failed to generate a sound.
else:
return None
mass = np.min([m1, m2])
# Re-scale the amplitude.
if self.object_modes[id2][id1].count == 0:
# Sample the modes.
sound, modes_1, modes_2 = self.make_impact_audio(
amp2re1,
mass,
mat1=mat1,
mat2=mat2,
id1=id1,
id2=id2,
resonance=resonance)
# Save collision info - we will need for later collisions.
amp = self.object_modes[id2][id1].amp
self.object_modes[id2][id1].init_speed = normal_speed
self.object_modes[id2][id1].obj1_modes = modes_1
self.object_modes[id2][id1].obj2_modes = modes_2
else:
amp = self.object_modes[id2][
id1].amp * normal_speed / self.object_modes[id2][id1].init_speed
# Adjust modes here so that two successive impacts are not identical.
modes_1 = self.object_modes[id2][id1].obj1_modes
modes_2 = self.object_modes[id2][id1].obj2_modes
modes_1.powers = modes_1.powers + np.random.normal(
0, 2, len(modes_1.powers))
modes_2.powers = modes_2.powers + np.random.normal(
0, 2, len(modes_2.powers))
sound = PyImpact.synth_impact_modes(modes_1, modes_2, mass,
resonance)
self.object_modes[id2][id1].obj1_modes = modes_1
self.object_modes[id2][id1].obj2_modes = modes_2
if self.logging:
mode_props = dict()
self.log_modes(self.object_modes[id2][id1].count, mode_props, id1,
id2, modes_1, modes_2, amp, str(mat1), str(mat2))
# On rare occasions, it is possible for PyImpact to fail to generate a sound.
if sound is None:
return None
# Count the collisions.
self.object_modes[id2][id1].count_collisions()
# Prevent distortion by clamping the amp.
if self.prevent_distortion and np.abs(amp) > 0.99:
amp = 0.99
sound = amp * sound / np.max(np.abs(sound))
return Base64Sound(sound)
def get_sound(self, collision: Union[Collision, EnvironmentCollision], rigidbodies: Rigidbodies, id1: int, mat1: str, id2: int, mat2: str, amp2re1: float) -> Optional[Base64Sound]:
"""
Produce sound of two colliding objects as a byte array.
:param collision: TDW `Collision` or `EnvironmentCollision` output data.
:param rigidbodies: TDW `Rigidbodies` output data.
:param id1: The object ID for one of the colliding objects.
:param mat1: The material label for one of the colliding objects.
:param id2: The object ID for the other object.
:param mat2: The material label for the other object.
:param amp2re1: The sound amplitude of object 2 relative to that of object 1.
:return Sound data as a Base64Sound object.
"""
# Set the object modes.
if id2 not in self.object_modes:
self.object_modes.update({id2: {}})
if id1 not in self.object_modes[id2]:
self.object_modes[id2].update({id1: CollisionInfo(self._get_object_modes(mat2),
self._get_object_modes(mat1),
amp=self.initial_amp)})
obj_col = isinstance(collision, Collision)
# Unpack useful parameters.
# Compute normal velocity at impact.
vel = 0
if obj_col:
vel = collision.get_relative_velocity()
else:
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == id2:
vel = rigidbodies.get_velocity(i)
break
vel = np.asarray(vel)
speed = np.square(vel)
speed = np.sum(speed)
speed = math.sqrt(speed)
nvel = vel / np.linalg.norm(vel)
num_contacts = collision.get_num_contacts()
nspd = []
for jc in range(0, num_contacts):
tmp = np.asarray(collision.get_contact_normal(jc))
tmp = tmp / np.linalg.norm(tmp)
tmp = np.arccos(np.clip(np.dot(tmp, nvel), -1.0, 1.0))
# Scale the speed by the angle (i.e. we want speed Normal to the surface).
tmp = speed * np.cos(tmp)
nspd.append(tmp)
normal_speed = np.mean(nspd)
# Get indices of objects in collisions
id1_index = None
id2_index = None
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == id1:
id1_index = i
if rigidbodies.get_id(i) == id2:
id2_index = i
# Make sure both IDs were found. If they aren't, don't return a sound.
if obj_col and (id1_index is None or id2_index is None):
return None
m1 = rigidbodies.get_mass(id1_index) if obj_col else 1000
m2 = rigidbodies.get_mass(id2_index)
mass = np.min([m1, m2])
# Re-scale the amplitude.
if self.object_modes[id2][id1].count == 0:
# Sample the modes.
sound, modes_1, modes_2 = self.make_impact_audio(amp2re1, mass, mat1=mat1, mat2=mat2, id1=id1, id2=id2)
# Save collision info - we will need for later collisions.
amp = self.object_modes[id2][id1].amp
self.object_modes[id2][id1].init_speed = normal_speed
self.object_modes[id2][id1].obj1_modes = modes_1
self.object_modes[id2][id1].obj2_modes = modes_2
else:
amp = self.object_modes[id2][id1].amp * normal_speed / self.object_modes[id2][id1].init_speed
# Adjust modes here so that two successive impacts are not identical.
modes_1 = self.object_modes[id2][id1].obj1_modes
modes_2 = self.object_modes[id2][id1].obj2_modes
modes_1.powers = modes_1.powers + np.random.normal(0, 2, len(modes_1.powers))
modes_2.power = modes_2.powers + np.random.normal(0, 2, len(modes_1.powers))
sound = PyImpact.synth_impact_modes(modes_1, modes_2, mass)
self.object_modes[id2][id1].obj1_modes = modes_1
self.object_modes[id2][id1].obj2_modes = modes_2
# On rare occasions, it is possible for PyImpact to fail to generate a sound.
if sound is None:
return None
# Count the collisions.
self.object_modes[id2][id1].count_collisions()
# Prevent distortion by clamping the amp.
if self.prevent_distortion and amp > 0.99:
amp = 0.99
sound = amp * sound / np.max(np.abs(sound))
return Base64Sound(sound)
def run(self, c: Controller):
"""
Run the trial and save the output.
:param c: The controller.
"""
print(f"Images will be saved to: {self.output_dir}")
# Initialize the scene.
resp = c.communicate(self.init_commands)
# Get a map of the segmentation colors.
segm = SegmentationColors(resp[0])
for i in range(segm.get_num()):
for obj in self.moving_objects:
if obj.object_id == segm.get_object_id(i):
obj.possibility.segmentation_color = segm.get_object_color(
i)
# Request scene data and images per frame.
frame_data: List[dict] = []
resp = c.communicate([{
"$type": "send_images",
"frequency": "always"
}, {
"$type": "send_transforms",
"ids": self.m_ids,
"frequency": "always"
}, {
"$type": "send_rigidbodies",
"ids": self.m_ids,
"frequency": "always"
}])
# Run the trial.
for frame in range(self.num_frames):
colors: Dict[int, Tuple[int, int, int]] = {}
transforms: Dict[int, Tuple[float, float, float]] = {}
transform_data = None
rigidbody_data = None
# Parse the output data.
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Record all Transforms data.
if r_id == "tran":
transform_data = Transforms(r)
for i in range(transform_data.get_num()):
transforms.update({
transform_data.get_id(i):
transform_data.get_position(i)
})
# Record all Rigidbodies data.
elif r_id == "rigi":
rigidbody_data = Rigidbodies(r)
# Save the images.
elif r_id == "imag":
images = Images(r)
for p in range(images.get_num_passes()):
if images.get_pass_mask(p) == "_id":
image_colors = TDWUtils.get_pil_image(
images, p).getcolors()
for ic in image_colors:
color = ic[1]
for obj in self.moving_objects:
if obj.possibility.segmentation_color == color:
colors.update({obj.object_id: color})
TDWUtils.save_images(Images(r),
TDWUtils.zero_padding(frame),
output_directory=self.output_dir)
# Append frame data.
frame_data.append(
Trial._get_frame_state(transform_data, rigidbody_data, frame))
# Build the frame state.
state = State(colors, transforms, frame)
# Apply object actions.
commands = []
for o in self.occluders:
commands.extend(o.get_frame_commands(state))
for mo in self.moving_objects:
commands.extend(mo.get_frame_commands(state))
if len(commands) == 0:
commands = [{"$type": "do_nothing"}]
# Send the commands and update the state.
resp = c.communicate(commands)
# Cleanup.
c.communicate([{
"$type": "destroy_all_objects"
}, {
"$type": "unload_asset_bundles"
}, {
"$type": "send_images",
"frequency": "never"
}, {
"$type": "send_transforms",
"ids": self.m_ids,
"frequency": "never"
}, {
"$type": "send_rigidbodies",
"ids": self.m_ids,
"frequency": "never"
}])
print("\tGenerated images.")
# Output the scene metadata.
Path(self.output_dir).joinpath("state.json").write_text(
json.dumps({"frames": frame_data}), encoding="utf-8")
print("\tWrote state file.")
# Get _id passes with randomized colors.
self._randomize_segmentation_colors()
print("\tCreated random segmentation colors.")
# Organize the images.
self._organize_output()
print("\tOrganized files")
def __init__(self,
object_id: int,
rigidbodies: Rigidbodies,
segmentation_colors: SegmentationColors,
bounds: Bounds,
audio: ObjectInfo,
target_object: bool = False):
"""
:param object_id: The unique ID of the object.
:param rigidbodies: Rigidbodies output data.
:param bounds: Bounds output data.
:param segmentation_colors: Segmentation colors output data.
"""
self.object_id = object_id
self.model_name = audio.name
self.container = self.model_name in StaticObjectInfo.CONTAINERS
self.kinematic = self.model_name in StaticObjectInfo._KINEMATIC
self.target_object = target_object
self.category = ""
# This is a sub-object of a composite object.
if audio.library == "":
# Get the record of the composite object.
for k in StaticObjectInfo._COMPOSITE_OBJECTS:
for v in StaticObjectInfo._COMPOSITE_OBJECTS[k]:
if v == audio.name:
record = TransformInitData.LIBRARIES[
"models_core.json"].get_record(k)
# Get the semantic category.
self.category = record.wcategory
break
else:
# Get the model record from the audio data.
record = TransformInitData.LIBRARIES[audio.library].get_record(
audio.name)
# Get the semantic category.
self.category = record.wcategory
# Get the segmentation color.
self.segmentation_color: Optional[np.array] = None
for i in range(segmentation_colors.get_num()):
if segmentation_colors.get_object_id(i) == self.object_id:
self.segmentation_color = np.array(
segmentation_colors.get_object_color(i))
break
assert self.segmentation_color is not None, f"Segmentation color not found: {self.object_id}"
# Get the size of the object.
self.size = np.array([0, 0, 0])
for i in range(bounds.get_num()):
if bounds.get_id(i) == self.object_id:
self.size = np.array([
float(
np.abs(bounds.get_right(i)[0] -
bounds.get_left(i)[0])),
float(
np.abs(bounds.get_top(i)[1] -
bounds.get_bottom(i)[1])),
float(
np.abs(bounds.get_front(i)[2] - bounds.get_back(i)[2]))
])
break
assert np.linalg.norm(
self.size) > 0, f"Bounds data not found for: {self.object_id}"
# Get the mass.
self.mass: float = -1
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == self.object_id:
self.mass = rigidbodies.get_mass(i)
break
assert self.mass >= 0, f"Mass not found: {self.object_id}"