def get_trial_initialization_commands(self) -> List[dict]:
commands = []
# randomization across trials
if not(self.randomize):
self.trial_seed = (self.MAX_TRIALS * self.seed) + self._trial_num
random.seed(self.trial_seed)
else:
self.trial_seed = -1 # not used
# Place target zone
commands.extend(self._place_target_zone())
# Choose and drop an object.
commands.extend(self._place_drop_object())
# Choose and place a middle object.
commands.extend(self._place_intermediate_object())
# Teleport the avatar to a reasonable position based on the drop height.
a_pos = self.get_random_avatar_position(radius_min=self.camera_radius_range[0],
radius_max=self.camera_radius_range[1],
angle_min=self.camera_min_angle,
angle_max=self.camera_max_angle,
y_min=self.drop_height * self.camera_min_height,
y_max=self.drop_height * self.camera_max_height,
center=TDWUtils.VECTOR3_ZERO)
cam_aim = {"x": 0, "y": self.drop_height * 0.5, "z": 0}
commands.extend([
{"$type": "teleport_avatar_to",
"position": a_pos},
{"$type": "look_at_position",
"position": cam_aim},
{"$type": "set_focus_distance",
"focus_distance": TDWUtils.get_distance(a_pos, cam_aim)}
])
# Set the camera parameters
self._set_avatar_attributes(a_pos)
self.camera_position = a_pos
self.camera_rotation = np.degrees(np.arctan2(a_pos['z'], a_pos['x']))
dist = TDWUtils.get_distance(a_pos, self.camera_aim)
self.camera_altitude = np.degrees(np.arcsin((a_pos['y'] - self.camera_aim['y'])/dist))
# For distractor placements
self.middle_scale = self.zone_scale
# Place distractor objects in the background
commands.extend(self._place_background_distractors(z_pos_scale=1))
# Place occluder objects in the background
commands.extend(self._place_occluders(z_pos_scale=1))
# test mode colors
if self.use_test_mode_colors:
self._set_test_mode_colors(commands)
return commands
def __init__(self, c_0: Dict[str, float], c_1: Dict[str, float]):
"""
:param c_0: Center of a subsector.
:param c_1: Center of a subsector.
"""
self.c_0 = TDWUtils.vector3_to_array(c_0)
self.c_1 = TDWUtils.vector3_to_array(c_1)
def _get_object_position(object_positions: List[ObjectPosition], max_tries: int = 1000, radius: float = 2) -> \
Dict[str, float]:
"""
Try to get a valid random position that doesn't interpentrate with other objects.
:param object_positions: The positions and radii of all objects so far that will be added to the scene.
:param max_tries: Try this many times to get a valid position before giving up.
:param radius: The radius to pick a position in.
:return: A valid position that doesn't interpentrate with other objects.
"""
o_pos = TDWUtils.array_to_vector3(
TDWUtils.get_random_point_in_circle(center=np.array([0, 0, 0]),
radius=radius))
# Pick a position away from other objects.
ok = False
count = 0
while not ok and count < max_tries:
count += 1
ok = True
for o in object_positions:
# If the object is too close to another object, try another position.
if TDWUtils.get_distance(o.position, o_pos) <= o.radius:
ok = False
o_pos = TDWUtils.array_to_vector3(
TDWUtils.get_random_point_in_circle(center=np.array(
[0, 0, 0]),
radius=radius))
return o_pos
def trial_loop(self,
num: int,
output_dir: str,
temp_path: str) -> None:
output_dir = Path(output_dir)
if not output_dir.exists():
output_dir.mkdir(parents=True)
temp_path = Path(temp_path)
if not temp_path.parent.exists():
temp_path.parent.mkdir(parents=True)
# Remove an incomplete temp path.
if temp_path.exists():
temp_path.unlink()
pbar = tqdm(total=num)
# Skip trials that aren't on the disk, and presumably have been uploaded; jump to the highest number.
exists_up_to = 0
for f in output_dir.glob("*.hdf5"):
if int(f.stem) > exists_up_to:
exists_up_to = int(f.stem)
if exists_up_to > 0:
print('Trials up to %d already exist, skipping those' % exists_up_to)
pbar.update(exists_up_to)
for i in range(exists_up_to, num):
filepath = output_dir.joinpath(TDWUtils.zero_padding(i, 4) + ".hdf5")
self.stimulus_name = '_'.join([filepath.parent.name, str(Path(filepath.name).with_suffix(''))])
if not filepath.exists():
# Save out images
self.png_dir = None
if any([pa in PASSES for pa in self.save_passes]):
self.png_dir = output_dir.joinpath("pngs_" + TDWUtils.zero_padding(i, 4))
if not self.png_dir.exists():
self.png_dir.mkdir(parents=True)
# Do the trial.
self.trial(filepath=filepath,
temp_path=temp_path,
trial_num=i)
# Save an MP4 of the stimulus
if self.save_movies:
for pass_mask in self.save_passes:
mp4_filename = str(filepath).split('.hdf5')[0] + pass_mask
cmd, stdout, stderr = pngs_to_mp4(
filename=mp4_filename,
image_stem=pass_mask[1:]+'_',
png_dir=self.png_dir,
size=[self._height, self._width],
overwrite=True,
remove_pngs=True,
use_parent_dir=False)
rm = subprocess.run('rm -rf ' + str(self.png_dir), shell=True)
pbar.update(1)
pbar.close()
def _push(self, position: Dict[str, float], target: Dict[str, float],
force_mag: float) -> List[dict]:
"""
:param position: The initial position.
:param target: The target position (used for the force vector).
:param force_mag: The force magnitutde.
:return: A list of commands to push the object along the floor.
"""
commands, soft_id = self._get_squishable(position=position,
rotation={
"x": 0,
"y":
random.uniform(0, 360),
"z": 0
})
# Get a force vector towards the target.
p0 = TDWUtils.vector3_to_array(position)
p1 = TDWUtils.vector3_to_array(target)
force = ((p1 - p0) / (np.linalg.norm(p1 - p0))) * force_mag
# Add the force.
commands.append({
"$type": "apply_force_to_flex_object",
"force": TDWUtils.array_to_vector3(force),
"id": soft_id
})
return commands
def run(self):
self.start()
width = 12
length = 12
# Create an empty room.
self.communicate(TDWUtils.create_empty_room(width, length))
# Disable physics.
self.communicate({"$type": "simulate_physics", "value": False})
# Add the avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": -5,
"y": 10,
"z": -2
},
look_at=TDWUtils.VECTOR3_ZERO))
# Create the dining table.
self.dining_table()
# Create the sofa.
self.sofa(width, length)
# Bake the NavMesh.
self.communicate({"$type": "bake_nav_mesh", "carve_type": "none"})
# Get a random position on the NavMesh.
position = TDWUtils.get_random_position_on_nav_mesh(
self, width, length)
position = TDWUtils.array_to_vector3(position)
# Add a side table.
self.side_table(position)
def run(self):
self.start()
init_setup_commands = [{
"$type": "set_screen_size",
"width": 1280,
"height": 962
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
self.communicate(init_setup_commands)
# Create an empty room.
self.load_streamed_scene(scene="tdw_room_2018")
# Disable physics.
self.communicate({"$type": "simulate_physics", "value": False})
# Add the avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": 0.9,
"y": 1.5,
"z": 0.9
},
look_at=TDWUtils.array_to_vector3(
[0, 0.8, 0]),
avatar_id="avatar"))
# Create the dining table.
self.dining_table()
def run(self):
self.start()
# Create an empty room.
self.communicate(TDWUtils.create_empty_room(12, 12))
# Create the avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": 0,
"y": 3,
"z": -4
},
look_at=TDWUtils.VECTOR3_ZERO))
# Set the pass mask to _depth only.
# Get an image.
resp = self.communicate([
self.get_add_object("rh10", object_id=0), {
"$type": "set_pass_masks",
"avatar_id": "a",
"pass_masks": ["_depth"]
}, {
"$type": "send_images",
"frequency": "once"
}
])
images = Images(resp[0])
# Get the depth values of each pixel.
depth = TDWUtils.get_depth_values(images.get_image(0))
print(depth)
def run(self, target: str):
"""
Take a screenshot per record.
:param target: The target model. If none, target all models.
"""
# Get an image of just one model.
if args.target:
records = [self._lib.get_record(target)]
# Get an image of all models.
else:
records = self._lib.records
pbar = tqdm(total=len(records))
for record in records:
if not self._record_is_ok(record):
pbar.update(1)
continue
record_name = record.name
# Skip an image that exists.
if not Path(self.output_dir).joinpath(record_name + ".jpg").exists():
# Capture a screenshot.
images, frame = self.get_image(record)
# Save the image.
TDWUtils.save_images(images, record_name,
output_directory=self.output_dir, append_pass=False)
pbar.update(1)
# Terminate the build.
self.communicate({"$type": "terminate"})
pbar.close()
def run(self):
self.start()
# Create an office.
self.communicate(create_office(12, 24, 4, 4, 6))
# Get a random position on the NavMesh.
x, y, z = TDWUtils.get_random_position_on_nav_mesh(self, 12, 24)
# Index of the next destination.
d = 0
# Create the avatar.
self.communicate(
TDWUtils.create_avatar(avatar_type="A_Nav_Mesh",
position={
"x": x,
"y": 0.5,
"z": 0
}))
# Teleport the avatar.
# Set the speed of the avatar.
# Set the destination of the avatar.
self.communicate([{
"$type": "set_nav_mesh_avatar",
"avatar_id": "a",
"speed": 2
}, {
"$type": "set_nav_mesh_avatar_destination",
"avatar_id": "a",
"destination": NavMesh.DESTINATIONS[d]
}])
# Go to 3 destinations.
i = 0
while i < 3:
resp = self.communicate({"$type": "do_nothing"})
# If there is only 1 element in the response, it is an empty frame.
if len(resp) == 1:
continue
# The avatar arrived at the destination!
else:
assert OutputData.get_data_type_id(resp[0]) == "anmd"
# Parse the output data.
data = ArrivedAtNavMeshDestination(resp[0])
print(data.get_avatar_id())
print("")
i += 1
d += 1
if d >= len(NavMesh.DESTINATIONS):
d = 0
# Set a new destination.
self.communicate({
"$type": "set_nav_mesh_avatar_destination",
"avatar_id": "a",
"destination": NavMesh.DESTINATIONS[d]
})
def get_p_1(self) -> Dict[str, float]:
"""
:return: A position in the c_1 sub-sector.
"""
return TDWUtils.array_to_vector3(
TDWUtils.get_random_point_in_circle(center=self.c_1,
radius=self.RADIUS))
def run(self):
self.start()
commands = [TDWUtils.create_empty_room(self.size, self.size)]
commands.extend(
TDWUtils.create_avatar(position={
"x": 0,
"y": 1.5,
"z": 0
},
avatar_id="a"))
commands.extend([{
"$type": "set_pass_masks",
"avatar_id": "a",
"pass_masks": ["_img"]
}, {
"$type": "send_images",
"frequency": "always"
}])
self.communicate(commands)
# Listen for keys.
# Turn left.
self.listen(key="left",
commands={
"$type": "rotate_sensor_container_by",
"axis": "yaw",
"angle": self.angle * -1,
"sensor_name": "SensorContainer",
"avatar_id": "a"
})
# Turn right.
self.listen(key="right",
commands={
"$type": "rotate_sensor_container_by",
"axis": "yaw",
"angle": self.angle,
"sensor_name": "SensorContainer",
"avatar_id": "a"
})
# Quit the application.
self.listen(key="esc", function=self.stop)
# Continue until the quit key is pressed.
i = 0
while not self.done:
# Listen for keyboard input. Receive output data.
resp = self.communicate([])
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Save images.
if r_id == "imag":
TDWUtils.save_images(
images=Images(r),
output_directory=self.images_directory,
filename=TDWUtils.zero_padding(i, width=4))
# Increment the image number.
i += 1
self.communicate({"$type": "terminate"})
def run(self):
# Create the output directory.
output_directory = "example_output"
if os.path.exists(output_directory):
shutil.rmtree(output_directory)
sleep(0.5)
os.mkdir(output_directory)
print(f"Images will be saved to: {output_directory}")
self.start()
# Create an empty room.
self.communicate(TDWUtils.create_empty_room(12, 12))
# Add the objects.
lamp_id = self.add_object("alma_floor_lamp",
position={"x": 1, "y": 0, "z": 0},
rotation={"x": 0, "y": 90, "z": 0})
self.add_object("live_edge_coffee_table",
position={"x": 1.5, "y": 0, "z": 1.5},
rotation={"x": 0, "y": 30, "z": 0})
self.add_object("small_table_green_marble",
position={"x": -0.9, "y": 0, "z": -1.35})
# Create the avatar.
self.communicate({"$type": "create_avatar",
"type": "A_Img_Caps_Kinematic",
"id": "a"})
# Try to find a valid position on the NavMesh.
x, y, z = TDWUtils.get_random_position_on_nav_mesh(self, 12, 12)
# Teleport the avatar to the valid position.
# Apply a force to the lamp.
# Set the pass masks to _img.
# Enable image capture.
self.communicate([{"$type": "teleport_avatar_to",
"avatar_id": "a",
"position": {"x": x, "y": 1.5, "z": z}},
{"$type": "apply_force_to_object",
"force": {"x": 2, "y": 1, "z": 0},
"id": lamp_id},
{"$type": "set_pass_masks",
"avatar_id": "a",
"pass_masks": ["_img", "_id"]},
{"$type": "send_images",
"frequency": "always"},
])
# Capture 100 images.
for i in range(100):
# Look at the lamp.
resp = self.communicate({"$type": "look_at",
"avatar_id": "a",
"object_id": lamp_id,
"use_centroid": True})
images = Images(resp[0])
# Save the image.
TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=output_directory)
def run(self):
self.start()
init_setup_commands = [{"$type": "set_screen_size",
"width": 1280,
"height": 962},
{"$type": "set_render_quality",
"render_quality": 5}]
self.communicate(init_setup_commands)
# Create an empty room.
self.load_streamed_scene(scene="tdw_room_2018")
# Disable physics.
self.communicate({"$type": "set_gravity",
"value": False})
# Add the avatar.
self.communicate(TDWUtils.create_avatar(position={"x": 0, "y": 1, "z": 0},
look_at=TDWUtils.array_to_vector3([2, 0.8, -1.5]),
avatar_id="avatar"))
# self.communicate({"$type": "set_field_of_view",
# "field_of_view": 68.0,
# "avatar_id": "avatar"})
table_id = self.add_object(model_name="glass_table",
position={"x": 2, "y": 0, "z": -1.5},
library="models_full.json")
table_bounds = self.get_bounds_data(table_id)
top = table_bounds.get_top(0)
self.add_object(model_name="flower_vase_set_2",
position={"x": 1.6, "y": top[1], "z": -1.5},
library="models_full.json")
self.add_object(model_name="flower_vase_set_2",
position={"x": 2, "y": top[1], "z": -1.5},
library="models_full.json")
self.add_object(model_name="flower_vase_set_2",
position={"x": 2.4, "y": top[1], "z": -1.5},
library="models_full.json")
# Enable image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
scene_data = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": TDWUtils.array_to_vector3([2, 0.8, -1.5])})
images = Images(scene_data[0])
TDWUtils.save_images(images, "flowers_row", output_directory="replicated_images")
def create_scene(self):
"""
Initialize a blank room.
"""
self.start()
commands = [TDWUtils.create_empty_room(12, 12)]
commands.extend(TDWUtils.create_avatar(position={"x": 5, "y": 10, "z": 5}, look_at=TDWUtils.VECTOR3_ZERO))
self.communicate(commands)
def is_done(self, resp: List[bytes], frame: int) -> bool:
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# If the ball reaches or overshoots the destination, the trial is done.
if r_id == "tran":
t = Transforms(r)
d0 = TDWUtils.get_distance(
TDWUtils.array_to_vector3(t.get_position(0)), self._p0)
d1 = TDWUtils.get_distance(self._p0, self._p1)
return d0 > d1 * 1.5
return False
def communicate(self, commands: Union[dict, List[dict]]) -> List[bytes]:
resp = super().communicate(commands)
# Save images per frame.
images = get_data(resp=resp, d_type=Images)
if images is not None:
TDWUtils.save_images(images=images,
filename=TDWUtils.zero_padding(
self.frame_count, width=4),
output_directory=self.output_dir)
self.frame_count += 1
return resp
def run(self, drag_on_land, angular_drag_on_land):
self.start()
self.communicate(TDWUtils.create_empty_room(40, 40))
# Create the avatar.
self.communicate(TDWUtils.create_avatar(avatar_type="A_Simple_Body",
position={"x": 0, "y": 1, "z": 0}))
# When the y value of the avatar's position is this value, the avatar is on the ground.
ground_y = 1
# When the y value of the avatar's position exceeds ground_y by this much, it is in the air.
in_air_threshold = 0.05
# If true, the avatar is in the air.
is_in_air = False
# Set high default drag values.
# Set a low mass.
# Apply a force.
# Request avatar data.
resp = self.communicate([{"$type": "set_avatar_drag",
"avatar_id": "a",
"drag": 30,
"angular_drag": 80},
{"$type": "send_avatars",
"frequency": "always"},
{"$type": "apply_force_to_avatar",
"avatar_id": "a",
"direction": {"x": 0.3, "y": 0.7, "z": 0},
"magnitude": 900}])
for i in range(500):
x, y, z = AvatarSimpleBody(resp[0]).get_position()
# If the avatar just launched in to the air, set its drag values to 0.
if y > ground_y + in_air_threshold and not is_in_air:
is_in_air = True
resp = self.communicate({"$type": "set_avatar_drag",
"avatar_id": "a",
"drag": 0,
"angular_drag": 0})
print("The avatar is in the air on frame: " + str(Controller.get_frame(resp[-1])))
# If the avatar just landed on the ground, set drag values.
elif y <= ground_y + in_air_threshold and is_in_air:
resp = self.communicate({"$type": "set_avatar_drag",
"avatar_id": "a",
"drag": drag_on_land,
"angular_drag": angular_drag_on_land})
is_in_air = False
print("The avatar is on the ground on frame: " + str(Controller.get_frame(resp[-1])))
# Wait.
else:
resp = self.communicate({"$type": "do_nothing"})
def run(self):
self.start()
# Create the room from an image.
self.communicate(TDWUtils.create_room_from_image("room.png"))
self.communicate({"$type": "set_post_process", "value": False})
self.communicate(
TDWUtils.create_avatar(position={
"x": 0,
"y": 40,
"z": 0
},
look_at=TDWUtils.VECTOR3_ZERO))
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)
width = 8
length = 8
# Create an empty room.
self.load_streamed_scene("box_room_2018")
# Disable physics.
self.communicate({"$type": "set_gravity",
"value": False})
# Add the avatar.
self.communicate(TDWUtils.create_avatar(position={"x": 0, "y": 1.7, "z": 0},
look_at={"x": 0.3, "y": 0.8, "z": 2.2},
avatar_id="avatar"))
self.communicate({"$type": "set_field_of_view",
"field_of_view": 68.0,
"avatar_id": "avatar"})
self.add_object(model_name="b06_bikenew",
position={"x": 0.3, "y": 0, "z": 2.2},
rotation={"x": 0, "y": 180, "z": 0},
library="models_full.json")
self.add_object(model_name="animal_dog_rtsit_1280",
position={"x": 0.35, "y": 0, "z": 1.8},
library="models_full.json")
# Enable image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
scene_data = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": {"x": 0.3, "y": 0.8, "z": 2.2}})
images = Images(scene_data[0])
TDWUtils.save_images(images, "dog_bike", output_directory="replicated_images/interior")
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", "_albedo"],
"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()
self.communicate(TDWUtils.create_empty_room(20, 20))
self.communicate(TDWUtils.create_avatar(position={"x": 0, "y": 3, "z": -6},
look_at=TDWUtils.VECTOR3_ZERO))
model_name = "rh10"
z = -3
x = -1.5
print("With the add_object command (complex syntax but you have maximum control):")
record = ModelLibrarian().get_record(model_name)
self.communicate({"$type": "add_object",
"name": model_name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": {"x": x, "y": 0, "z": z},
"rotation": TDWUtils.VECTOR3_ZERO,
"category": record.wcategory,
"id": self.get_unique_id()})
x = 0
print("With the wrapper function Controller.add_object() "
"(easy to use, but you can't add additional commands to this frame):")
self.add_object(model_name=model_name,
position={"x": x, "y": 0, "z": z},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_core.json")
x = 1.5
print("With the wrapper function Controller.get_add_object() "
"(harder to use, but you can add commands to this frame):")
self.communicate(self.get_add_object(model_name=model_name,
object_id=self.get_unique_id(),
position={"x": x, "y": 0, "z": 0},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_core.json"))
print("With the add_object command, minus all optional parameters (the model won't scale properly!):")
self.communicate({"$type": "add_object",
"name": model_name,
"url": record.get_url(),
"id": self.get_unique_id()})
print("With the wrapper function Controller.add_object(), minus all optional parameters:")
self.add_object(model_name)
print("With the wrapper function Controller.get_add_object(), minus all optional parameters:")
self.communicate(self.get_add_object(model_name=model_name,
object_id=self.get_unique_id()))
def get_rotated_target(self, target: np.array) -> np.array:
"""
Rotate the target by the avatar's forward directional vector.
:param target: The target position.
:return: The rotated position.
"""
angle = TDWUtils.get_angle_between(v1=FORWARD,
v2=self.frame.get_forward())
return TDWUtils.rotate_position_around(position=target -
self.frame.get_position(),
angle=-angle)
def push_into_other(self) -> List[dict]:
"""
:return: A list of commands to push one object into another object.
"""
# Get two points some distance from each other.
p0 = np.array([0, 0, 0])
p1 = np.array([0, 0, 0])
center = np.array([0, 0, 0])
count = 0
while count < 1000 and np.linalg.norm(np.abs(p1 - p0)) < 1.5:
p0 = TDWUtils.get_random_point_in_circle(center=center, radius=2)
p1 = TDWUtils.get_random_point_in_circle(center=center, radius=2)
count += 1
p0 = TDWUtils.array_to_vector3(p0)
p0["y"] = 0.2
p1 = TDWUtils.array_to_vector3(p1)
p1["y"] = 0.1
# Get commands to create the first object and push it.
commands = self._push(position=p0,
target=p1,
force_mag=random.uniform(1000, 2000))
# Add commands for the second object.
second_object_commands, s_id = self._get_squishable(position=p1,
rotation={
"x":
0,
"y":
random.uniform(
0, 360),
"z":
0
})
commands.extend(second_object_commands)
# Set the avatar.
# Get the midpoint between the objects and use this to place the avatar.
p_med = TDWUtils.array_to_vector3(
np.array([(p0["x"] + p1["x"]) / 2, 0, (p0["z"] + p1["z"]) / 2]))
commands.extend(
self._set_avatar(a_pos=self.get_random_avatar_position(
radius_min=1.8,
radius_max=2.3,
y_min=1,
y_max=1.75,
center=p_med),
cam_aim=p1))
return commands
def get_trial_initialization_commands(self) -> List[dict]:
commands = []
# place a ramp
commands.extend(self._place_ramp_object())
# Choose and place a target object.
# commands.extend(self._place_target_object())
# Choose and drop an object.
#commands.extend(self._place_drop_object())
# Teleport the avatar to a reasonable position based on the drop height.
a_pos = self.get_random_avatar_position(radius_min=0.5,
radius_max=1,
angle_min=0,
angle_max=10,
y_min=0.5,
y_max=1,
center=TDWUtils.VECTOR3_ZERO)
print("avatar pos", a_pos)
cam_aim = {"x": 0, "y": 0, "z": 0}
commands.extend([
{"$type": "teleport_avatar_to",
"position": a_pos},
{"$type": "look_at_position",
"position": cam_aim},
{"$type": "set_focus_distance",
"focus_distance": TDWUtils.get_distance(a_pos, cam_aim)}
])
return commands
def _set_avatar(a_pos: Dict[str, float],
cam_aim: Dict[str, float]) -> List[dict]:
"""
:param a_pos: The avatar position.
:param cam_aim: The camera aim point.
:return: A list of commands to teleport the avatar and rotate the sensor container.
"""
return [{
"$type": "teleport_avatar_to",
"position": a_pos
}, {
"$type": "look_at_position",
"position": cam_aim
}, {
"$type": "set_focus_distance",
"focus_distance": TDWUtils.get_distance(a_pos, cam_aim)
}, {
"$type": "rotate_sensor_container_by",
"axis": "pitch",
"angle": random.uniform(-3, 3)
}, {
"$type": "rotate_sensor_container_by",
"axis": "yaw",
"angle": random.uniform(-3, 3)
}]
def get_falling_commands(self, mass: float = 3) -> List[List[dict]]:
"""
:param mass: Objects with <= this mass might receive a force.
:return: A list of lists; per-frame commands to make small objects fly up.
"""
per_frame_commands: List[List[dict]] = []
# Get a list of all small objects.
small_ids = self.get_objects_by_mass(mass)
random.shuffle(small_ids)
max_num_objects = len(small_ids) if len(small_ids) < 8 else 8
min_num_objects = max_num_objects - 3
if min_num_objects <= 0:
min_num_objects = 1
# Add some objects.
for i in range(random.randint(min_num_objects, max_num_objects)):
o_id = small_ids.pop(0)
force_dir = np.array([random.uniform(-0.125, 0.125), random.uniform(0.7, 1), random.uniform(-0.125, 0.125)])
force_dir = force_dir / np.linalg.norm(force_dir)
min_force = self.physics_info[o_id].mass * 2
max_force = self.physics_info[o_id].mass * 4
force = TDWUtils.array_to_vector3(force_dir * random.uniform(min_force, max_force))
per_frame_commands.append([{"$type": "apply_force_to_object",
"force": force,
"id": o_id}])
# Wait some frames.
for j in range(10, 30):
per_frame_commands.append([])
return per_frame_commands
def _place_camera(self) -> List[dict]:
commands = []
a_pos = self.get_random_avatar_position(
radius_min=self.camera_radius_range[0],
radius_max=self.camera_radius_range[1],
angle_min=self.camera_min_angle,
angle_max=self.camera_max_angle,
y_min=self.camera_min_height,
y_max=self.camera_max_height,
center=TDWUtils.VECTOR3_ZERO,
reflections=self.camera_left_right_reflections)
self._set_avatar_attributes(a_pos)
commands.extend([{
"$type": "teleport_avatar_to",
"position": self.camera_position
}, {
"$type": "look_at_position",
"position": self.camera_aim
}, {
"$type":
"set_focus_distance",
"focus_distance":
TDWUtils.get_distance(a_pos, self.camera_aim)
}])
return commands
def run():
# Create the asset bundle and the record.
asset_bundle_paths, record_path = AssetBundleCreator(
).create_asset_bundle("cube.fbx", True, 123, "", 1)
# Get the name of the bundle for this platform. For example, Windows -> "StandaloneWindows64"
bundle = SYSTEM_TO_UNITY[system()]
# Get the correct asset bundle path.
for p in asset_bundle_paths:
# Get the path to the asset bundle.
if bundle in str(p.parent.resolve()):
url = "file:///" + str(p.resolve())
# Launch the controller.
c = Controller()
c.start()
# Create the environment.
# Add the object.
commands = [{
"$type": "create_empty_environment"
}, {
"$type": "add_object",
"name": "cube",
"url": url,
"scale_factor": 1,
"id": c.get_unique_id()
}]
# Create the avatar.
commands.extend(
TDWUtils.create_avatar(position={
"x": 0,
"y": 0,
"z": -3.6
}))
c.communicate(commands)
return
def get_per_frame_commands(self, resp: List[bytes], frame) -> List[dict]:
commands = []
# Apply a force.
if frame < self._num_force_frames:
for r in resp[:-1]:
if FlexParticles.get_data_type_id(r) == "flex":
fp = FlexParticles(r)
for i in range(fp.get_num_objects()):
# Find the cloth.
if fp.get_id(i) == self.cloth_id:
forces = []
p_id = 0
for p in fp.get_particles(i):
# Add a force if this is a "corner particle".
if np.abs(np.linalg.norm(p[:-1] - self._corner)
) <= self._corner_radius:
# Calculate the force.
pos = np.array(p[:-1])
force = ((pos - self._corner) / np.linalg.norm(pos - self._corner)) * self.\
_force_per_frame
# Add the force and particle ID.
forces.extend(force)
forces.append(p_id)
p_id += 1
# Encode and send the force.
commands.extend([{
"$type":
"apply_forces_to_flex_object_base64",
"forces_and_ids_base64":
TDWUtils.get_base64_flex_particle_forces(
forces),
"id":
self.cloth_id
}])
return commands
