def run(self):
# Load the streamed scene and add controller rig.
self.load_streamed_scene(scene="tdw_room")
self.communicate({"$type": "create_vr_rig"})
# Add the table object and make it kinematic.
table_id = self.add_object("small_table_green_marble",
position={"x": 0, "y": 0, "z": 0.5})
self.communicate([{"$type": "set_kinematic_state",
"id": table_id,
"is_kinematic": True}])
# Add a box object and make it graspable.
graspable_id_box = self.add_object("woven_box",
position={"x": 0.2, "y": 1.0, "z": 0.5},
library="models_core.json")
self.communicate([{"$type": "set_graspable",
"id": graspable_id_box}])
# Add the ball object and make it graspable.
graspable_id = self.add_object("prim_sphere",
position={"x": 0.2, "y": 3.0, "z": 0.5},
library="models_special.json")
self.communicate([{"$type": "set_graspable",
"id": graspable_id},
{"$type": "scale_object",
"scale_factor": {"x": 0.2, "y": 0.2, "z": 0.2},
"id": graspable_id}])
# Receive VR data.
resp = self.communicate({"$type": "send_vr_rig",
"frequency": "once"})
assert len(resp) > 1 and OutputData.get_data_type_id(resp[0]) == "vrri"
vr = VRRig(resp[0])
print(vr.get_head_rotation())
# Start an infinite loop to allow the build to simulate physics.
while True:
self.communicate({"$type": "do_nothing"})
def get_object_target_collision(self, obj_id: int, target_id: int,
resp: List[bytes]):
contact_points = []
contact_normals = []
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "coll":
co = Collision(r)
coll_ids = [co.get_collider_id(), co.get_collidee_id()]
if [obj_id, target_id] == coll_ids or [target_id, obj_id
] == coll_ids:
contact_points = [
co.get_contact_point(i)
for i in range(co.get_num_contacts())
]
contact_normals = [
co.get_contact_normal(i)
for i in range(co.get_num_contacts())
]
return (contact_points, contact_normals)
def communicate(self, commands: Union[dict, List[dict]]) -> List[bytes]:
"""
See `Magnebot.communicate()`.
Images are saved per-frame.
"""
resp = super().communicate(commands=commands)
if not self._debug:
# Save all images.
got_images = False
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "imag":
got_images = True
images = Images(resp[i])
avatar_id = images.get_avatar_id()
if avatar_id in self.image_directories:
TDWUtils.save_images(filename=TDWUtils.zero_padding(self._image_count, 8),
output_directory=self.image_directories[avatar_id],
images=images)
if got_images:
self._image_count += 1
return resp
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
# Add a camera to the scene.
commands.extend(
TDWUtils.create_avatar(position={
"x": -2.49,
"y": 4,
"z": 0
},
look_at={
"x": 0,
"y": 0,
"z": 0
}))
resp = c.communicate(commands)
wheel_ids = []
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "srob":
sr = StaticRobot(resp[i])
for j in range(sr.get_num_joints()):
joint_id = sr.get_joint_id(j)
joint_name = sr.get_joint_name(j)
print(joint_name, joint_id)
# Find all of the wheels.
if "wheel" in joint_name:
wheel_ids.append(joint_id)
# Move the wheels forward.
commands = []
for wheel_id in wheel_ids:
commands.append({
"$type": "set_revolute_target",
probe_material=args.pmaterial,
zone_material=args.zmaterial,
zone_color=args.zcolor,
zone_friction=args.zfriction,
distractor_types=args.distractor,
distractor_categories=args.distractor_categories,
num_distractors=args.num_distractors,
occluder_types=args.occluder,
occluder_categories=args.occluder_categories,
num_occluders=args.num_occluders,
flex_only=args.only_use_flex_objects,
no_moving_distractors=args.no_moving_distractors,
use_test_mode_colors=args.use_test_mode_colors
)
if bool(args.run):
DC.run(num=args.num,
output_dir=args.dir,
temp_path=args.temp,
width=args.width,
height=args.height,
write_passes=args.write_passes.split(','),
save_passes=args.save_passes.split(','),
save_movies=args.save_movies,
save_labels=args.save_labels,
save_meshes=args.save_meshes,
args_dict=vars(args))
else:
end = DC.communicate({"$type": "terminate"})
print([OutputData.get_data_type_id(r) for r in end])
def _write_frame(self, frames_grp: h5py.Group, resp: List[bytes], frame_num: int) -> \
Tuple[h5py.Group, h5py.Group, dict, bool]:
num_objects = len(self.object_ids)
# Create a group for this frame.
frame = frames_grp.create_group(TDWUtils.zero_padding(frame_num, 4))
# Create a group for images.
images = frame.create_group("images")
# Transforms data.
positions = np.empty(dtype=np.float32, shape=(num_objects, 3))
forwards = np.empty(dtype=np.float32, shape=(num_objects, 3))
rotations = np.empty(dtype=np.float32, shape=(num_objects, 4))
camera_matrices = frame.create_group("camera_matrices")
# Parse the data in an ordered manner so that it can be mapped back to the object IDs.
tr_dict = dict()
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "tran":
tr = Transforms(r)
for i in range(tr.get_num()):
pos = tr.get_position(i)
tr_dict.update({
tr.get_id(i): {
"pos": pos,
"for": tr.get_forward(i),
"rot": tr.get_rotation(i)
}
})
# Add the Transforms data.
for o_id, i in zip(self.object_ids, range(num_objects)):
if o_id not in tr_dict:
continue
positions[i] = tr_dict[o_id]["pos"]
forwards[i] = tr_dict[o_id]["for"]
rotations[i] = tr_dict[o_id]["rot"]
elif r_id == "imag":
im = Images(r)
# Add each image.
for i in range(im.get_num_passes()):
images.create_dataset(im.get_pass_mask(i),
data=im.get_image(i),
compression="gzip")
# Add the camera matrices.
elif OutputData.get_data_type_id(r) == "cama":
matrices = CameraMatrices(r)
camera_matrices.create_dataset(
"projection_matrix", data=matrices.get_projection_matrix())
camera_matrices.create_dataset(
"camera_matrix", data=matrices.get_camera_matrix())
objs = frame.create_group("objects")
objs.create_dataset("positions",
data=positions.reshape(num_objects, 3),
compression="gzip")
objs.create_dataset("forwards",
data=forwards.reshape(num_objects, 3),
compression="gzip")
objs.create_dataset("rotations",
data=rotations.reshape(num_objects, 4),
compression="gzip")
return frame, objs, tr_dict, False
def on_frame(self, resp: List[bytes]) -> List[dict]:
"""
Update the avatar based on its current arm-bending goals and its state.
If the avatar has achieved a goal (for example, picking up an object), it will stop moving that arm.
Update the avatar's state as needed.
:param resp: The response from the build.
:return: A list of commands to pick up, stop moving, etc.
"""
def _get_mitten_position(a: Arm) -> np.array:
"""
:param a: The arm.
:return: The position of a mitten.
"""
if a == Arm.left:
return np.array(frame.get_mitten_center_left_position())
else:
return np.array(frame.get_mitten_center_right_position())
# Update dynamic data.
frame = self._get_frame(resp=resp)
# Update dynamic collision data.
self.collisions.clear()
self.env_collisions.clear()
# Get each collision.
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "coll":
coll = Collision(resp[i])
collider_id = coll.get_collider_id()
collidee_id = coll.get_collidee_id()
# Check if this was a mitten, if we're supposed to stop if there's a collision,
# and if the collision was not with the target.
for arm in self._ik_goals:
if self._ik_goals[arm] is not None:
if (collider_id == self.mitten_ids[arm] or collidee_id == self.mitten_ids[arm]) and \
(collider_id not in frame.get_held_right() and
collider_id not in frame.get_held_left() and
collidee_id not in frame.get_held_right() and
collidee_id not in frame.get_held_left()) and \
self._ik_goals[arm].stop_on_mitten_collision and \
(self._ik_goals[arm].target is None or
(self._ik_goals[arm].pick_up_id != collidee_id and
self._ik_goals[arm].pick_up_id != collider_id)) and \
(collidee_id not in self.body_parts_static or
collider_id not in self.body_parts_static):
self.status = TaskStatus.mitten_collision
self._ik_goals[arm] = None
if self._debug:
print(
"Stopping because the mitten collided with something."
)
return self._stop_arm(arm=arm)
# Check if the collision includes a body part.
if collider_id in self.body_parts_static and collidee_id not in self.body_parts_static:
if collider_id not in self.collisions:
self.collisions[collider_id] = []
self.collisions[collider_id].append(collidee_id)
elif collidee_id in self.body_parts_static and collider_id not in self.body_parts_static:
if collidee_id not in self.collisions:
self.collisions[collidee_id] = []
self.collisions[collidee_id].append(collider_id)
elif r_id == "enco":
coll = EnvironmentCollision(resp[i])
collider_id = coll.get_object_id()
if collider_id in self.body_parts_static:
self.env_collisions.append(collider_id)
# Check if IK goals are done.
temp_goals: Dict[Arm, Optional[_IKGoal]] = dict()
# Get commands for the next frame.
commands: List[dict] = []
for arm in self._ik_goals:
# No IK goal on this arm.
if self._ik_goals[arm] is None:
temp_goals[arm] = None
# This is a dummy IK goal. Let it run.
elif self._ik_goals[arm].target is None:
temp_goals[arm] = self._ik_goals[arm]
else:
# Is the arm at the target?
mitten_position = _get_mitten_position(arm)
# If we're not trying to pick something up, check if we are at the target position.
if self._ik_goals[arm].pick_up_id is None:
# If we're at the position, stop.
d = np.linalg.norm(mitten_position -
self._ik_goals[arm].target)
if d < self._ik_goals[arm].precision:
if self._debug:
print(
f"{arm.name} mitten is at target position {self._ik_goals[arm].target}. Stopping."
)
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self.status = TaskStatus.success
# Keep bending the arm.
else:
temp_goals[arm] = self._ik_goals[arm]
self._ik_goals[arm].previous_distance = d
# If we're trying to pick something, check if it was picked up on the previous frame.
else:
if self._ik_goals[arm].pick_up_id in frame.get_held_left() or self._ik_goals[arm]. \
pick_up_id in frame.get_held_right():
if self._debug:
print(
f"{arm.name} mitten picked up {self._ik_goals[arm].pick_up_id}. Stopping."
)
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self.status = TaskStatus.success
# Keep bending the arm and trying to pick up the object.
else:
commands.extend([{
"$type":
"pick_up_proximity",
"distance":
0.02,
"radius":
0.05,
"grip":
1000,
"is_left":
arm == Arm.left,
"avatar_id":
self.id,
"object_ids": [self._ik_goals[arm].pick_up_id]
}, {
"$type":
"pick_up",
"grip":
1000,
"is_left":
arm == Arm.left,
"object_ids": [self._ik_goals[arm].pick_up_id],
"avatar_id":
self.id
}])
temp_goals[arm] = self._ik_goals[arm]
self._ik_goals = temp_goals
# Check if the arms are still moving.
temp_goals: Dict[Arm, Optional[_IKGoal]] = dict()
for arm in self._ik_goals:
# No IK goal on this arm.
if self._ik_goals[arm] is None:
temp_goals[arm] = None
else:
# Get the past and present angles.
if arm == Arm.left:
angles_0 = self.frame.get_angles_left()
angles_1 = frame.get_angles_left()
else:
angles_0 = self.frame.get_angles_right()
angles_1 = frame.get_angles_right()
# Try to stop any moving joints.
if self._ik_goals[arm].rotations is not None and self._ik_goals[
arm].pick_up_id is not None:
joint_profile = self._get_default_sticky_mitten_profile()
for angle, joint_name in zip(angles_1, Avatar.ANGLE_ORDER):
target_angle = self._ik_goals[arm].rotations[
joint_name]
# Check if the joint stopped moving. Ignore if the joint already stopped.
if target_angle > 0.01 and np.abs(angle - target_angle) < 0.01 and \
joint_name in self._ik_goals[arm].moving_joints:
self._ik_goals[arm].moving_joints.remove(
joint_name)
j = joint_name.split("_")
j_name = f"{j[0]}_{arm.name}"
axis = j[1]
# Set the name of the elbow to the expected profile key.
if "elbow" in joint_name:
profile_key = "elbow"
else:
profile_key = joint_name
if self._debug:
print(
f"{joint_name} {arm.name} slowing down: {np.abs(angle - target_angle)}"
)
# Stop the joint from moving any more.
# Set the damper, force, and angular drag to "default" (non-moving) values.
commands.extend([{
"$type":
"set_joint_damper",
"joint":
j_name,
"axis":
axis,
"damper":
joint_profile[profile_key]["damper"],
"avatar_id":
self.id
}, {
"$type":
"set_joint_force",
"joint":
j_name,
"axis":
axis,
"force":
joint_profile[profile_key]["force"],
"avatar_id":
self.id
}, {
"$type":
"set_joint_angular_drag",
"joint":
j_name,
"axis":
axis,
"angular_drag":
joint_profile[profile_key]["angular_drag"],
"avatar_id":
self.id
}])
# Is any joint still moving?
moving = False
for a0, a1 in zip(angles_0, angles_1):
if np.abs(a0 - a1) > 0.03:
moving = True
break
# Keep moving.
if moving:
temp_goals[arm] = self._ik_goals[arm]
else:
if self._ik_goals[arm].rotations is not None:
# This is a reset arm action.
if self._ik_goals[arm].target is None:
mitten_position = _get_mitten_position(
arm) - frame.get_position()
d = np.linalg.norm(
self._initial_mitten_positions[arm] -
mitten_position)
# The reset arm action ended with the mitten very close to the initial position.
if d < self._ik_goals[arm].precision:
self.status = TaskStatus.success
else:
self.status = TaskStatus.no_longer_bending
# This is a regular action.
# It ended with the arm no longer moving but having never reached the target.
else:
if self._debug:
print(
f"{arm.name} is no longer bending. Cancelling."
)
self.status = TaskStatus.no_longer_bending
commands.extend(self._stop_arm(arm=arm))
temp_goals[arm] = None
self._ik_goals = temp_goals
self.frame = frame
return commands
def run(self):
# Create the scene.
# Set image encoding globals.
self.start()
commands = [
TDWUtils.create_empty_room(12, 12), {
"$type": "set_screen_size",
"width": 128,
"height": 128
}, {
"$type": "set_img_pass_encoding",
"value": False
}
]
# Create the avatar.
commands.extend(
TDWUtils.create_avatar(position={
"x": 2.478,
"y": 1.602,
"z": 1.412
},
look_at=TDWUtils.VECTOR3_ZERO,
avatar_id="a"))
# Enable all pass masks. Request an image for this frame only.
commands.extend([{
"$type":
"set_pass_masks",
"pass_masks": [
"_img", "_id", "_category", "_mask", "_depth", "_normals",
"_flow", "_depth_simple"
],
"avatar_id":
"a"
}, {
"$type": "send_images",
"ids": ["a"],
"frequency": "once"
}])
# Add objects.
commands.append(
self.get_add_object("small_table_green_marble",
position=TDWUtils.VECTOR3_ZERO,
rotation=TDWUtils.VECTOR3_ZERO,
object_id=0))
commands.append(
self.get_add_object("rh10",
position={
"x": 0.7,
"y": 0,
"z": 0.4
},
rotation={
"x": 0,
"y": 30,
"z": 0
},
object_id=1))
commands.append(
self.get_add_object("jug01",
position={
"x": -0.3,
"y": 0.9,
"z": 0.2
},
rotation=TDWUtils.VECTOR3_ZERO,
object_id=3))
commands.append(
self.get_add_object("jug05",
position={
"x": 0.3,
"y": 0.9,
"z": -0.2
},
rotation=TDWUtils.VECTOR3_ZERO,
object_id=4))
# Send the commands.
resp = self.communicate(commands)
# Save the images.
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "imag":
# Save the images.
TDWUtils.save_images(output_directory="dist",
images=Images(r),
filename="0")
print(f"Images saved to: {Path('dist').resolve()}")
# Stop the build.
self.communicate({"$type": "terminate"})
def run(self):
self.start()
object_id = self.get_unique_id()
resp = self.communicate([
TDWUtils.create_empty_room(12, 12),
{
"$type": "create_flex_container",
"particle_size": 0.1,
"collision_distance": 0.025,
"solid_rest": 0.1
},
# {"$type": "create_flex_container",
# "collision_distance": 0.001,
# "static_friction": 1.0,
# "dynamic_friction": 1.0,
# "radius": 0.1875,
# "max_particles": 200000},
self.get_add_object(model_name="linbrazil_diz_armchair",
position={
"x": 0.0,
"y": 2.0,
"z": 0.0
},
rotation={
"x": 25.0,
"y": 45.0,
"z": -40.0
},
object_id=object_id),
{
"$type": "set_kinematic_state",
"id": object_id
},
{
"$type": "set_flex_solid_actor",
"id": object_id,
"mesh_expansion": 0,
"particle_spacing": 0.025,
"mass_scale": 1
},
# {"$type": "set_flex_soft_actor",
# "id": object_id,
# "skinning_falloff": 0.5,
# "volume_sampling": 1.0,
# "mass_scale": 1.0,
# "cluster_stiffness": 0.2,
# "cluster_spacing": 0.2,
# "cluster_radius": 0.2,
# "link_radius": 0,
# "link_stiffness": 1.0,
# "particle_spacing": 0.025},
{
"$type": "assign_flex_container",
"id": object_id,
"container_id": 0
},
{
'$type': 'load_primitive_from_resources',
'id': 2,
'primitive_type': 'Cube',
'position': {
'x': 0,
'y': 0.5,
'z': 0
},
'orientation': {
'x': 0,
'y': 0,
'z': 0
}
},
{
"$type": "scale_object",
"id": 2,
"scale_factor": {
"x": 0.5,
"y": 0.5,
"z": 0.5
}
},
{
'$type': 'set_kinematic_state',
'id': 2,
'is_kinematic': False
},
{
"$type": "set_flex_solid_actor",
"id": 2,
"mesh_expansion": 0,
"particle_spacing": 0.025,
"mass_scale": 1
},
# {'$type': 'set_flex_soft_actor',
# 'id': 2,
# 'draw_particles': False,
# 'particle_spacing': 0.125,
# 'cluster_stiffness': 0.22055267521432875},
{
'$type': 'assign_flex_container',
'id': 2,
'container_id': 0
},
{
'$type': 'set_flex_particles_mass',
'id': 2,
'mass': 15.625
},
{
"$type": "send_flex_particles",
"frequency": "always"
},
{
"$type": "send_collisions",
"enter": True,
"stay": True,
"exit": True,
"collision_types": ["obj", "env"]
}
])
for i in range(100):
for j in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[j])
if r_id == "flex":
flex = FlexParticles(resp[j])
print(i, "flex", flex.get_id(0))
elif r_id == "enco":
enco = EnvironmentCollision(resp[j])
print(i, "enco", enco.get_state())
if r_id == "coll":
coll = Collision(resp[j])
print(i, "coll", coll.get_state())
if coll.get_state() == 'enter':
print("BAM!")
resp = self.communicate([])
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 run(self):
self.start()
self.communicate(TDWUtils.create_empty_room(12, 12))
# Find the local asset bundle for this platform.
url = "file:///" + str(
Path("composite_objects/" + platform.system() +
"/test_composite_object").resolve())
# Add the local object.
o_id = self.get_unique_id()
self.communicate([{
"$type": "add_object",
"name": "test_composite_object",
"url": url,
"scale_factor": 1,
"id": o_id
}, {
"$type": "set_mass",
"id": o_id,
"mass": 100
}])
self.communicate(
TDWUtils.create_avatar(position={
"x": 0,
"y": 1.49,
"z": -2.77
}))
# Test that each sub-object has a unique segmentation color.
resp = self.communicate({
"$type": "send_segmentation_colors",
"ids": []
})
colors = SegmentationColors(resp[0])
segmentation_colors = []
# There are 4 objects: The parent object, the motor, the "base" of the motor, and the light.
assert colors.get_num() == 4, colors.get_num()
print("Segmentation colors:")
for i in range(colors.get_num()):
print(colors.get_object_id(i), colors.get_object_color(i))
# Cache the color for the next test.
segmentation_colors.append(colors.get_object_color(i))
# Test that each sub-object is visible to the avatar.
resp = self.communicate({
"$type": "send_id_pass_segmentation_colors",
"frequency": "once"
})
colors = IdPassSegmentationColors(resp[0])
# There are three visible objects (the light won't be visible in the _id pass).
assert colors.get_num_segmentation_colors() == 3
print("\nObserved colors:")
# Test that the colors observed by the avatar are in the cache.
for i in range(colors.get_num_segmentation_colors()):
observed_color = colors.get_segmentation_color(i)
assert observed_color in segmentation_colors
print(observed_color)
# Get composite objects data.
resp = self.communicate({"$type": "send_composite_objects"})
assert len(resp) > 1
assert OutputData.get_data_type_id(resp[0]) == "comp"
o = CompositeObjects(resp[0])
assert o.get_num() == 1
# There are 3 sub-objects: The motor, the "base" of the motor, and the light.
assert o.get_num_sub_objects(0) == 3
print("\nCompositeObjects: ")
commands = []
lights = []
# Iterate through each sub-object.
for s in range(o.get_num_sub_objects(0)):
sub_object_id = o.get_sub_object_id(0, s)
s_type = o.get_sub_object_machine_type(0, s)
print(sub_object_id, s_type)
# Add commands depending on the type of sub-object.
if s_type == "motor":
# Start the motor.
commands.append({
"$type": "set_motor",
"target_velocity": 500,
"force": 500,
"id": sub_object_id
})
elif s_type == "light":
commands.append({
"$type": "set_sub_object_light",
"is_on": True,
"id": sub_object_id
})
lights.append(sub_object_id)
self.communicate(commands)
is_on = True
for i in range(1000):
commands = []
# Every 50 frames, blink the lights on and off.
if i % 50 == 0:
is_on = not is_on
for light in lights:
commands.append({
"$type": "set_sub_object_light",
"is_on": is_on,
"id": light
})
else:
commands.append({"$type": "do_nothing"})
self.communicate(commands)
def run(self):
o_id = 0
self.start()
# 1. Create a room.
# 2. Create a Flex container.
# 3. Add a cloth object and set the cloth actor.
# 4. Request Flex particle output data (this frame only).
commands = [
TDWUtils.create_empty_room(12, 12), {
"$type": "create_flex_container",
"collision_distance": 0.001,
"static_friction": 1.0,
"dynamic_friction": 1.0,
"iteration_count": 12,
"substep_count": 12,
"radius": 0.1875,
"damping": 0,
"drag": 0
},
self.get_add_object("cloth_square",
library="models_special.json",
position={
"x": 0,
"y": 2,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
object_id=o_id), {
"$type": "set_flex_cloth_actor",
"id": o_id,
"mesh_tesselation": 1,
"stretch_stiffness": 0.5620341548096974,
"bend_stiffness": 0.6528988964052056,
"tether_stiffness": 0.7984931184979334,
"tether_give": 0,
"pressure": 0,
"mass_scale": 1
}, {
"$type": "assign_flex_container",
"container_id": 0,
"id": o_id
}, {
"$type": "send_flex_particles",
"frequency": "once"
}
]
# Add to the list of commands: create the avatar, teleport it to a position, look at a position.
commands.extend(
TDWUtils.create_avatar(position={
"x": -2.75,
"y": 2.3,
"z": -2
},
look_at={
"x": 0,
"y": 0.25,
"z": 0
}))
# Send the commands.
resp = self.communicate(commands)
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# This is FlexParticles data.
if r_id == "flex":
fp = FlexParticles(resp[0])
# The maximum distance between a particle and the center of the cloth.
max_d = 0
# The ID of the particle at the maximum distance.
max_id = 0
p_id = 0
for p in fp.get_particles(0):
# Get the particle that is furthest from the center.
d = np.linalg.norm(p[:-1] - 1)
if d > max_d:
max_d = d
max_id = p_id
p_id += 1
# Set the particle as "fixed".
self.communicate({
"$type": "set_flex_particle_fixed",
"is_fixed": True,
"particle_id": max_id,
"id": o_id
})
for i in range(800):
self.communicate([])
def run(self):
self.start()
self.communicate(TDWUtils.create_empty_room(12, 12))
# Create the objects.
statue_id = self.add_object("satiro_sculpture",
position={
"x": 2,
"y": 0,
"z": 1
})
# Request segmentation colors.
resp = self.communicate({"$type": "send_segmentation_colors"})
segmentation_colors = SegmentationColors(resp[0])
# Get the segmentation color of the sculpture.
statue_color = None
for i in range(segmentation_colors.get_num()):
if segmentation_colors.get_object_id(i) == statue_id:
statue_color = segmentation_colors.get_object_color(i)
break
# Create the VR Rig.
self.communicate({"$type": "create_vr_rig"})
# Attach an avatar to the VR rig.
self.communicate({"$type": "attach_avatar_to_vr_rig", "id": "a"})
# Request the colors of objects currently observed by the avatar per frame.
# Request VR rig data per frame.
# Reduce render quality in order to improve framerate.
self.communicate([{
"$type": "send_id_pass_segmentation_colors",
"frequency": "always"
}, {
"$type": "send_vr_rig",
"frequency": "always"
}, {
"$type": "set_post_process",
"value": False
}, {
"$type": "set_render_quality",
"render_quality": 0
}])
while True:
resp = self.communicate({"$type": "do_nothing"})
head_rotation = None
can_see_statue = False
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Get the head rotation.
if r_id == "vrri":
head_rotation = VRRig(r).get_head_rotation()
# Check if we can see the statue.
elif r_id == "ipsc":
observed_objects = IdPassSegmentationColors(r)
for i in range(
observed_objects.get_num_segmentation_colors()):
if observed_objects.get_segmentation_color(
i) == statue_color:
can_see_statue = True
if can_see_statue:
print("You can see the object!\nHead rotation: " +
str(head_rotation) + "\n")
def run(self):
self.start()
init_setup_commands = [{
"$type": "set_screen_size",
"width": 600,
"height": 480
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
self.communicate(init_setup_commands)
self.communicate(self.get_add_hdri_skybox("industrial_sunset_4k"))
# Create an empty room.
self.communicate(TDWUtils.create_empty_room(8, 8))
# Disable physics.
self.communicate({"$type": "set_gravity", "value": False})
# Add the avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": 1,
"y": 1.5,
"z": 0.3
},
look_at=TDWUtils.array_to_vector3(
[2.5, 0.5, 0.5]),
avatar_id="avatar"))
lib = MaterialLibrarian(library="materials_med.json")
record = lib.get_record("concrete_049")
self.communicate({
"$type": "add_material",
"name": "concrete_049",
"url": record.get_url()
})
self.communicate({
"$type": "set_proc_gen_floor_material",
"name": "concrete_049"
})
# self.communicate({"$type": "set_proc_gen_walls_material", "name": "concrete"})
# self.communicate({"$type": "set_field_of_view",
# "field_of_view": 68.0,
# "avatar_id": "avatar"})
bench = self.add_object(model_name="cgaxis_models_51_19_01",
position={
"x": 2.5,
"y": 0,
"z": 0.5
},
rotation={
"x": 0,
"y": 90,
"z": 0
},
library="models_full.json")
bench_bounds = self.get_bounds_data(bench)
top = bench_bounds.get_top(0)
self.add_object(model_name="b05_cat_model_3dmax2012",
position={
"x": 2.5,
"y": top[1],
"z": 0.2
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library="models_full.json")
self.add_object(model_name="azor",
position={
"x": 3,
"y": 0,
"z": 0.5
},
rotation={
"x": 0,
"y": 90,
"z": 0
},
library="models_full.json")
""" Post-scene construction """
# Enable image capture
self.communicate({
"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_id"]
})
scene_data = self.communicate({
"$type":
"look_at_position",
"avatar_id":
"avatar",
"position":
TDWUtils.array_to_vector3([3, 0.5, 0.5])
})
# images = Images(scene_data[0])
im = cv2.imread(
"/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/"
"compare_COCO_TDW/replicated_images/exterior/id_bench_bike.png")
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
# WRITE THIS TO JSON!
segmentation_data = []
COCO_to_TDW = pd.read_csv(
'/Users/leonard/Desktop/coco/COCO_to_TDW.csv',
header=None,
index_col=0,
squeeze=True).to_dict()
segmentation_colors = self.communicate({
"$type": "send_segmentation_colors",
"frequency": "once"
})
for r in segmentation_colors[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "segm":
s = SegmentationColors(r)
for i in range(s.get_num()):
name = s.get_object_name(i)
category = get_COCO(name, COCO_to_TDW)
color = s.get_object_color(i)
indices = np.where(im != 0) # Segmentation pixels
print(color)
print(indices)
print(im[indices[0][0]]
) # some random RGB val -- MUSt match one of the 3
segmentation_data.append([indices, category])
# TDWUtils.save_images(images, "bench_bike",
# output_directory="/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/compare_COCO_TDW/replicated_images/exterior")
# Consider using zip() to clean up coordinates
return segmentation_data
def run(self) -> None:
self.start()
# Get IDs for the object, the table, and the avatars.
o_id = self.get_unique_id()
table_id = self.get_unique_id()
table_position = {"x": 0, "y": 0, "z": 6.8}
a = "a"
# 1. Create the room.
commands = [TDWUtils.create_empty_room(20, 20)]
# 2. Add the avatar.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_StickyMitten_Adult",
avatar_id="a"))
# 3. Add the objects. Scale the objects. Set a high mass for the table.
# 4. Disable the avatar's cameras.
# 5. Set the stickiness of the avatar's left mitten. Set a high drag. Rotate the head.
# 6. Request AvatarChildrenNames data and Bounds data (this frame only).
# 7. Strengthen the left shoulder a bit.
commands.extend([
self.get_add_object("jug05",
position={
"x": -0.417,
"y": 0.197,
"z": 0.139
},
rotation={
"x": 90,
"y": 0,
"z": 0
},
object_id=o_id), {
"$type": "scale_object",
"scale_factor": {
"x": 2,
"y": 2,
"z": 2
},
"id": o_id
},
self.get_add_object("small_table_green_marble",
position={
"x": 0,
"y": 0,
"z": 6.8
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
object_id=table_id), {
"$type": "scale_object",
"scale_factor": {
"x": 2,
"y": 0.5,
"z": 2
},
"id": table_id
},
{
"$type": "set_object_collision_detection_mode",
"id": table_id,
"mode": "continuous_dynamic"
}, {
"$type": "set_object_collision_detection_mode",
"id": o_id,
"mode": "continuous_dynamic"
}, {
"$type": "set_mass",
"mass": 100,
"id": table_id
}, {
"$type": "toggle_image_sensor",
"sensor_name": "SensorContainer",
"avatar_id": a
}, {
"$type": "toggle_image_sensor",
"sensor_name": "FollowCamera",
"avatar_id": a
}, {
"$type": "set_stickiness",
"sub_mitten": "back",
"sticky": True,
"is_left": True,
"avatar_id": a
}, {
"$type": "set_avatar_drag",
"drag": 1000,
"angular_drag": 1000,
"avatar_id": a
}, {
"$type": "rotate_head_by",
"axis": "pitch",
"angle": 5
}, {
"$type": "send_avatar_children_names",
"ids": [a],
"frequency": "once"
}, {
"$type": "send_bounds",
"ids": [table_id],
"frequency": "once"
}, {
"$type": "adjust_joint_force_by",
"delta": 2,
"joint": "shoulder_left",
"axis": "pitch"
}
])
# 7. Add a 3rd-person camera.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
avatar_id="c",
position={
"x": -3.9,
"y": 2.3,
"z": 4.3
}))
# 8. Request StickyMittenAvatar and Images data per-frame.
commands.extend([{
"$type": "set_pass_masks",
"pass_masks": ["_img"],
"avatar_id": self.cam_id
}, {
"$type": "send_avatars",
"ids": [a],
"frequency": "always"
}, {
"$type": "send_images",
"ids": [self.cam_id],
"frequency": "always"
}])
resp = self.communicate(commands)
# Get the object ID of the left mitten, the size of the table, and the avatar's position.
mitten_left_id = None
table_size: Tuple[float, float, float] = (0, 0, 0)
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Get the mitten ID.
if r_id == "avcn":
avcn = AvatarChildrenNames(r)
for i in range(avcn.get_num_children()):
if avcn.get_child_name(i) == "mitten_left":
mitten_left_id = avcn.get_child_id(i)
# Get the table bounds.
elif r_id == "boun":
boun = Bounds(r)
for i in range(boun.get_num()):
if boun.get_id(i) == table_id:
table_size = (boun.get_right(i)[0] -
boun.get_left(i)[0], boun.get_top(i)[1] -
boun.get_bottom(i)[1],
boun.get_front(i)[2] -
boun.get_back(i)[2])
# Get the avatar's position.
elif r_id == "avsm":
self.avatar_position = AvatarStickyMitten(r).get_position()
# Pick up the object.
# Lift up the object.
self._bend_arm_joints([{
"$type": "pick_up_proximity",
"distance": 20,
"grip": 10000,
"is_left": True,
"avatar_id": a
}, {
"$type": "rotate_head_by",
"axis": "pitch",
"angle": 20,
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 25,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": -25,
"joint": "shoulder_left",
"axis": "yaw",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 60,
"joint": "shoulder_left",
"axis": "roll",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 100,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Allow the avatar to move again.
self._do_frame({
"$type": "set_avatar_drag",
"drag": 0.125,
"angular_drag": 1000,
"avatar_id": a
})
# Move to the table.
move_to_table = True
# The position of the side of the table the avatar is aiming for.
table_side_position = {
"x": table_position["x"],
"y": 0,
"z": table_position["z"] - table_size[2]
}
while move_to_table:
# Stop moving if we are close enough.
if TDWUtils.get_distance(
table_side_position,
TDWUtils.array_to_vector3(self.avatar_position)) < 0.7:
move_to_table = False
# Keep moving forward.
else:
self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": 20
})
# Stop.
# Allow the avatar to move again.
self._do_frame({
"$type": "set_avatar_drag",
"drag": 1000,
"angular_drag": 1000,
"avatar_id": a
})
mitten_over_table = False
target_z = table_side_position["z"] + 0.17
# Keep lifting until the mitten is over the table.
while not mitten_over_table:
avsm = self._bend_arm_joints([{
"$type": "bend_arm_joint_by",
"angle": 15,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": -10,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten is over the table.
mitten_over_table = avsm.get_body_part_position(
i)[2] >= target_z
# Drop the arm gently.
mitten_near_table = False
while not mitten_near_table:
avsm = self._bend_arm_joints([{
"$type": "bend_arm_joint_by",
"angle": -25,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}])
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten very close to the table surface.
mitten_near_table = avsm.get_body_part_position(
i)[1] - table_size[1] <= 0.1
# Drop the object.
# Allow the avatar to move again.
self._do_frame([{
"$type": "put_down",
"is_left": True,
"avatar_id": a
}, {
"$type": "set_avatar_drag",
"drag": 0.125,
"angular_drag": 1000,
"avatar_id": a
}])
# Back away from the table.
mitten_away_from_table = False
while not mitten_away_from_table:
avsm = self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": -50
})
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten has cleared the table.
mitten_away_from_table = avsm.get_body_part_position(
i)[2] < target_z
# Drop the arm.
self._do_frame([{
"$type": "rotate_head_by",
"axis": "pitch",
"angle": -25
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "yaw",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "roll",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Back away from the table.
away_from_table = False
while not away_from_table:
self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": -20
})
away_from_table = TDWUtils.get_distance(
TDWUtils.array_to_vector3(self.avatar_position),
table_position) > 1.8
# Let the joints drop.
self._bend_arm_joints([])
self.communicate({"$type": "terminate"})
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))
collision_states = np.empty(dtype=str, shape=(0, 1))
# 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
# rtypes = [OutputData.get_data_type_id(r) for r in resp[:-1]]
# print(frame_num, "COLLISION" if 'coll' in rtypes else "")
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)):
try:
velocities[i] = ri_dict[o_id]["vel"]
angular_velocities[i] = ri_dict[o_id]["ang"]
except KeyError:
print("Couldn't store velocity data for object %d" %
o_id)
print("frame num", frame_num)
print("ri_dict", ri_dict)
print([OutputData.get_data_type_id(r) for r in resp])
elif r_id == "coll":
co = Collision(r)
collision_states = np.append(collision_states, co.get_state())
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")
collisions.create_dataset("states",
data=collision_states.astype('S'),
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 _write_frame(self, frames_grp: h5py.Group, resp: List[bytes], frame_num: int) -> \
Tuple[h5py.Group, h5py.Group, dict, bool]:
num_objects = len(self.object_ids)
# Create a group for this frame.
frame = frames_grp.create_group(TDWUtils.zero_padding(frame_num, 4))
# Create a group for images.
images = frame.create_group("images")
# Transforms data.
positions = np.empty(dtype=np.float32, shape=(num_objects, 3))
forwards = np.empty(dtype=np.float32, shape=(num_objects, 3))
rotations = np.empty(dtype=np.float32, shape=(num_objects, 4))
camera_matrices = frame.create_group("camera_matrices")
# Parse the data in an ordered manner so that it can be mapped back to the object IDs.
tr_dict = dict()
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "tran":
tr = Transforms(r)
for i in range(tr.get_num()):
pos = tr.get_position(i)
tr_dict.update({
tr.get_id(i): {
"pos": pos,
"for": tr.get_forward(i),
"rot": tr.get_rotation(i)
}
})
# Add the Transforms data.
for o_id, i in zip(self.object_ids, range(num_objects)):
if o_id not in tr_dict:
continue
positions[i] = tr_dict[o_id]["pos"]
forwards[i] = tr_dict[o_id]["for"]
rotations[i] = tr_dict[o_id]["rot"]
elif r_id == "imag":
im = Images(r)
# Add each image.
for i in range(im.get_num_passes()):
pass_mask = im.get_pass_mask(i)
# Reshape the depth pass array.
if pass_mask == "_depth":
image_data = TDWUtils.get_shaped_depth_pass(images=im,
index=i)
else:
image_data = im.get_image(i)
images.create_dataset(pass_mask,
data=image_data,
compression="gzip")
# Save PNGs
if pass_mask in self.save_passes:
filename = pass_mask[1:] + "_" + TDWUtils.zero_padding(
frame_num, 4) + "." + im.get_extension(i)
path = self.png_dir.joinpath(filename)
if pass_mask in ["_depth", "_depth_simple"]:
Image.fromarray(
TDWUtils.get_shaped_depth_pass(
images=im, index=i)).save(path)
else:
with open(path, "wb") as f:
f.write(im.get_image(i))
# Add the camera matrices.
elif OutputData.get_data_type_id(r) == "cama":
matrices = CameraMatrices(r)
camera_matrices.create_dataset(
"projection_matrix", data=matrices.get_projection_matrix())
camera_matrices.create_dataset(
"camera_matrix", data=matrices.get_camera_matrix())
objs = frame.create_group("objects")
objs.create_dataset("positions",
data=positions.reshape(num_objects, 3),
compression="gzip")
objs.create_dataset("forwards",
data=forwards.reshape(num_objects, 3),
compression="gzip")
objs.create_dataset("rotations",
data=rotations.reshape(num_objects, 4),
compression="gzip")
return frame, objs, tr_dict, False
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_real_grayscale(
self, o_id: int, a_id: str,
e: Environment) -> (float, float, dict, dict, dict, dict):
"""
Get the "real" grayscale value of an image we hope to capture.
:param o_id: The ID of the object.
:param a_id: The ID of the avatar.
:param e: The environment.
:return: (grayscale, distance, avatar_position, object_position, object_rotation, avatar_rotation)
"""
# Get a random position for the avatar.
a_p = np.array([
RNG.uniform(-e.w / 2, e.w / 2) + e.x,
RNG.uniform(0.4, e.h * self.max_height),
RNG.uniform(-e.l / 2, e.l / 2) + e.z
])
# Get a random distance from the avatar.
d = RNG.uniform(0.8, 3)
# Get a random position for the object constrained to the environment bounds.
o_p = SingleObject.sample_spherical() * d
# Clamp the y value to positive.
o_p[1] = abs(o_p[1])
o_p = a_p + o_p
# Clamp the y value of the object.
if o_p[1] > e.h * self.max_height:
o_p[1] = e.h * self.max_height
# Convert the avatar's position to a Vector3.
a_p = TDWUtils.array_to_vector3(a_p)
# Set random camera rotations.
yaw = RNG.uniform(-15, 15)
pitch = RNG.uniform(-15, 15)
# Convert the object position to a Vector3.
o_p = TDWUtils.array_to_vector3(o_p)
# Add rotation commands.
# If we're clamping the rotation, rotate the object within +/- 30 degrees on each axis.
if self.clamp_rotation:
o_rot = None
commands = [{
"$type": "rotate_object_to",
"id": o_id,
"rotation": {
"x": 0,
"y": 0,
"z": 0,
"w": 0
}
}, {
"$type": "rotate_object_by",
"id": o_id,
"angle": RNG.uniform(-30, 30),
"axis": "pitch"
}, {
"$type": "rotate_object_by",
"id": o_id,
"angle": RNG.uniform(-30, 30),
"axis": "yaw"
}, {
"$type": "rotate_object_by",
"id": o_id,
"angle": RNG.uniform(-30, 30),
"axis": "roll"
}]
# Set a totally random rotation.
else:
o_rot = {
"x": RNG.uniform(-360, 360),
"y": RNG.uniform(-360, 360),
"z": RNG.uniform(-360, 360),
"w": RNG.uniform(-360, 360),
}
commands = [{
"$type": "rotate_object_to",
"id": o_id,
"rotation": o_rot
}]
# After rotating the object:
# 1. Teleport the object.
# 2. Teleport the avatar.
# 3. Look at the object.
# 4. Perturb the camera slightly.
# 5. Send grayscale data and image sensor data.
commands.extend([{
"$type": "teleport_object",
"id": o_id,
"position": o_p
}, {
"$type": "teleport_avatar_to",
"avatar_id": a_id,
"position": a_p
}, {
"$type": "look_at",
"avatar_id": a_id,
"object_id": o_id,
"use_centroid": True
}, {
"$type": "rotate_sensor_container_by",
"angle": pitch,
"avatar_id": "a",
"axis": "pitch"
}, {
"$type": "rotate_sensor_container_by",
"angle": yaw,
"avatar_id": "a",
"axis": "yaw"
}, {
"$type": "send_id_pass_grayscale",
"frequency": "once"
}, {
"$type": "send_image_sensors",
"frequency": "once"
}])
resp = self.communicate(commands)
# Parse the output data:
# 1. The grayscale value of the image.
# 2. The camera rotation.
grayscale = 0
cam_rot = None
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
if r_id == "idgs":
grayscale = IdPassGrayscale(r).get_grayscale()
elif r_id == "imse":
cam_rot = ImageSensors(r).get_sensor_rotation(0)
cam_rot = {
"x": cam_rot[0],
"y": cam_rot[1],
"z": cam_rot[2],
"w": cam_rot[3]
}
# If we clamped the rotation of the object, we need to know its quaternion.
if self.clamp_rotation:
resp = self.communicate({
"$type": "send_transforms",
"frequency": "once",
"ids": [o_id]
})
t = Transforms(resp[0])
o_rot = t.get_rotation(0)
o_rot = {
"x": o_rot[0],
"y": o_rot[1],
"z": o_rot[2],
"w": o_rot[3],
}
return grayscale, d, a_p, o_p, o_rot, cam_rot
def run(self) -> None:
robot_id = 0
self.start()
# Create the scene. Add a robot.
# Request static robot data for this frame only.
# Request dynamic robot data per frame.
commands = [
TDWUtils.create_empty_room(12, 12),
self.get_add_robot(name="ur3", robot_id=robot_id), {
"$type": "send_static_robots",
"frequency": "once"
}, {
"$type": "send_robots",
"frequency": "always"
}
]
# Add an avatar to render the scene (just for demo purposes).
commands.extend(
TDWUtils.create_avatar(look_at=TDWUtils.VECTOR3_ZERO,
position={
"x": -0.881,
"y": 0.836,
"z": -1.396
}))
resp = self.communicate(commands)
# Parse the output data for static robot data.
static_robot: Optional[StaticRobot] = None
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "srob":
r = StaticRobot(resp[i])
if r.get_id() == robot_id:
static_robot = r
break
assert static_robot is not None, f"No static robot data: {resp}"
# Get the IDs of the shoulder and the elbow.
body_part_ids: Dict[str, int] = dict()
for i in range(static_robot.get_num_joints()):
b_id = static_robot.get_joint_id(i)
b_name = static_robot.get_joint_name(i)
body_part_ids[b_name] = b_id
assert "Shoulder" in body_part_ids
assert "Elbow" in body_part_ids
# Rotate the shoulder and the elbow for two motions.
# The values in this array are for the angle that the [shoulder, elbow] should rotate to per action.
# For more complex actions, you will probably want to organize your commands differently.
for angles in [[70, 90], [-30, -25]]:
resp = self.communicate([{
"$type": "set_revolute_target",
"id": robot_id,
"joint_id": body_part_ids["Shoulder"],
"target": angles[0]
}, {
"$type": "set_revolute_target",
"id": robot_id,
"joint_id": body_part_ids["Elbow"],
"target": angles[1]
}])
# Get the robot output data.
robot = self.get_robot(resp=resp)
angles_0 = self.get_joint_angles(robot=robot)
# Wait for the joints to stop moving.
moving = True
while moving:
robot = self.get_robot(resp=self.communicate([]))
angles_1 = self.get_joint_angles(robot=robot)
# Compare the current joint angles to the angles of the previous frame.
moving = False
for angle_0, angle_1 in zip(angles_0, angles_1):
if np.abs(angle_0 - angle_1) > 0.001:
moving = True
break
angles_0 = angles_1
def run(self):
self.start()
positions_list = [] # Stores current model locations and radii
scene_dimensions = [] # Store scene/environment dimensions
init_setup_commands = [{
"$type": "set_screen_size",
"width": 1280,
"height": 962
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
self.communicate(init_setup_commands)
scene_lib = SceneLibrarian()
# Disable physics when adding in new objects (objects float)
self.communicate({"$type": "simulate_physics", "value": False})
for scene in scenes[1:]:
# Load in scene
print("Scene", scene[0])
if scene[3] == "interior" and scene[0] == "box_room_2018":
self.start()
scene_record = scene_lib.get_record(scene[0])
self.communicate({
"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()
})
# Gets dimensions of environments (e.g. inside, outside) in the scene
# This command returns environment data in the form of a list of serialized byte arrays
scene_bytes = self.communicate({
"$type": "send_environments",
"frequency": "once"
})
# Iterating through data and parsing byte array
# Ignoring the last element (the frame count)
for b in scene_bytes[:-1]:
e = Environments(b)
for i in range(e.get_num()):
center = e.get_center(i)
bounds = e.get_bounds(i)
env_id = e.get_id(i)
scene_dimensions = [center, bounds,
env_id] # Center, bounds are tuples
# Must come before set_pass_masks
avatar_position = TDWUtils.array_to_vector3([
0.9 * scene_dimensions[1][0] / 2,
scene_dimensions[1][1] / 2, 0
])
print("Avatar Position:", avatar_position)
self.communicate(
TDWUtils.create_avatar(avatar_id="avatar",
position=avatar_position,
look_at={
"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0
}))
# Set collision mode
self.communicate({
"$type": "set_avatar_collision_detection_mode",
"mode": "continuous_speculative",
"avatar_id": "avatar"
})
# Alter FOV
self.communicate({
"$type": "set_field_of_view",
"field_of_view": 80,
"avatar_id": "avatar"
})
# Gets rid of header (Model: Category)
objects = models[1:]
random.shuffle(objects)
obj_count = 0
obj_overlaps = 0 # Number of failed attempts to place object due to over-dense objects area
while obj_count < 30 and obj_overlaps < 5:
# Need to have random position for Bounds Data to return meaningful info
valid_obj_pos = {
"x":
random.uniform(-1 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][0] / 2),
"y":
scene_dimensions[1][1] / 4,
"z":
random.uniform(-0.9 * scene_dimensions[1][2] / 2,
0.9 * scene_dimensions[1][2] / 2)
}
# Add in the object at random position
# Object will later be removed or updated accordingly after performing collision calculations
record = ModelLibrarian(
library="models_full.json").get_record(
objects[obj_count][0])
self.communicate({
"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": {
"x": 0,
"y": 0,
"z": 0
},
"category": record.wcategory,
"id": obj_count
})
# Returns bound data for added object
bounds_data = self.communicate({
"$type": "send_bounds",
"frequency": "once"
})
# Appends object, with information on position and obj_radius, to positions_list
# Length of buffer array should be 1
print("Bounds Data:", bounds_data)
for b in bounds_data[:-1]:
print("Buffer Loop:", b)
b_id = OutputData.get_data_type_id(b)
if b_id == "boun":
print("BOUNDS")
o = Bounds(b)
print("# of Objects:", o.get_num())
print("# of Failed Attempts:", obj_overlaps)
print("Buffer Array:", b)
print("Bounds Object:", o)
for i in range(o.get_num()):
print("Object ID:", o.get_id(i))
print("obj_count:", obj_count)
print("Object:", objects[obj_count][0],
"Category:", objects[obj_count][1])
print("Object Center:", o.get_center(i))
# Only want to compute valid_position for object we are about to add
# Skip any computation if this is not the case
if o.get_id(i) != obj_count:
continue
# Useful for detecting if object fits in environment
print(
"Calculating if object fits in environment"
)
width = distance.euclidean(
o.get_left(i), o.get_right(i))
depth = distance.euclidean(
o.get_front(i), o.get_back(i))
height = distance.euclidean(
o.get_top(i), o.get_bottom(i))
print("Width:", width)
print("Depth:", depth)
print("Height:", height)
# Useful for avoiding object overlap
print("Calculating Object Bounds")
center_to_top = distance.euclidean(
o.get_center(i), o.get_top(i))
center_to_bottom = distance.euclidean(
o.get_center(i), o.get_bottom(i))
center_to_left = distance.euclidean(
o.get_center(i), o.get_left(i))
center_to_right = distance.euclidean(
o.get_center(i), o.get_right(i))
center_to_front = distance.euclidean(
o.get_center(i), o.get_front(i))
center_to_back = distance.euclidean(
o.get_center(i), o.get_back(i))
# Max object radius (center to diagonal of bounding box)
obj_radius = \
max(math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_back ** 2))
print("Obj_Radius:", obj_radius)
# Set sweeping radius, based on scene plane dimensions
l_radius = random.uniform(
0,
min(0.9 * scene_dimensions[1][0] / 2,
0.9 * scene_dimensions[1][2] / 2))
# Checking that object fits in scene viewing
if (width > min(0.7 * scene_dimensions[1][0],
0.7 * scene_dimensions[1][2])
or depth > min(
0.7 * scene_dimensions[1][0],
0.7 * scene_dimensions[1][2]) or
height > 0.7 * scene_dimensions[1][1]):
print("Object does not fit in scene")
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Not possible to find valid object position -- too many overlapping objects
elif (not self._get_object_position(
scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[0]):
print(
"Could not calculate valid object location"
)
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
obj_overlaps += 1
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Find appropriate, non-overlapping object position
# Reset object position to the valid position
else:
print("Object fits in scene")
valid_obj_pos = self._get_object_position(
scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[1]
print("Position calculated")
positions_list.append(
ObjectPosition(valid_obj_pos,
obj_radius))
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
print("Object ready to reset")
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type":
"add_object",
"name":
objects[obj_count][0],
"url":
record.get_url(),
"scale_factor":
record.scale_factor,
"position":
valid_obj_pos,
"rotation": {
"x": 0,
"y": 0,
"z": 0
},
"category":
record.wcategory,
"id":
obj_count
}])
print("Object reset")
# Rotate the object randomly
print("Rotating object")
self.communicate({
"$type":
"rotate_object_by",
"angle":
random.uniform(-45, 45),
"axis":
"yaw",
"id":
obj_count,
"is_world":
True
})
# Don't rotate the object if doing so will result in overlap into scene
if not (o.get_bottom(i)[1] < 0
or o.get_top(i)[1] >
0.9 * scene_dimensions[1][1]):
pitch_angle = random.uniform(-45, 45)
self.communicate({
"$type": "rotate_object_by",
"angle": pitch_angle,
"axis": "pitch",
"id": obj_count,
"is_world": True
})
roll_angle = random.uniform(-45, 45)
self.communicate({
"$type": "rotate_object_by",
"angle": roll_angle,
"axis": "roll",
"id": obj_count,
"is_world": True
})
# Setting random materials/textures
# Looping through sub-objects and materials
sub_count = 0
for sub_object in record.substructure:
# Loop through materials in sub-objects
for j in range(len(sub_object)):
# Get random material and load in
material = random.choice(
materials[1:])
self.load_material(material)
print("Material loaded")
# Set random material on material of sub-object
self.communicate({
"$type":
"set_visual_material",
"material_index":
j,
"material_name":
material[0],
"object_name":
sub_object['name'],
"id":
obj_count
})
print("Material set")
sub_count += 1
if sub_count > 10:
break
break
print("Updating count")
obj_count += 1
print("Breaking out of object_id loop")
break
# Break out of buffer loop
print("Breaking out of buffer loop")
break
# Move onto next iteration of while loop (next object to load in)
print("Object added - next while loop iteration")
continue
# Enable image capture
self.communicate({
"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]
})
self.communicate({
"$type": "send_images",
"frequency": "always"
})
# Capture scene
scene_data = self.communicate({
"$type": "look_at_position",
"avatar_id": "avatar",
"position": {
"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0
}
})
images = Images(scene_data[0])
TDWUtils.save_images(images,
TDWUtils.zero_padding(i),
output_directory=path)
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
}), {
"$type": "send_bounds",
"ids": [table_id],
"frequency": "once"
}])
# Get the top of the table.
top_y = 0
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Find the bounds data.
if r_id == "boun":
b = Bounds(r)
# We only requested the table, so it is object 0:
_, top_y, _ = b.get_top(0)
box_record = lib.get_record("puzzle_box_composite")
box_id = 1
c.communicate(
c.get_add_object(model_name=box_record.name,
object_id=box_id,
position={
"x": 0,
"y": top_y,
"z": 0
def run(self,
num: int,
output_dir: str,
temp_path: str,
width: int,
height: int,
write_passes: List[str] = PASSES,
save_passes: List[str] = [],
save_movies: bool = False,
save_labels: bool = False,
args_dict: dict={}) -> None:
"""
Create the dataset.
:param num: The number of trials in the dataset.
:param output_dir: The root output directory.
:param temp_path: Temporary path to a file being written.
:param width: Screen width in pixels.
:param height: Screen height in pixels.
:param save_passes: a list of which passes to save out as PNGs (or convert to MP4)
:param save_movies: whether to save out a movie of each trial
:param save_labels: whether to save out JSON labels for the full trial set.
"""
self._height, self._width = height, width
# which passes to write to the HDF5
self.write_passes = write_passes
if isinstance(self.write_passes, str):
self.write_passes = self.write_passes.split(',')
self.write_passes = [p for p in self.write_passes if (p in PASSES)]
# which passes to save as an MP4
self.save_passes = save_passes
if isinstance(self.save_passes, str):
self.save_passes = self.save_passes.split(',')
self.save_passes = [p for p in self.save_passes if (p in self.write_passes)]
self.save_movies = save_movies
print("write passes", self.write_passes)
print("save passes", self.save_passes)
print("save movies", self.save_movies)
if self.save_movies:
assert len(self.save_passes),\
"You need to pass \'--save_passes [PASSES]\' to save out movies, where [PASSES] is a comma-separated list of items from %s" % PASSES
# whether to save a JSON of trial-level labels
self.save_labels = save_labels
if self.save_labels:
self.meta_file = Path(output_dir).joinpath('metadata.json')
if self.meta_file.exists():
self.trial_metadata = json.loads(self.meta_file.read_text())
else:
self.trial_metadata = []
initialization_commands = self.get_initialization_commands(width=width, height=height)
# Initialize the scene.
self.communicate(initialization_commands)
# Run trials
self.trial_loop(num, output_dir, temp_path)
# Terminate TDW
end = self.communicate({"$type": "terminate"})
end_types = [OutputData.get_data_type_id(r) for r in end]
end_count = 0
print("TOLD BUILD TO TERMINATE, TRY %d" % end_count)
print(end_types)
while ('tre\x04' in end_types) and end_count <= 10:
end_count += 1
print([OutputData.get_data_type_id(r) for r in end])
print("TOLD BUILD TO TERMINATE, TRY %d" % end_count)
end = self.communicate({"$type": "terminate"})
# Save the command line args
if self.save_args:
self.args_dict = copy.deepcopy(args_dict)
self.save_command_line_args(output_dir)
# Save the across-trial stats
if self.save_labels:
hdf5_paths = glob.glob(str(output_dir) + '/*.hdf5')
stats = get_across_trial_stats_from(
hdf5_paths, funcs=self.get_controller_label_funcs(classname=type(self).__name__))
stats["num_trials"] = int(len(hdf5_paths))
stats_str = json.dumps(stats, indent=4)
stats_file = Path(output_dir).joinpath('trial_stats.json')
stats_file.write_text(stats_str, encoding='utf-8')
print("ACROSS TRIAL STATS")
print(stats_str)
def trial(self,
filepath: Path,
temp_path: Path,
trial_num: int) -> None:
"""
Run a trial. Write static and per-frame data to disk until the trial is done.
:param filepath: The path to this trial's hdf5 file.
:param temp_path: The path to the temporary file.
:param trial_num: The number of the current trial.
"""
# Clear the object IDs and other static data
self.clear_static_data()
self._trial_num = trial_num
# Create the .hdf5 file.
f = h5py.File(str(temp_path.resolve()), "a")
commands = []
# Remove asset bundles (to prevent a memory leak).
if trial_num % 100 == 0:
commands.append({"$type": "unload_asset_bundles"})
# Add commands to start the trial.
commands.extend(self.get_trial_initialization_commands())
# Add commands to request output data.
commands.extend(self._get_send_data_commands())
# Send the commands and start the trial.
r_types = ['']
count = 0
while 'imag' not in r_types:
resp = self.communicate(commands)
r_types = [OutputData.get_data_type_id(r) for r in resp]
count += 1
print("INITIAL RESPONSE FROM BUILD ON TRY %d" % count)
print(r_types)
self._set_segmentation_colors(resp)
frame = 0
# Write static data to disk.
static_group = f.create_group("static")
self._write_static_data(static_group)
# Add the first frame.
done = False
frames_grp = f.create_group("frames")
frame_grp, _, _, _ = self._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame)
self._write_frame_labels(frame_grp, resp, -1, False)
# Continue the trial. Send commands, and parse output data.
while not done:
frame += 1
# print('frame %d' % frame)
resp = self.communicate(self.get_per_frame_commands(resp, frame))
r_ids = [OutputData.get_data_type_id(r) for r in resp[:-1]]
# Sometimes the build freezes and has to reopen the socket.
# This prevents such errors from throwing off the frame numbering
if ('imag' not in r_ids) or ('tran' not in r_ids):
print("retrying frame %d, response only had %s" % (frame, r_ids))
frame -= 1
continue
frame_grp, objs_grp, tr_dict, done = self._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame)
# Write whether this frame completed the trial and any other trial-level data
labels_grp, _, _, done = self._write_frame_labels(frame_grp, resp, frame, done)
# Cleanup.
commands = []
for o_id in self.object_ids:
commands.append({"$type": self._get_destroy_object_command_name(o_id),
"id": int(o_id)})
self.communicate(commands)
# Compute the trial-level metadata. Save it per trial in case of failure mid-trial loop
if self.save_labels:
meta = OrderedDict()
meta = get_labels_from(f, label_funcs=self.get_controller_label_funcs(type(self).__name__), res=meta)
self.trial_metadata.append(meta)
# Save the trial-level metadata
json_str =json.dumps(self.trial_metadata, indent=4)
self.meta_file.write_text(json_str, encoding='utf-8')
print("TRIAL %d LABELS" % self._trial_num)
print(json.dumps(self.trial_metadata[-1], indent=4))
# Close the file.
f.close()
# Move the file.
try:
temp_path.replace(filepath)
except OSError:
shutil.move(temp_path, filepath)
def trial(self, d_f, s_f, b_f, d_c, s_c, b_c, collision_types):
"""
Collide a chair with a fridge.
:param d_f: The dynamic friction of the fridge.
:param s_f: The static friction of the fridge.
:param b_f: The bounciness of the fridge.
:param d_c: The dynamic friction of the chair.
:param s_c: The static friction of the chair.
:param b_c: The bounciness of the chair.
:param collision_types: The types of collisions to listen for.
"""
print("###\n\nNew Trial\n\n###\n")
fridge_id = 0
chair_id = 1
# Destroy all objects currently in the scene.
init_commands = [{"$type": "destroy_all_objects"}]
# Create the avatar.
init_commands.extend(
TDWUtils.create_avatar(position={
"x": 1,
"y": 2.5,
"z": 5
},
look_at=TDWUtils.VECTOR3_ZERO))
# Add the objects.
# Set the masses and physic materials.
# Apply a force to the chair.
# Receive collision data (note that by setting "stay" to True, you will receive a LOT of data;
# see "Performance Optimizations" documentation.)
init_commands.extend([
self.get_add_object("fridge_large", object_id=fridge_id),
self.get_add_object("chair_billiani_doll",
object_id=chair_id,
position={
"x": 4,
"y": 0,
"z": 0
}), {
"$type": "set_mass",
"id": fridge_id,
"mass": 40
}, {
"$type": "set_mass",
"id": chair_id,
"mass": 20
}, {
"$type": "set_physic_material",
"id": fridge_id,
"dynamic_friction": d_f,
"static_friction": s_f,
"bounciness": b_f
}, {
"$type": "set_physic_material",
"id": chair_id,
"dynamic_friction": d_c,
"static_friction": s_c,
"bounciness": b_c
}, {
"$type": "apply_force_to_object",
"force": {
"x": -200,
"y": 0,
"z": 0
},
"id": chair_id
}, {
"$type": "send_collisions",
"collision_types": collision_types,
"enter": True,
"exit": True,
"stay": True
}
])
self.communicate(init_commands)
# Iterate through 500 frames.
# Every frame, listen for collisions, and parse the output data.
for i in range(500):
resp = self.communicate([])
if len(resp) > 1:
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# There was a collision between two objects.
if r_id == "coll":
collision = Collision(r)
print("Collision between two objects:")
print("\tEvent: " + collision.get_state())
print("\tCollider: " +
str(collision.get_collider_id()))
print("\tCollidee: " +
str(collision.get_collidee_id()))
print("\tRelative velocity: " +
str(collision.get_relative_velocity()))
print("\tContacts:")
for j in range(collision.get_num_contacts()):
print(
str(collision.get_contact_normal(j)) + "\t" +
str(collision.get_contact_point(j)))
# There was a collision between an object and the environment.
elif r_id == "enco":
collision = EnvironmentCollision(r)
print(
"Collision between an object and the environment:")
print("\tEvent: " + collision.get_state())
print("\tCollider: " + str(collision.get_object_id()))
print("\tContacts:")
for j in range(collision.get_num_contacts()):
print(
str(collision.get_contact_normal(j)) + "\t" +
str(collision.get_contact_point(j)))
else:
raise Exception(r_id)
def __init__(self,
resp: List[bytes],
avatar_id: str = "a",
debug: bool = False):
"""
:param resp: The response from the build after creating the avatar.
:param avatar_id: The ID of the avatar.
:param debug: If True, print debug statements.
"""
self.id = avatar_id
self._debug = debug
# Set the arm chains.
self._arms: Dict[Arm, Chain] = {
Arm.left: self._get_left_arm(),
Arm.right: self._get_right_arm()
}
self._initial_mitten_positions = self._get_initial_mitten_positions()
# Any current IK goals.
self._ik_goals: Dict[Arm, Optional[_IKGoal]] = {
Arm.left: None,
Arm.right: None
}
smsc: Optional[AvatarStickyMittenSegmentationColors] = None
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
if r_id == "smsc":
q = AvatarStickyMittenSegmentationColors(resp[i])
if q.get_id() == avatar_id:
smsc = q
break
assert smsc is not None, f"No avatar segmentation colors found for {avatar_id}"
# Get data for the current frame.
self.frame = self._get_frame(resp)
# Get the masses of each body part.
body_part_masses: Dict[int, float] = dict()
for i in range(self.frame.get_num_rigidbody_parts()):
body_part_masses[self.frame.get_rigidbody_part_id(
i)] = self.frame.get_rigidbody_part_mass(i)
# Cache static data of body parts.
self.body_parts_static: Dict[int, BodyPartStatic] = dict()
self.base_id = 0
self.mitten_ids: Dict[Arm, int] = dict()
for i in range(smsc.get_num_body_parts()):
body_part_id = smsc.get_body_part_id(i)
if body_part_id in body_part_masses:
mass = body_part_masses[body_part_id]
else:
mass = 0.1
name = smsc.get_body_part_name(i)
# Cache the base ID and the mitten IDs.
if name.startswith("A_StickyMitten"):
self.base_id = body_part_id
elif name == "mitten_left":
self.mitten_ids[Arm.left] = body_part_id
elif name == "mitten_right":
self.mitten_ids[Arm.right] = body_part_id
bps = BodyPartStatic(
object_id=body_part_id,
segmentation_color=smsc.get_body_part_segmentation_color(i),
name=name,
mass=mass)
self.body_parts_static[body_part_id] = bps
# Start dynamic data.
self.collisions: Dict[int, List[int]] = dict()
self.env_collisions: List[int] = list()
self.status = TaskStatus.idle
self._mass = self._get_mass()
import json
from pathlib import Path
from tdw.controller import Controller
from tdw.output_data import OutputData
c = Controller(launch_build=True)
p = Path.home().joinpath("Downloads/tdw_commands (1).json")
txt = p.read_text()
for i, line in enumerate(txt.split("\n")):
try:
if not line.startswith("["):
continue
commands = json.loads(line.split(" trial ")[0])
if not isinstance(commands, list):
continue
resp = c.communicate(commands)
r_ids = [OutputData.get_data_type_id(r) for r in resp[:-1]]
assert "imag" in r_ids, (i, line, r_ids)
except json.JSONDecodeError:
continue
def run(self):
""" Generate room using COCO_TDW dataset """
objects_in_scene = 15
object_ids = []
# Get Category-Object mapping
TDW_COCO_models = TDW_relationships.get_COCO_TDW_mapping()
# print("TDWCOCO:", TDW_COCO_models)
# print("KEYS:", TDW_COCO_models.keys())
# Gets COCO categories co-occurring in a scene
# +5 is for dealing with failed object insertion attempts
COCO_configurations = msCOCO_matrix.get_max_co_occurrence(5, int(objects_in_scene + 5))
configuration_1 = COCO_configurations[0]
print("Config 1:", configuration_1)
# TDW models/objects
objects = []
for COCO_object in configuration_1:
print(COCO_object)
print(COCO_object.split())
if len(COCO_object.split()) > 1:
COCO_object = COCO_object.split()[-1]
print(COCO_object)
# Check if COCO category is a key in the COCO-TDW mapping
if COCO_object in TDW_COCO_models.keys():
# Gets random TDW model (from COCO-to-TDW map) based on COCO category key
print(TDW_COCO_models[COCO_object])
model = TDW_COCO_models[COCO_object][random.randint(0, len(TDW_COCO_models[COCO_object]) - 1)]
objects.append(model)
print("COCO to TDW objects:", objects)
# print(len(objects))
# Stores object categories that other objects can be placed upon (e.g. table, chair, couch, bed)
surface_properties_list = TDW_COCO_models['table'] + TDW_COCO_models['chair'] + \
TDW_COCO_models['bed'] + TDW_COCO_models['couch'] + \
TDW_COCO_models['bench'] + TDW_COCO_models['refrigerator']
surface_categories = []
for surface_properties in surface_properties_list:
surface_categories.append(surface_properties[0])
print("Surface Categories:", surface_categories)
# Stores the actual surface object instances/ids alongside number of objects on the surface
surface_object_ids = {}
self.start()
positions_list = [] # Stores current model locations and radii
scene_dimensions = [] # Store scene/environment dimensions
init_setup_commands = [{"$type": "set_screen_size",
"width": 640,
"height": 481},
{"$type": "set_render_quality",
"render_quality": 1}]
self.communicate(init_setup_commands)
scene_lib = SceneLibrarian()
# Disable physics when adding in new objects (objects float)
self.communicate({"$type": "simulate_physics",
"value": False})
for scene in scenes[1:]:
# Load in scene
# print("Scene", scene[0])
if scene[3] == "interior" and scene[0] == "box_room_2018":
self.start()
scene_record = scene_lib.get_record(scene[0])
self.communicate({"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()})
# Gets dimensions of environments (e.g. inside, outside) in the scene
# This command returns environment data in the form of a list of serialized byte arrays
scene_bytes = self.communicate({"$type": "send_environments",
"frequency": "once"})
# Iterating through data and parsing byte array
# Ignoring the last element (the frame count)
for b in scene_bytes[:-1]:
e = Environments(b)
for i in range(e.get_num()):
center = e.get_center(i)
bounds = e.get_bounds(i)
env_id = e.get_id(i)
scene_dimensions = [center, bounds, env_id] # Center, bounds are tuples
# Must come before set_pass_masks
avatar_position = TDWUtils.array_to_vector3([0.9 * scene_dimensions[1][0] / 2,
scene_dimensions[1][1] / 2,
0])
# print("Avatar Position:", avatar_position)
self.communicate(TDWUtils.create_avatar(avatar_id="avatar",
position=avatar_position,
look_at={"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0}))
# Set collision mode
self.communicate({"$type": "set_avatar_collision_detection_mode",
"mode": "continuous_speculative",
"avatar_id": "avatar"})
# Alter FOV
self.communicate({"$type": "set_field_of_view",
"field_of_view": 80,
"avatar_id": "avatar"})
# # Gets rid of header (Model: Category)
# objects = TDW_COCO_models[1:]
# random.shuffle(objects)
obj_count = 0
obj_overlaps = 0 # Number of failed attempts to place object due to over-dense objects area
while obj_count < objects_in_scene and obj_overlaps < 5:
# Handles if object has been added to a flat surface
added_to_surface = False
print("Object COUNT:", obj_count)
# Need to have random position for Bounds Data to return meaningful info
valid_obj_pos = {"x": random.uniform(-1 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][0] / 2),
"y": scene_dimensions[1][1] / 4,
"z": random.uniform(-0.9 * scene_dimensions[1][2] / 2,
0.9 * scene_dimensions[1][2] / 2)}
print("First random position")
# Add in the object at random position
# Object will later be removed or updated accordingly after performing collision calculations
record = ModelLibrarian(library="models_full.json").get_record(objects[obj_count][0])
print("Record gotten")
print(objects[obj_count][0])
o_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": TDWUtils.VECTOR3_ZERO,
"category": record.wcategory,
"id": obj_count})
print("Random first add")
# Returns bound data for added object
bounds_data = self.communicate({"$type": "send_bounds",
"frequency": "once"})
print("Bounds returned")
# Appends object, with information on position and obj_radius, to positions_list
# Length of buffer array should be 1
# print("Bounds Data:", bounds_data)
for b in bounds_data[:-1]:
# print("Buffer Loop:", b)
b_id = OutputData.get_data_type_id(b)
if b_id == "boun":
# print("BOUNDS")
o = Bounds(b)
# print("# of Objects:", o.get_num())
# print("# of Failed Attempts:", obj_overlaps)
# print("Buffer Array:", b)
# print("Bounds Object:", o)
for i in range(o.get_num()):
print("Object ID:", o.get_id(i))
print("obj_count:", obj_count)
print("Object:", objects[obj_count][0], "Category:", objects[obj_count][1])
# print("Object Center:", o.get_center(i))
# Only want to compute valid_position for object we are about to add
# Skip any computation if this is not the case
if o.get_id(i) != obj_count:
continue
# Useful for detecting if object fits in environment
# print("Calculating if object fits in environment")
width = distance.euclidean(o.get_left(i), o.get_right(i))
depth = distance.euclidean(o.get_front(i), o.get_back(i))
height = distance.euclidean(o.get_top(i), o.get_bottom(i))
# print("Width:", width)
# print("Depth:", depth)
# ("Height:", height)
# Useful for avoiding object overlap
# print("Calculating Object Bounds")
center_to_top = distance.euclidean(o.get_center(i), o.get_top(i))
center_to_bottom = distance.euclidean(o.get_center(i), o.get_bottom(i))
center_to_left = distance.euclidean(o.get_center(i), o.get_left(i))
center_to_right = distance.euclidean(o.get_center(i), o.get_right(i))
center_to_front = distance.euclidean(o.get_center(i), o.get_front(i))
center_to_back = distance.euclidean(o.get_center(i), o.get_back(i))
# Max object radius (center to diagonal of bounding box)
obj_radius = \
max(math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_back ** 2))
# print("Obj_Radius:", obj_radius)
# Set sweeping radius, based on scene plane dimensions
l_radius = random.uniform(0, min(0.5 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][2] / 2))
# Checking that object fits in scene viewing
if (width > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or
depth > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or
height > 0.7 * scene_dimensions[1][1]):
print("Object does not fit in scene")
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Not possible to find valid object position -- too many overlapping objects
elif (not self._get_object_position(scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[0]):
print("Could not calculate valid object location")
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
obj_overlaps += 1
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Find appropriate, non-overlapping object position
# Reset object position to the valid position
else:
print("Object fits in scene")
# Check if object fits on table, chair, couch, etc.
# Add object if it fits, place it somewhere on top of the surface
for surface_id in surface_object_ids.keys():
print("Surface ID:", surface_id)
# Skip placement feasibility if the object is already a surface-type object
# Ex. no chair on top of a table
if objects[obj_count][0] in surface_categories:
print("Object: %s is already a surface object" % objects[obj_count][0])
break
# Check how many objects are on surface
if surface_object_ids[surface_id] >= 3:
print("Too many objects on surface")
print("From surface objects dict:", surface_object_ids[surface_id])
continue
surface_bounds = self.get_bounds_data(surface_id)
surface_area = distance.euclidean(surface_bounds.get_left(0),
surface_bounds.get_right(0)) * \
distance.euclidean(surface_bounds.get_front(0),
surface_bounds.get_back(0))
obj_area = width * height
if obj_area < surface_area:
s_center_to_top = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_top(0))
s_center_to_bottom = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_bottom(0))
s_center_to_left = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_left(0))
s_center_to_right = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_right(0))
s_center_to_front = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_front(0))
s_center_to_back = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_back(0))
surface_radius = \
max(math.sqrt(
s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2))
print("Surface-type object")
self.communicate({"$type": "destroy_object",
"id": obj_count})
# Adding the object to the top of the surface
on_pos = surface_bounds.get_top(0)
on_y = on_pos[1]
on_pos = TDWUtils.get_random_point_in_circle(np.array(on_pos),
0.7 * surface_radius)
on_pos[1] = on_y
on_pos = TDWUtils.array_to_vector3(on_pos)
on_rot = {"x": 0, "y": random.uniform(-45, 45), "z": 0}
# Add the object.
print("Model Name on Surface:", objects[obj_count][0])
record = ModelLibrarian(library="models_full.json").get_record(
objects[obj_count][0])
on_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": on_pos,
"rotation": on_rot,
"category": record.wcategory,
"id": obj_count})
obj_count += 1
surface_object_ids[surface_id] += 1
object_ids.append(obj_count)
print("Object added on top of surface")
added_to_surface = True
# Breaking out of surface objects loop
break
if added_to_surface:
print("Breaking out of object loop")
# Breaking out of object loop
break
print("Post-surface")
valid_obj_pos = self._get_object_position(scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[1]
print("Position calculated")
positions_list.append(ObjectPosition(valid_obj_pos, obj_radius))
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
added_object_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": {"x": 0, "y": 0, "z": 0},
"category": record.wcategory,
"id": obj_count})
# print("Object ID:", added_object_id)
print("Regular object add")
object_ids.append(added_object_id)
# If TDW model belongs to surface categories, store id_information
if objects[obj_count][0] in surface_categories:
surface_object_ids[obj_count] = 0
# Rotate the object randomly
print("Rotating object")
self.communicate({"$type": "rotate_object_by",
"angle": random.uniform(-45, 45),
"axis": "yaw",
"id": obj_count,
"is_world": True})
# Minimal rotating for position differences
# Don't rotate the object if doing so will result in overlap into scene
if not (o.get_bottom(i)[1] < 0 or o.get_top(i)[1] > 0.9 * scene_dimensions[1][1]):
pitch_angle = random.uniform(-45, 45)
self.communicate({"$type": "rotate_object_by",
"angle": pitch_angle,
"axis": "pitch",
"id": obj_count,
"is_world": True})
roll_angle = random.uniform(-45, 45)
self.communicate({"$type": "rotate_object_by",
"angle": roll_angle,
"axis": "roll",
"id": obj_count,
"is_world": True})
# Don't need this for just changing positions
# Setting random materials/textures
# Looping through sub-objects and materials
sub_count = 0
for sub_object in record.substructure:
# Loop through materials in sub-objects
for j in range(len(sub_object)):
# Get random material and load in
material = random.choice(materials[1:])
self.load_material(material)
print("Material loaded")
# Set random material on material of sub-object
self.communicate({"$type": "set_visual_material",
"material_index": j,
"material_name": material[0],
"object_name": sub_object['name'],
"id": obj_count})
print("Material set")
sub_count += 1
if sub_count > 10:
break
break
print("Updating count")
obj_count += 1
print("Breaking out of object_id loop")
break
# Break out of buffer loop
print("Breaking out of buffer loop")
break
# Move onto next iteration of while loop (next object to load in)
print("Object added - next while loop iteration")
continue
# for i in range(200):
# self.communicate({"$type": "simulate_physics",
# "value": False})
# Enable image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
# Capture scene
# NOTE: THESE SCENES GET REPLACED IN THE TARGET DIRECTORY
scene_data = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": {"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0}})
images = Images(scene_data[0])
TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=path)
print("Object ids:", object_ids)
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)