def __init__(self):
self.obj_name_dict: Dict[int, str] = {}
# Set up the object transform data.
object_setup_data = json.loads(Path("object_setup.json").read_text())
self.object_setups: List[_ObjectSetup] = []
for o in object_setup_data:
combo = _ObjectSetup(id=int(o),
model_name=object_setup_data[o]["model_name"],
position=object_setup_data[o]["position"],
rotation=object_setup_data[o]["rotation"],
scale=object_setup_data[o]["scale"])
self.object_setups.append(combo)
# Parse the default objects.csv spreadsheet.
self.object_audio_data = PyImpact.get_object_info()
self.temp_amp = 0.1
# Keep track of the current trial number, for logging purposes.
self.current_trial_num = 0
# Fetch the ball and board model's records; we will need them later to change its material.
self.special_models = ModelLibrarian(library="models_special.json")
self.full_models = ModelLibrarian(library="models_full.json")
self.ball_record = self.special_models.get_record("prim_sphere")
self.board_record = self.full_models.get_record("wood_board")
# Set path to write out logging info.
self.root_dest_dir = Path("dist/mode_properties_logs")
if not self.root_dest_dir.exists():
self.root_dest_dir.mkdir(parents=True)
super().__init__()
class AudioInitData(RigidbodyInitData):
"""
A subclass of `RigidbodyInitData` that includes [audio values](py_impact.md#objectinfo).
Physics values are derived from these audio values.
"""
_DYNAMIC_FRICTION = {AudioMaterial.ceramic: 0.47,
AudioMaterial.hardwood: 0.35,
AudioMaterial.wood: 0.35,
AudioMaterial.cardboard: 0.47,
AudioMaterial.glass: 0.65,
AudioMaterial.metal: 0.43}
_STATIC_FRICTION = {AudioMaterial.ceramic: 0.47,
AudioMaterial.hardwood: 0.4,
AudioMaterial.wood: 0.4,
AudioMaterial.cardboard: 0.47,
AudioMaterial.glass: 0.65,
AudioMaterial.metal: 0.52}
AUDIO = PyImpact.get_object_info()
def __init__(self, name: str, library: str = "models_core.json", scale_factor: Dict[str, float] = None, position: Dict[str, float] = None, rotation: Dict[str, float] = None, kinematic: bool = False, gravity: bool = True, audio: ObjectInfo = None):
"""
:param name: The name of the model.
:param library: The filename of the library containing the model's record.
:param scale_factor: The [scale factor](../api/command_api.md#scale_object).
:param position: The initial position. If None, defaults to: `{"x": 0, "y": 0, "z": 0`}.
:param rotation: The initial rotation as Euler angles or a quaternion. If None, defaults to: `{"w": 1, "x": 0, "y": 0, "z": 0}`
:param kinematic: If True, the object will be [kinematic](../api/command_api.md#set_kinematic_state).
:param gravity: If True, the object won't respond to [gravity](../api/command_api.md#set_kinematic_state).
:param audio: If None, derive physics data from the audio data in `PyImpact.get_object_info()` (if the object isn't in this dictionary, this constructor will throw an error). If not None, use these values instead of the default audio values.
"""
if audio is None:
self.audio = AudioInitData.AUDIO[name]
else:
self.audio = audio
super().__init__(name=name, library=library, scale_factor=scale_factor, position=position, rotation=rotation,
kinematic=kinematic, gravity=gravity, mass=self.audio.mass,
dynamic_friction=AudioInitData._DYNAMIC_FRICTION[self.audio.material],
static_friction=AudioInitData._STATIC_FRICTION[self.audio.material],
bounciness=self.audio.bounciness)
def get_commands(self) -> Tuple[int, List[dict]]:
"""
:return: Tuple: The ID of the object; a list of commands to create the object: `[add_object, rotate_object_to, scale_object, set_kinematic_state, set_object_collision_detection_mode, set_mass, set_physic_material]`
"""
return super().get_commands()
def __init__(self):
self.obj_name_dict: Dict[int, str] = {}
# Set up the object transform data.
object_setup_data = json.loads(Path("object_setup.json").read_text())
self.object_setups: List[_ObjectSetup] = []
for o in object_setup_data:
combo = _ObjectSetup(id=int(o),
model_name=object_setup_data[o]["model_name"],
position=object_setup_data[o]["position"],
rotation=object_setup_data[o]["rotation"],
scale=object_setup_data[o]["scale"])
self.object_setups.append(combo)
# Parse the default objects.csv spreadsheet.
self.object_audio_data = PyImpact.get_object_info()
# Fetch the ball and board model's records; we will need them later to change its material.
self.special_models = ModelLibrarian(library="models_special.json")
self.full_models = ModelLibrarian(library="models_full.json")
self.ball_record = self.special_models.get_record("prim_sphere")
self.board_record = self.full_models.get_record("wood_board")
super().__init__()
def do_trial(self):
# Initialize PyImpact and pass in the "master gain" amplitude value. This value must be betweem 0 and 1.
# The relative amplitudes of all scene objects involved in collisions will be scaled relative to this value.
# Note -- if this value is too high, waveform clipping can occur and the resultant audio will be distorted.
# For this reason, the value used here is considerably smaller than the corresponding value used in the
# impact_sounds.py example controller. Here we have a large number of closely-occuring collisions resulting in
# a rapid series of "clustered" impact sounds, as opposed to a single object falling from a height;
# using a higher value such as the 0.5 used in the example controller will definitely result in unpleasant
# distortion of the audio.
# Note that logging is also enabled.
# Keep track of trial number.
self.current_trial_num += 1
# Create folder for this trial's logging info.
dest_dir = self.root_dest_dir.joinpath(str(self.current_trial_num))
if not dest_dir.exists():
dest_dir.mkdir(parents=True)
dest_dir_str = str(dest_dir.resolve())
p = PyImpact(0.25, logging=True)
self.add_all_objects()
# "Aim" the ball at the monkey and apply the force.
# Note that this force value was arrived at through a number of trial-and-error iterations.
resp = self.communicate([{"$type": "object_look_at_position",
"id": 0,
"position": {"x": -12.95, "y": 1.591, "z": -5.1}},
{"$type": "apply_force_magnitude_to_object",
"id": 0,
"magnitude": 98.0}])
for i in range(400):
collisions, environment_collision, rigidbodies = PyImpact.get_collisions(resp)
# Sort the objects by mass.
masses: Dict[int, float] = {}
for j in range(rigidbodies.get_num()):
masses.update({rigidbodies.get_id(j): rigidbodies.get_mass(j)})
# If there was a collision, create an impact sound.
if len(collisions) > 0 and PyImpact.is_valid_collision(collisions[0]):
collider_id = collisions[0].get_collider_id()
collidee_id = collisions[0].get_collidee_id()
# The "target" object should have less mass than the "other" object.
if masses[collider_id] < masses[collidee_id]:
target_id = collider_id
other_id = collidee_id
else:
target_id = collidee_id
other_id = collider_id
target_name = self.obj_name_dict[target_id]
other_name = self.obj_name_dict[other_id]
impact_sound_command = p.get_impact_sound_command(
collision=collisions[0],
rigidbodies=rigidbodies,
target_id=target_id,
target_mat=self.object_audio_data[target_name].material,
other_id=other_id,
other_mat=self.object_audio_data[other_name].material,
target_amp=self.object_audio_data[target_name].amp,
other_amp=self.object_audio_data[other_name].amp,
resonance=self.object_audio_data[target_name].resonance)
resp = self.communicate(impact_sound_command)
# Continue to run the trial.
else:
resp = self.communicate([])
# Get the logging info for this trial and write it out.
log = p.get_log()
json_dest = dest_dir.joinpath("mode_properties_log.json")
json_dest.write_text(json.dumps(log, indent=2))
for obj_setup in self.object_setups:
self.communicate({"$type": "destroy_object", "id": obj_setup.id})
def trial(self):
"""
Select random objects and collide them to produce impact sounds.
"""
p = PyImpact(initial_amp=0.25)
# Set the environment audio materials.
floor = AudioMaterial.ceramic
wall = AudioMaterial.wood
# Create an empty room.
# Set the screen size.
# Adjust framerate so that physics is closer to realtime.
commands = [
TDWUtils.create_empty_room(12, 12), {
"$type": "destroy_all_objects"
}, {
"$type": "set_screen_size",
"width": 1024,
"height": 1024
}, {
"$type": "set_target_framerate",
"framerate": 100
}
]
# Create the avatar.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
position={
"x": 1,
"y": 1.2,
"z": 1.2
},
look_at=TDWUtils.VECTOR3_ZERO))
# Add the audio sensor.
commands.extend([{
"$type": "add_audio_sensor",
"avatar_id": "a"
}, {
"$type": "set_focus_distance",
"focus_distance": 2
}])
# Select a random pair of objects.
obj1_names = [
"trapezoidal_table", "glass_table_round", "yellow_side_chair",
"table_square", "marble_table"
]
obj2_names = ["vase_06", "spoon1", "glass3", "jug02"]
obj1_name = random.choice(obj1_names)
obj2_name = random.choice(obj2_names)
# Get initialization data from the default audio data (which includes mass, friction values, etc.).
obj1_init_data = AudioInitData(name=obj1_name)
obj2_init_data = AudioInitData(name=obj2_name,
position={
"x": 0,
"y": 2,
"z": 0
},
rotation={
"x": 135,
"y": 0,
"z": 30
})
# Convert the initialization data to commands.
obj1_id, obj1_commands = obj1_init_data.get_commands()
obj2_id, obj2_commands = obj2_init_data.get_commands()
# Cache the IDs and names of each object for PyImpact.
object_names = {obj1_id: obj1_name, obj2_id: obj2_name}
p.set_default_audio_info(object_names=object_names)
# Add the objects.
commands.extend(obj1_commands)
commands.extend(obj2_commands)
# Apply a small force to the dropped object.
# Request collision and rigidbody output data.
commands.extend([{
"$type": "apply_force_to_object",
"force": {
"x": 0,
"y": -0.01,
"z": 0
},
"id": obj2_id
}, {
"$type": "send_collisions",
"enter": True,
"stay": False,
"exit": False,
"collision_types": ["obj", "env"]
}, {
"$type": "send_rigidbodies",
"frequency": "always",
"ids": [obj2_id, obj1_id]
}])
# Send all of the commands.
resp = self.communicate(commands)
# Iterate through 200 frames.
# Every frame, listen for collisions, and parse the output data.
for i in range(200):
# Use PyImpact to generate audio from the output data and then convert the audio to TDW commands.
# If no audio is generated, then `commands` is an empty list.
commands = p.get_audio_commands(resp=resp, floor=floor, wall=wall)
# Send the commands to TDW in order to play the audio.
resp = self.communicate(commands)
def trial(self):
"""
Select random objects and collide them to produce impact sounds.
"""
p = PyImpact(initial_amp=0.5)
# Destroy all objects currently in the scene.
# Set the screen size.
# Adjust physics timestep for more real-time physics behavior.
commands = [{
"$type": "destroy_all_objects"
}, {
"$type": "set_screen_size",
"width": 1024,
"height": 1024
}, {
"$type": "set_time_step",
"time_step": 0.02
}]
# Create the avatar.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
position={
"x": 1,
"y": 1.2,
"z": 1.2
},
look_at=TDWUtils.VECTOR3_ZERO))
# Add the audio sensor.
# Set the target framerate.
# Make sure that post-processing is enabled and render quality is set to max.
commands.extend([{
"$type": "add_audio_sensor",
"avatar_id": "a"
}, {
"$type": "set_target_framerate",
"framerate": 60
}, {
"$type": "set_post_process",
"value": True
}, {
"$type": "set_focus_distance",
"focus_distance": 2
}, {
"$type": "set_render_quality",
"render_quality": 5
}])
# Select a random pair of objects.
objects = PyImpact.get_object_info()
obj1_names = [
"trapezoidal_table", "glass_table_round", "yellow_side_chair",
"table_square", "marble_table"
]
obj2_names = ["vase_06", "spoon1", "glass3", "jug02"]
obj1_name = random.choice(obj1_names)
obj2_name = random.choice(obj2_names)
obj1_id = 0
obj2_id = 1
# Add the objects.
# Set their masses from the audio info data.
# Set a physics material for the second object.
# Apply a force to the second object.
# Listen for collisions, and object properties.
commands.extend([
self.get_add_object(model_name=obj1_name,
object_id=obj1_id,
library=objects[obj1_name].library), {
"$type": "set_mass",
"id": obj1_id,
"mass": objects[obj1_name].mass
},
self.get_add_object(model_name=obj2_name,
object_id=obj2_id,
library=objects[obj2_name].library,
rotation={
"x": 135,
"y": 0,
"z": 30
},
position={
"x": 0,
"y": 2,
"z": 0
}), {
"$type": "set_mass",
"id": obj2_id,
"mass": objects[obj2_name].mass
}, {
"$type": "set_physic_material",
"id": obj2_id,
"bounciness":
objects[obj2_name].bounciness,
"dynamic_friction": 0.8
}, {
"$type": "apply_force_to_object",
"force": {
"x": 0,
"y": -0.01,
"z": 0
},
"id": obj2_id
}, {
"$type": "send_collisions",
"enter": True,
"stay": False,
"exit": False
}, {
"$type": "send_rigidbodies",
"frequency": "always",
"ids": [obj2_id, obj1_id]
}
])
# Send all of the commands.
resp = self.communicate(commands)
# Iterate through 200 frames.
# Every frame, listen for collisions, and parse the output data.
for i in range(200):
collisions, environment_collision, rigidbodies = PyImpact.get_collisions(
resp)
# If there was a collision, create an impact sound.
if len(collisions) > 0 and PyImpact.is_valid_collision(
collisions[0]):
impact_sound_command = p.get_impact_sound_command(
collision=collisions[0],
rigidbodies=rigidbodies,
target_id=obj2_id,
target_mat=objects[obj2_name].material.name,
target_amp=objects[obj2_name].amp,
other_id=obj1_id,
other_amp=objects[obj1_name].amp,
other_mat=objects[obj1_name].material.name,
resonance=objects[obj1_name].resonance)
resp = self.communicate(impact_sound_command)
# Continue to run the trial.
else:
resp = self.communicate([])
# Stop listening for collisions and rigidbodies.
self.communicate([{
"$type": "send_collisions",
"frequency": "never"
}, {
"$type": "send_rigidbodies",
"frequency": "never"
}])
def trial(self, scene: Scene, record: ModelRecord, output_path: Path,
scene_index: int) -> None:
"""
Run a trial in a scene that has been initialized.
:param scene: Data for the current scene.
:param record: The model's metadata record.
:param output_path: Write the .wav file to this path.
:param scene_index: The scene identifier.
"""
self.py_impact.reset(initial_amp=0.05)
# Initialize the scene, positioning objects, furniture, etc.
resp = self.communicate(scene.initialize_scene(self))
center = scene.get_center(self)
max_y = scene.get_max_y()
# The object's initial position.
o_x = RNG.uniform(center["x"] - 0.15, center["x"] + 0.15)
o_y = RNG.uniform(max_y - 0.5, max_y)
o_z = RNG.uniform(center["z"] - 0.15, center["z"] + 0.15)
# Physics values.
mass = self.object_info[record.name].mass + RNG.uniform(
self.object_info[record.name].mass * -0.15,
self.object_info[record.name].mass * 0.15)
static_friction = RNG.uniform(0.1, 0.3)
dynamic_friction = RNG.uniform(0.7, 0.9)
# Angles of rotation.
yaw = RNG.uniform(-30, 30)
pitch = RNG.uniform(0, 45)
roll = RNG.uniform(-45, 45)
# The force applied to the object.
force = RNG.uniform(0, 5)
# The avatar's position.
a_r = RNG.uniform(1.5, 2.2)
a_x = center["x"] + a_r
a_y = RNG.uniform(1.5, 3)
a_z = center["z"] + a_r
cam_angle_min, cam_angle_max = scene.get_camera_angles()
theta = np.radians(RNG.uniform(cam_angle_min, cam_angle_max))
a_x = np.cos(theta) * (a_x - center["x"]) - np.sin(theta) * (
a_z - center["z"]) + center["x"]
a_z = np.sin(theta) * (a_x - center["x"]) + np.cos(theta) * (
a_z - center["z"]) + center["z"]
o_id = 0
# Create the object and apply a force.
commands = [{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": {
"x": o_x,
"y": o_y,
"z": o_z
},
"category": record.wcategory,
"id": o_id
}, {
"$type": "set_mass",
"id": o_id,
"mass": mass
}, {
"$type": "set_physic_material",
"id": o_id,
"bounciness": self.object_info[record.name].bounciness,
"static_friction": static_friction,
"dynamic_friction": dynamic_friction
}, {
"$type": "rotate_object_by",
"angle": yaw,
"id": o_id,
"axis": "yaw",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": pitch,
"id": o_id,
"axis": "pitch",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": roll,
"id": o_id,
"axis": "roll",
"is_world": True
}, {
"$type": "apply_force_magnitude_to_object",
"magnitude": force,
"id": o_id
}, {
"$type": "send_rigidbodies",
"frequency": "always"
}, {
"$type": "send_collisions",
"enter": True,
"exit": False,
"stay": False,
"collision_types": ["obj", "env"]
}, {
"$type": "send_transforms",
"frequency": "always"
}]
# Parse bounds data to get the centroid of all objects currently in the scene.
bounds = Bounds(resp[0])
if bounds.get_num() == 0:
look_at = {"x": center["x"], "y": 0.1, "z": center["z"]}
else:
centers = []
for i in range(bounds.get_num()):
centers.append(bounds.get_center(i))
centers_x, centers_y, centers_z = zip(*centers)
centers_len = len(centers_x)
look_at = {
"x": sum(centers_x) / centers_len,
"y": sum(centers_y) / centers_len,
"z": sum(centers_z) / centers_len
}
# Add the avatar.
# Set the position at a given distance (r) from the center of the scene.
# Rotate around that position to a random angle constrained by the scene's min and max angles.
commands.extend([{
"$type": "teleport_avatar_to",
"position": {
"x": a_x,
"y": a_y,
"z": a_z
}
}, {
"$type": "look_at_position",
"position": look_at
}])
# Send the commands.
resp = self.communicate(commands)
AudioUtils.start(output_path=output_path, until=(0, 10))
# Loop until all objects are sleeping.
done = False
while not done and AudioUtils.is_recording():
commands = []
collisions, environment_collisions, rigidbodies = PyImpact.get_collisions(
resp)
# Create impact sounds from object-object collisions.
for collision in collisions:
if PyImpact.is_valid_collision(collision):
# Get the audio material and amp.
collider_id = collision.get_collider_id()
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
collidee_id = collision.get_collider_id()
collidee_material, collidee_amp = self._get_object_info(
collidee_id, Scene.OBJECT_IDS, record.name)
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collidee_id,
target_amp=collidee_amp,
target_mat=collidee_material.name,
other_id=collider_id,
other_mat=collider_material.name,
other_amp=collider_amp,
play_audio_data=False)
commands.append(impact_sound_command)
# Create impact sounds from object-environment collisions.
for collision in environment_collisions:
collider_id = collision.get_object_id()
if self._get_velocity(rigidbodies, collider_id) > 0:
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
surface_material = scene.get_surface_material()
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collider_id,
target_amp=collider_amp,
target_mat=collider_material.name,
other_id=-1,
other_amp=0.01,
other_mat=surface_material.name,
play_audio_data=False)
commands.append(impact_sound_command)
# If there were no collisions, check for movement. If nothing is moving, the trial is done.
if len(commands) == 0:
transforms = AudioDataset._get_transforms(resp)
done = True
for i in range(rigidbodies.get_num()):
if self._is_moving(rigidbodies.get_id(i), transforms,
rigidbodies):
done = False
break
# Continue the trial.
if not done:
resp = self.communicate(commands)
# Stop listening for anything except audio data..
resp = self.communicate([{
"$type": "send_rigidbodies",
"frequency": "never"
}, {
"$type": "send_transforms",
"frequency": "never"
}, {
"$type": "send_collisions",
"enter": False,
"exit": False,
"stay": False,
"collision_types": []
}, {
"$type": "send_audio_sources",
"frequency": "always"
}])
# Wait for the audio to finish.
done = False
while not done and AudioUtils.is_recording():
done = True
for r in resp[:-1]:
if OutputData.get_data_type_id(r) == "audi":
audio_sources = AudioSources(r)
for i in range(audio_sources.get_num()):
if audio_sources.get_is_playing(i):
done = False
if not done:
resp = self.communicate([])
# Cleanup.
commands = [{
"$type": "send_audio_sources",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": o_id
}]
for scene_object_id in Scene.OBJECT_IDS:
commands.append({"$type": "destroy_object", "id": scene_object_id})
self.communicate(commands)
# Insert the trial's values into the database.
self.db_c.execute(
"INSERT INTO sound20k VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
(output_path.name, scene_index, a_x, a_y, a_z, o_x, o_y, o_z, mass,
static_friction, dynamic_friction, yaw, pitch, roll, force))
self.conn.commit()
def __init__(self,
output_dir: Path = Path("D:/audio_dataset"),
total: int = 28602,
port: int = 1071):
"""
:param output_dir: The output directory for the files.
:param port: The socket port.
:param total: The total number of files per sub-set.
"""
self.total = total
self.output_dir = output_dir
if not self.output_dir.exists():
self.output_dir.mkdir(parents=True)
db_path = self.output_dir.joinpath('results.db')
self.conn = sqlite3.connect(str(db_path.resolve()))
self.db_c = self.conn.cursor()
# Sound20K table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sound20k'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE sound20k (path text, scene integer, cam_x real, cam_y real, cam_z real,"
"obj_x real, obj_y real, obj_z real, mass real, static_friction real, dynamic_friction "
"real, yaw real, pitch real, roll real, force real)")
# Scenes table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='scenes'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE scenes (id integer, commands text)")
self.py_impact = PyImpact()
self.object_info = PyImpact.get_object_info()
sound20k_object_info = PyImpact.get_object_info(
Path("models/object_info.csv"))
for obj_info in sound20k_object_info:
if obj_info in self.object_info:
continue
else:
self.object_info.update(
{obj_info: sound20k_object_info[obj_info]})
self.libs: Dict[str, ModelLibrarian] = {}
# Load all model libraries into memory.
for lib_name in ModelLibrarian.get_library_filenames():
self.libs.update({lib_name: ModelLibrarian(lib_name)})
# Add the custom model library.
self.libs.update({
"models/models.json":
ModelLibrarian(str(Path("models/models.json").resolve()))
})
super().__init__(port=port)
# Global settings.
self.communicate([{
"$type": "set_screen_size",
"width": 256,
"height": 256
}, {
"$type": "set_time_step",
"time_step": 0.02
}, {
"$type": "set_target_framerate",
"framerate": 60
}, {
"$type": "set_physics_solver_iterations",
"iterations": 20
}])
class AudioDataset(Controller):
def __init__(self,
output_dir: Path = Path("D:/audio_dataset"),
total: int = 28602,
port: int = 1071):
"""
:param output_dir: The output directory for the files.
:param port: The socket port.
:param total: The total number of files per sub-set.
"""
self.total = total
self.output_dir = output_dir
if not self.output_dir.exists():
self.output_dir.mkdir(parents=True)
db_path = self.output_dir.joinpath('results.db')
self.conn = sqlite3.connect(str(db_path.resolve()))
self.db_c = self.conn.cursor()
# Sound20K table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sound20k'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE sound20k (path text, scene integer, cam_x real, cam_y real, cam_z real,"
"obj_x real, obj_y real, obj_z real, mass real, static_friction real, dynamic_friction "
"real, yaw real, pitch real, roll real, force real)")
# Scenes table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='scenes'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE scenes (id integer, commands text)")
self.py_impact = PyImpact()
self.object_info = PyImpact.get_object_info()
sound20k_object_info = PyImpact.get_object_info(
Path("models/object_info.csv"))
for obj_info in sound20k_object_info:
if obj_info in self.object_info:
continue
else:
self.object_info.update(
{obj_info: sound20k_object_info[obj_info]})
self.libs: Dict[str, ModelLibrarian] = {}
# Load all model libraries into memory.
for lib_name in ModelLibrarian.get_library_filenames():
self.libs.update({lib_name: ModelLibrarian(lib_name)})
# Add the custom model library.
self.libs.update({
"models/models.json":
ModelLibrarian(str(Path("models/models.json").resolve()))
})
super().__init__(port=port)
# Global settings.
self.communicate([{
"$type": "set_screen_size",
"width": 256,
"height": 256
}, {
"$type": "set_time_step",
"time_step": 0.02
}, {
"$type": "set_target_framerate",
"framerate": 60
}, {
"$type": "set_physics_solver_iterations",
"iterations": 20
}])
def remove_output_directory(self) -> None:
"""
Delete the old directory.
"""
dir_util.remove_tree(str(self.output_dir.resolve()))
def process_sub_set(self, name: str, models_mat_file: str,
init_commands: List[dict],
scenes: List[Scene]) -> None:
"""
Process a sub-set of the complete dataset (e.g. all of Sound20K).
:param name: The name of the sub-set.
:param models_mat_file: The models per material data filename.
:param init_commands: The commands used to initialize the entire process (this is sent only once).
:param scenes: The scenes that can be loaded.
"""
print(name)
# Load models by wnid.
materials: Dict[str, List[str]] = loads(
Path(f"models/{models_mat_file}.json").read_text(encoding="utf-8"))
num_per_material = int(self.total / len(materials))
# Load the scene.
self.communicate(init_commands)
pbar = tqdm(total=self.total)
for material in materials:
pbar.set_description(material)
self.process_material(root_dir=self.output_dir.joinpath(name),
scenes=scenes,
material=material,
models=materials[material],
num_total=num_per_material,
pbar=pbar)
pbar.close()
def sound20k_set(self) -> None:
"""
Generate a dataset analogous to Sound20K.
"""
sound20k_init_commands = [
{
"$type": "load_scene"
},
TDWUtils.create_empty_room(12, 12), {
"$type": "set_proc_gen_walls_scale",
"walls": TDWUtils.get_box(12, 12),
"scale": {
"x": 1,
"y": 4,
"z": 1
}
}, {
"$type": "set_reverb_space_simple",
"env_id": 0,
"reverb_floor_material": "parquet",
"reverb_ceiling_material": "acousticTile",
"reverb_front_wall_material": "smoothPlaster",
"reverb_back_wall_material": "smoothPlaster",
"reverb_left_wall_material": "smoothPlaster",
"reverb_right_wall_material": "smoothPlaster"
}, {
"$type": "create_avatar",
"type": "A_Img_Caps_Kinematic",
"id": "a"
}, {
"$type": "add_environ_audio_sensor"
}, {
"$type": "toggle_image_sensor"
}
]
self.process_sub_set("Sound20K", "models_per_material_sound20k",
sound20k_init_commands, get_sound20k_scenes())
def tdw_set(self) -> None:
self.process_sub_set("TDW", "models_per_material_tdw", [],
get_tdw_scenes())
def process_material(self, root_dir: Path, scenes: List[Scene],
models: List[str], material: str, num_total: int,
pbar: Optional[tqdm]) -> None:
"""
Generate .wav files from all models with the material.
:param root_dir: The root output directory.
:param scenes: The scenes that a trial can use.
:param models: The names of the models in the category and their libraries.
:param num_total: The total number of files to generate for this category.
:param material: The name of the material.
:param pbar: The progress bar.
"""
num_images_per_model = int(num_total / len(models))
num_scenes_per_model = int(num_images_per_model / len(scenes))
# The number of files generated for the wnid.
count = 0
# The number of files generated for the current model.
model_count = 0
# The model being used to generate files.
model_index = 0
# The number of files for the current model that have used the current scene.
scene_count = 0
# The scene being used to generate files.
scene_index = 0
output_dir = root_dir.joinpath(material)
if not output_dir.exists():
output_dir.mkdir(parents=True)
while count < num_total:
obj_name = models[model_index]
filename = output_dir.joinpath(
TDWUtils.zero_padding(count, 4) + ".wav")
# Get the expected output path.
output_path = output_dir.joinpath(filename)
# Do a trial if the file doesn't exist yet.
if not output_path.exists():
try:
self.trial(
scene=scenes[scene_index],
record=self.libs[self.object_info[
models[model_index]].library].get_record(obj_name),
output_path=output_path,
scene_index=scene_index)
finally:
# Stop recording audio.
AudioUtils.stop()
count += 1
# Iterate through scenes.
scene_count += 1
if scene_count > num_scenes_per_model:
scene_id = scenes[scene_index].get_id()
# Add the scene to the database.
scene_db = self.db_c.execute("SELECT * FROM scenes WHERE id=?",
(scene_id, )).fetchone()
if scene_db is None:
self.db_c.execute(
"INSERT INTO scenes VALUES(?,?)",
(scene_id,
json.dumps(
scenes[scene_index].initialize_scene(self))))
self.conn.commit()
scene_index += 1
scene_count = 0
if scene_index >= len(scenes):
scene_index = 0
# Iterate through models.
model_count += 1
if model_count > num_images_per_model:
model_index += 1
if model_index >= len(models):
model_index = 0
model_count = 0
# If this is a new model, reset the scene count.
scene_index = 0
scene_count = 0
# Unload the asset bundles because we are done with this model.
self.communicate({"$type": "unload_asset_bundles"})
if pbar is not None:
pbar.update(1)
def trial(self, scene: Scene, record: ModelRecord, output_path: Path,
scene_index: int) -> None:
"""
Run a trial in a scene that has been initialized.
:param scene: Data for the current scene.
:param record: The model's metadata record.
:param output_path: Write the .wav file to this path.
:param scene_index: The scene identifier.
"""
self.py_impact.reset(initial_amp=0.05)
# Initialize the scene, positioning objects, furniture, etc.
resp = self.communicate(scene.initialize_scene(self))
center = scene.get_center(self)
max_y = scene.get_max_y()
# The object's initial position.
o_x = RNG.uniform(center["x"] - 0.15, center["x"] + 0.15)
o_y = RNG.uniform(max_y - 0.5, max_y)
o_z = RNG.uniform(center["z"] - 0.15, center["z"] + 0.15)
# Physics values.
mass = self.object_info[record.name].mass + RNG.uniform(
self.object_info[record.name].mass * -0.15,
self.object_info[record.name].mass * 0.15)
static_friction = RNG.uniform(0.1, 0.3)
dynamic_friction = RNG.uniform(0.7, 0.9)
# Angles of rotation.
yaw = RNG.uniform(-30, 30)
pitch = RNG.uniform(0, 45)
roll = RNG.uniform(-45, 45)
# The force applied to the object.
force = RNG.uniform(0, 5)
# The avatar's position.
a_r = RNG.uniform(1.5, 2.2)
a_x = center["x"] + a_r
a_y = RNG.uniform(1.5, 3)
a_z = center["z"] + a_r
cam_angle_min, cam_angle_max = scene.get_camera_angles()
theta = np.radians(RNG.uniform(cam_angle_min, cam_angle_max))
a_x = np.cos(theta) * (a_x - center["x"]) - np.sin(theta) * (
a_z - center["z"]) + center["x"]
a_z = np.sin(theta) * (a_x - center["x"]) + np.cos(theta) * (
a_z - center["z"]) + center["z"]
o_id = 0
# Create the object and apply a force.
commands = [{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": {
"x": o_x,
"y": o_y,
"z": o_z
},
"category": record.wcategory,
"id": o_id
}, {
"$type": "set_mass",
"id": o_id,
"mass": mass
}, {
"$type": "set_physic_material",
"id": o_id,
"bounciness": self.object_info[record.name].bounciness,
"static_friction": static_friction,
"dynamic_friction": dynamic_friction
}, {
"$type": "rotate_object_by",
"angle": yaw,
"id": o_id,
"axis": "yaw",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": pitch,
"id": o_id,
"axis": "pitch",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": roll,
"id": o_id,
"axis": "roll",
"is_world": True
}, {
"$type": "apply_force_magnitude_to_object",
"magnitude": force,
"id": o_id
}, {
"$type": "send_rigidbodies",
"frequency": "always"
}, {
"$type": "send_collisions",
"enter": True,
"exit": False,
"stay": False,
"collision_types": ["obj", "env"]
}, {
"$type": "send_transforms",
"frequency": "always"
}]
# Parse bounds data to get the centroid of all objects currently in the scene.
bounds = Bounds(resp[0])
if bounds.get_num() == 0:
look_at = {"x": center["x"], "y": 0.1, "z": center["z"]}
else:
centers = []
for i in range(bounds.get_num()):
centers.append(bounds.get_center(i))
centers_x, centers_y, centers_z = zip(*centers)
centers_len = len(centers_x)
look_at = {
"x": sum(centers_x) / centers_len,
"y": sum(centers_y) / centers_len,
"z": sum(centers_z) / centers_len
}
# Add the avatar.
# Set the position at a given distance (r) from the center of the scene.
# Rotate around that position to a random angle constrained by the scene's min and max angles.
commands.extend([{
"$type": "teleport_avatar_to",
"position": {
"x": a_x,
"y": a_y,
"z": a_z
}
}, {
"$type": "look_at_position",
"position": look_at
}])
# Send the commands.
resp = self.communicate(commands)
AudioUtils.start(output_path=output_path, until=(0, 10))
# Loop until all objects are sleeping.
done = False
while not done and AudioUtils.is_recording():
commands = []
collisions, environment_collisions, rigidbodies = PyImpact.get_collisions(
resp)
# Create impact sounds from object-object collisions.
for collision in collisions:
if PyImpact.is_valid_collision(collision):
# Get the audio material and amp.
collider_id = collision.get_collider_id()
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
collidee_id = collision.get_collider_id()
collidee_material, collidee_amp = self._get_object_info(
collidee_id, Scene.OBJECT_IDS, record.name)
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collidee_id,
target_amp=collidee_amp,
target_mat=collidee_material.name,
other_id=collider_id,
other_mat=collider_material.name,
other_amp=collider_amp,
play_audio_data=False)
commands.append(impact_sound_command)
# Create impact sounds from object-environment collisions.
for collision in environment_collisions:
collider_id = collision.get_object_id()
if self._get_velocity(rigidbodies, collider_id) > 0:
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
surface_material = scene.get_surface_material()
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collider_id,
target_amp=collider_amp,
target_mat=collider_material.name,
other_id=-1,
other_amp=0.01,
other_mat=surface_material.name,
play_audio_data=False)
commands.append(impact_sound_command)
# If there were no collisions, check for movement. If nothing is moving, the trial is done.
if len(commands) == 0:
transforms = AudioDataset._get_transforms(resp)
done = True
for i in range(rigidbodies.get_num()):
if self._is_moving(rigidbodies.get_id(i), transforms,
rigidbodies):
done = False
break
# Continue the trial.
if not done:
resp = self.communicate(commands)
# Stop listening for anything except audio data..
resp = self.communicate([{
"$type": "send_rigidbodies",
"frequency": "never"
}, {
"$type": "send_transforms",
"frequency": "never"
}, {
"$type": "send_collisions",
"enter": False,
"exit": False,
"stay": False,
"collision_types": []
}, {
"$type": "send_audio_sources",
"frequency": "always"
}])
# Wait for the audio to finish.
done = False
while not done and AudioUtils.is_recording():
done = True
for r in resp[:-1]:
if OutputData.get_data_type_id(r) == "audi":
audio_sources = AudioSources(r)
for i in range(audio_sources.get_num()):
if audio_sources.get_is_playing(i):
done = False
if not done:
resp = self.communicate([])
# Cleanup.
commands = [{
"$type": "send_audio_sources",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": o_id
}]
for scene_object_id in Scene.OBJECT_IDS:
commands.append({"$type": "destroy_object", "id": scene_object_id})
self.communicate(commands)
# Insert the trial's values into the database.
self.db_c.execute(
"INSERT INTO sound20k VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
(output_path.name, scene_index, a_x, a_y, a_z, o_x, o_y, o_z, mass,
static_friction, dynamic_friction, yaw, pitch, roll, force))
self.conn.commit()
def _get_object_info(self, o_id: int, object_ids: Dict[int, str],
drop_name: str) -> Tuple[AudioMaterial, float]:
"""
:param o_id: The object ID.
:param object_ids: The scene object IDs.
:param drop_name: The name of the dropped object.
:return: The audio material and amp associated with the object.
"""
if o_id in object_ids:
return self.object_info[object_ids[
o_id]].material, self.object_info[object_ids[o_id]].amp
else:
return self.object_info[drop_name].material, self.object_info[
drop_name].amp
@staticmethod
def _get_transforms(resp: List[bytes]) -> Transforms:
"""
:param resp: The output data response.
:return: Transforms data.
"""
for r in resp[:-1]:
if OutputData.get_data_type_id(r) == "tran":
return Transforms(r)
raise Exception("Transforms output data not found!")
@staticmethod
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))
@staticmethod
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
class Scene(ABC):
"""
A recipe to initialize a scene.
"""
_MODEL_LIBRARY_PATH = str(Path("models/models.json").resolve())
# A list of object IDs for the scene objects and the model names.
OBJECT_IDS: Dict[int, str] = {}
_OBJECT_INFO = PyImpact.get_object_info()
# Append custom data.
_custom_object_info = PyImpact.get_object_info(Path("models/object_info.csv"))
for obj in _custom_object_info:
_OBJECT_INFO.update({obj: _custom_object_info[obj]})
_SCENE_IDS: Dict[str, int] = {}
_SCENE_INDEX = 0
def initialize_scene(self, c: Controller) -> List[dict]:
"""
Add these commands to the beginning of the list of initialization commands.
:param c: The controller.
:return: A list of commands to initialize a scene.
"""
# Clean up all objects.
Scene.OBJECT_IDS.clear()
# Custom commands to initialize the scene.
commands = self._initialize_scene(c)[:]
# Send bounds data (for the new objects).
commands.append({"$type": "send_bounds",
"frequency": "once"})
return commands
@abstractmethod
def _initialize_scene(self, c: Controller) -> List[dict]:
"""
:param c: The controller.
:return: A list of commands to initialize a scene.
"""
raise Exception()
@abstractmethod
def get_center(self, c: Controller) -> Dict[str, float]:
"""
:param c: The controller.
:return: The "center" of the scene, as a Vector3 dictionary.
"""
raise Exception()
@abstractmethod
def get_max_y(self) -> float:
"""
:return: The maximum y value for the camera (avatar) and starting height for an object.
"""
raise Exception()
@staticmethod
def _init_object(c: Controller, name: str, pos: Dict[str, float], rot: Dict[str, float]) -> List[dict]:
"""
:param c: The controller.
:param name: The name of the model.
:param pos: The initial position of the model.
:param rot: The initial rotation of the model.
:return: A list of commands to instantiate an object from ObjectInfo values.
"""
o_id = c.get_unique_id()
Scene.OBJECT_IDS.update({o_id: name})
info = Scene._OBJECT_INFO[name]
return [c.get_add_object(name,
object_id=o_id,
position=pos,
rotation=rot,
library=info.library),
{"$type": "set_mass",
"id": o_id,
"mass": info.mass},
{"$type": "set_physic_material",
"id": o_id,
"bounciness": info.bounciness,
"static_friction": 0.1,
"dynamic_friction": 0.8}]
@staticmethod
def get_camera_angles() -> Tuple[float, float]:
"""
:return: Range of valid camera angles.
"""
return 0, 360
@abstractmethod
def get_surface_material(self) -> AudioMaterial:
"""
:return: The audio material of the surface.
"""
raise Exception()
def _get_name(self) -> str:
"""
:return: The name of this scene; used for indexing.
"""
return type(self).__name__
def get_id(self) -> int:
"""
:return: The unique ID of this type of scene.
"""
name = self._get_name()
# Index the name.
if name not in Scene._SCENE_IDS:
Scene._SCENE_IDS.update({name: Scene._SCENE_INDEX})
Scene._SCENE_INDEX += 1
return Scene._SCENE_IDS[name]
class FrameData:
"""
Per-frame data that an avatar can use to decide what action to do next.
Access this data from the [StickyMittenAvatarController](sma_controller.md):
```python
from sticky_mitten_avatar import StickyMittenAvatarController, Arm
c = StickyMittenAvatarController()
c.init_scene()
# Look towards the left arm.
c.rotate_camera_by(pitch=70, yaw=-45)
c.reach_for_target(target={"x": -0.2, "y": 0.21, "z": 0.385}, arm=Arm.left)
# Save each image from the start of the most recent API action to the end.
for frame in c.frames:
frame.save_images(output_directory="dist")
c.end()
```
***
## Fields
### Visual
- `image_pass` Rendered image of the scene as a numpy array.
- `id_pass` Image pass of object color segmentation as a numpy array.
- `depth_pass` Image pass of depth values per pixel as a numpy array.
- `projection_matrix` The [camera projection matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the avatar's camera as a numpy array.
- `camera_matrix` The [camera matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the avatar's camera as a numpy array.
### Audio
- `audio` A list of tuples of audio generated by impacts. The first element in the tuple is a [`Base64Sound` object](https://github.com/threedworld-mit/tdw/blob/master/Documentation/python/py_impact.md#base64sound).
The second element is the ID of the "target" (smaller) object.
```python
for frame in c.frames:
for audio in frame.audio:
# Get the audio data.
wav_data = audio[0].bytes
# Get the position of the object that generated the audio data.
object_id = audio[1]
position = frame.object_transforms[object_id].position
```
### Objects
- `object_transforms` The dictionary of object [transform data](transform.md). Key = the object ID.
```python
# Print the position of each object per frame.
for frame in c.frames:
for object_id in frame.object_transforms:
print(frame.object_transforms[object_id].position)
```
### Avatar
- `avatar_transform` The [transform data](transform.md) of the avatar.
```python
for frame in c.frames:
avatar_position = frame.avatar_transform.position
```
- `avatar_body_part_transforms` The [transform data](transform.md) of each body part of the avatar. Key = body part ID.
```python
for frame in c.frames:
# Get the position and segmentation color of each body part.
for body_part_id in frame.avatar_body_part_transforms:
position = frame.avatar_body_part_transforms[body_part_id]
segmentation_color = c.static_avatar_data[body_part_id]
```
- `avatar_object_collisions` A dictionary of objects the avatar collided with. Key = body part ID. Value = A list of object IDs.
```python
for frame in c.frames:
for body_part_id in frame.avatar_object_collisions:
body_part = c.static_avatar_info[body_part_id]
object_ids = frame.avatar_object_collisions[body_part_id]
for object_id in object_ids:
print(body_part.name + " collided with object " + str(object_id))
```
- `avatar_env_collisions` A list of body part IDs that collided with the environment (such as a wall).
```python
for frame in c.frames:
for body_part_id in frame.avatar_env_collisions:
body_part = c.static_avatar_info[body_part_id]
print(body_part.name + " collided with the environment.")
```
- `held_objects` A dictionary of IDs of objects held in each mitten. Key = arm:
```python
from sticky_mitten_avatar import StickyMittenAvatarController, Arm
c = StickyMittenAvatarController()
# Your code here.
# Prints all objects held by the left mitten at the last frame.
print(c.frames[-1].held_objects[Arm.left])
```
***
## Functions
"""
_P = PyImpact(initial_amp=0.01)
_SURFACE_MATERIAL: AudioMaterial = AudioMaterial.hardwood
def __init__(self, resp: List[bytes], objects: Dict[int, StaticObjectInfo],
avatar: Avatar):
"""
:param resp: The response from the build.
:param objects: Static object info per object. Key = the ID of the object in the scene.
:param avatar: The avatar in the scene.
"""
self._frame_count = Controller.get_frame(resp[-1])
self.audio: List[Tuple[Base64Sound, int]] = list()
collisions, env_collisions, rigidbodies = FrameData._P.get_collisions(
resp=resp)
# Record avatar collisions.
if avatar is not None:
self.avatar_object_collisions = avatar.collisions
self.avatar_env_collisions = avatar.env_collisions
self.held_objects = {
Arm.left: avatar.frame.get_held_left(),
Arm.right: avatar.frame.get_held_right()
}
else:
self.avatar_object_collisions = None
self.avatar_env_collisions = None
self.held_objects = None
# Get the object transform data.
self.object_transforms: Dict[int, Transform] = dict()
tr = get_data(resp=resp, d_type=Transforms)
for i in range(tr.get_num()):
o_id = tr.get_id(i)
self.object_transforms[o_id] = Transform(
position=np.array(tr.get_position(i)),
rotation=np.array(tr.get_rotation(i)),
forward=np.array(tr.get_forward(i)))
# Get camera matrix data.
matrices = get_data(resp=resp, d_type=CameraMatrices)
self.projection_matrix = matrices.get_projection_matrix()
self.camera_matrix = matrices.get_camera_matrix()
# Get the transform data of the avatar.
self.avatar_transform = Transform(
position=np.array(avatar.frame.get_position()),
rotation=np.array(avatar.frame.get_rotation()),
forward=np.array(avatar.frame.get_forward()))
self.avatar_body_part_transforms: Dict[int, Transform] = dict()
for i in range(avatar.frame.get_num_body_parts()):
self.avatar_body_part_transforms[avatar.frame.get_body_part_id(
i)] = Transform(
position=np.array(avatar.frame.get_body_part_position(i)),
rotation=np.array(avatar.frame.get_body_part_rotation(i)),
forward=np.array(avatar.frame.get_body_part_forward(i)))
# Get the audio of each collision.
for coll in collisions:
if not FrameData._P.is_valid_collision(coll):
continue
collider_id = coll.get_collider_id()
collidee_id = coll.get_collidee_id()
collider_info: Optional[ObjectInfo] = None
collidee_info: Optional[ObjectInfo] = None
if collider_id in objects:
collider_info = objects[collider_id].audio
# Check if the object is a body part.
else:
if collider_id in avatar.body_parts_static:
collider_info = avatar.body_parts_static[collider_id].audio
if collidee_id in objects:
collidee_info = objects[collidee_id].audio
# Check if the object is a body part.
else:
if collidee_id in avatar.body_parts_static:
collidee_info = avatar.body_parts_static[collidee_id].audio
# If either object isn't a cached object, don't try to add audio.
if collider_info is None or collidee_info is None:
continue
if collider_info.mass < collidee_info.mass:
target_id = collider_id
target_amp = collider_info.amp
target_mat = collider_info.material.name
other_id = collidee_id
other_amp = collidee_info.amp
other_mat = collider_info.material.name
else:
target_id = collidee_id
target_amp = collidee_info.amp
target_mat = collidee_info.material.name
other_id = collider_id
other_amp = collider_info.amp
other_mat = collider_info.material.name
rel_amp = other_amp / target_amp
audio = FrameData._P.get_sound(coll, rigidbodies, other_id,
other_mat, target_id, target_mat,
rel_amp)
self.audio.append((audio, target_id))
# Get the audio of each environment collision.
for coll in env_collisions:
collider_id = coll.get_object_id()
if collider_id not in objects:
continue
v = FrameData._get_velocity(rigidbodies, collider_id)
if (v is not None) and (v > 0):
collider_info = objects[collider_id].audio
audio = FrameData._P.get_sound(
coll, rigidbodies, 1, FrameData._SURFACE_MATERIAL.name,
collider_id, collider_info.material.name, 0.01)
self.audio.append((audio, collider_id))
# Get the image data.
self.id_pass: Optional[np.array] = None
self.depth_pass: Optional[np.array] = None
self.image_pass: Optional[np.array] = None
for i in range(0, len(resp) - 1):
if OutputData.get_data_type_id(resp[i]) == "imag":
images = Images(resp[i])
for j in range(images.get_num_passes()):
if images.get_pass_mask(j) == "_id":
self.id_pass = images.get_image(j)
elif images.get_pass_mask(j) == "_depth_simple":
self.depth_pass = images.get_image(j)
elif images.get_pass_mask(j) == "_img":
self.image_pass = images.get_image(j)
@staticmethod
def set_surface_material(surface_material: AudioMaterial) -> None:
"""
Set the surface material of the scene.
:param surface_material: The floor's [audio material](https://github.com/threedworld-mit/tdw/blob/master/Documentation/python/py_impact.md#audiomaterialenum).
"""
FrameData._SURFACE_MATERIAL = surface_material
def save_images(self, output_directory: Union[str, Path]) -> None:
"""
Save the ID pass (segmentation colors) and the depth pass to disk.
Images will be named: `[frame_number]_[pass_name].[extension]`
For example, the depth pass on the first frame will be named: `00000000_depth.png`
The image pass is a jpg file and the other passes are png files.
:param output_directory: The directory that the images will be saved to.
"""
if isinstance(output_directory, str):
output_directory = Path(output_directory)
if not output_directory.exists():
output_directory.mkdir(parents=True)
prefix = TDWUtils.zero_padding(self._frame_count, 8)
# Save each image.
for image, pass_name, ext in zip(
[self.image_pass, self.id_pass, self.depth_pass],
["img", "id", "depth"], ["jpg", "png", "png"]):
p = output_directory.joinpath(f"{prefix}_{pass_name}.{ext}")
with p.open("wb") as f:
f.write(image)
def get_pil_images(self) -> dict:
"""
Convert each image pass to PIL images.
:return: A dictionary of PIL images. Key = the name of the pass (img, id, depth)
"""
print(type(Image.open(BytesIO(self.image_pass))))
return {
"img": Image.open(BytesIO(self.image_pass)),
"id": Image.open(BytesIO(self.id_pass)),
"depth": Image.open(BytesIO(self.depth_pass))
}
@staticmethod
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))