def add_transforms_object(self,
record: ModelRecord,
position: Dict[str, float],
rotation: Dict[str, float],
o_id: Optional[int] = None) -> dict:
"""
This is a wrapper for `Controller.get_add_object()` and the `add_object` command.
This caches the ID of the object so that it can be easily cleaned up later.
:param record: The model record.
:param position: The initial position of the object.
:param rotation: The initial rotation of the object, in Euler angles.
:param o_id: The unique ID of the object. If None, a random ID is generated.
:return: An `add_object` command.
"""
if o_id is None:
o_id: int = Controller.get_unique_id()
# Log the static data.
self.object_ids = np.append(self.object_ids, o_id)
return {
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": position,
"rotation": rotation,
"category": record.wcategory,
"id": o_id
}
def get_image(self, record: ModelRecord):
o_id = Controller.get_unique_id()
self.communicate({"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"rotation": record.canonical_rotation,
"id": o_id})
s = TDWUtils.get_unit_scale(record) * 2
# Scale the model and get an image.
# Look at the model's centroid.
resp = self.communicate([{"$type": "scale_object",
"id": o_id,
"scale_factor": {"x": s, "y": s, "z": s}},
{"$type": "look_at",
"avatar_id": "a",
"object_id": o_id,
"use_centroid": True}])
# Destroy the model and unload the asset bundle.
self.communicate([{"$type": "destroy_object",
"id": o_id},
{"$type": "unload_asset_bundles"}])
return Images(resp[0]), resp[-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 process_model(self, record: ModelRecord, a: str, envs: list,
train_count: int, val_count: int, root_dir: str,
wnid: str) -> float:
"""
Capture images of a model.
:param record: The model record.
:param a: The ID of the avatar.
:param envs: All environment data.
:param train_count: Number of train images.
:param val_count: Number of val images.
:param root_dir: The root directory for saving images.
:param wnid: The wnid of the record.
:return The time elapsed.
"""
image_count = 0
# Get the filename index. If we shouldn't overwrite any images, start after the last image.
if self.no_overwrite:
# Check if any images exist.
wnid_dir = Path(root_dir).joinpath(f"train/{wnid}")
if wnid_dir.exists():
max_file_index = -1
for image in wnid_dir.iterdir():
if not image.is_file() or image.suffix != ".jpg" \
or not image.stem.startswith("img_") or image.stem[4:-5] != record.name:
continue
image_index = int(image.stem[-4:])
if image_index > max_file_index:
max_file_index = image_index
file_index = max_file_index + 1
else:
file_index = 0
else:
file_index = 0
image_positions = []
o_id = self.get_unique_id()
s = TDWUtils.get_unit_scale(record)
# Add the object.
# Set the screen size to 32x32 (to make the build run faster; we only need the average grayscale values).
# Toggle off pass masks.
# Set render quality to minimal.
# Scale the object to "unit size".
self.communicate([{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"category": record.wcategory,
"id": o_id
}, {
"$type": "set_screen_size",
"height": 32,
"width": 32
}, {
"$type": "set_pass_masks",
"avatar_id": a,
"pass_masks": []
}, {
"$type": "set_render_quality",
"render_quality": 0
}, {
"$type": "scale_object",
"id": o_id,
"scale_factor": {
"x": s,
"y": s,
"z": s
}
}])
# The index in the HDRI records array.
hdri_index = 0
# The number of iterations on this skybox so far.
skybox_count = 0
if self.skyboxes:
# The number of iterations per skybox for this model.
its_per_skybox = round(
(train_count + val_count) / len(self.skyboxes))
# Set the first skybox.
hdri_index, skybox_count, command = self.set_skybox(
self.skyboxes, its_per_skybox, hdri_index, skybox_count)
self.communicate(command)
else:
its_per_skybox = 0
while len(image_positions) < train_count + val_count:
e = RNG.choice(envs)
# Get the real grayscale.
g_r, d, a_p, o_p, o_rot, cam_rot = self.get_real_grayscale(
o_id, a, e)
if g_r > 0:
# Get the optimal grayscale.
g_o = self.get_optimal_grayscale(o_id, a, o_p, a_p)
if g_o > 0 and g_r / g_o > self.grayscale_threshold:
# Cache the position.
image_positions.append(
ImagePosition(a_p, cam_rot, o_p, o_rot))
# Send images.
# Set the screen size.
# Set render quality to maximum.
commands = [{
"$type": "send_images",
"frequency": "always"
}, {
"$type": "set_pass_masks",
"avatar_id": a,
"pass_masks": ["_img", "_id"] if self.id_pass else ["_img"]
}, {
"$type": "set_screen_size",
"height": self.screen_size,
"width": self.screen_size
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
# Hide the object maybe.
if not self.show_objects:
commands.append({"$type": "hide_object", "id": o_id})
self.communicate(commands)
t0 = time()
# Generate images from the cached spatial data.
train = 0
for p in image_positions:
# Teleport the avatar.
# Rotate the avatar's camera.
# Teleport the object.
# Rotate the object.
# Get the response.
commands = [{
"$type": "teleport_avatar_to",
"avatar_id": a,
"position": p.avatar_position
}, {
"$type": "rotate_sensor_container_to",
"avatar_id": a,
"rotation": p.camera_rotation
}, {
"$type": "teleport_object",
"id": o_id,
"position": p.object_position
}, {
"$type": "rotate_object_to",
"id": o_id,
"rotation": p.object_rotation
}]
# Set the visual materials.
if self.materials is not None:
if record.name not in self.substructures:
self.substructures.update(
{record.name: record.substructure})
for sub_object in self.substructures[record.name]:
for i in range(
len(self.substructures[record.name][
sub_object["name"]])):
material_name = self.materials[RNG.randint(
0, len(self.materials))].name
commands.extend([
self.get_add_material(material_name), {
"$type": "set_visual_material",
"id": o_id,
"material_name": material_name,
"object_name": sub_object["name"],
"material_index": i
}
])
# Maybe set a new skybox.
# Rotate the skybox.
if self.skyboxes:
hdri_index, skybox_count, command = self.set_skybox(
self.skyboxes, its_per_skybox, hdri_index, skybox_count)
if command:
commands.append(command)
commands.append({
"$type": "rotate_hdri_skybox_by",
"angle": RNG.uniform(0, 360)
})
resp = self.communicate(commands)
train += 1
# Create a thread to save the image.
t = Thread(target=self.save_image,
args=(resp, record, file_index, root_dir, wnid, train,
train_count))
t.daemon = True
t.start()
file_index += 1
image_count += 1
t1 = time()
# Stop sending images.
# Destroy the object.
# Unload asset bundles.
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": o_id
}, {
"$type": "unload_asset_bundles"
}])
return t1 - t0