def add_transforms_object(self,
record: ModelRecord,
position: Dict[str, float],
rotation: Dict[str, float],
o_id: Optional[int] = None) -> dict:
"""
This is a wrapper for `Controller.get_add_object()` and the `add_object` command.
This caches the ID of the object so that it can be easily cleaned up later.
:param record: The model record.
:param position: The initial position of the object.
:param rotation: The initial rotation of the object, in Euler angles.
:param o_id: The unique ID of the object. If None, a random ID is generated.
:return: An `add_object` command.
"""
if o_id is None:
o_id: int = Controller.get_unique_id()
# Log the static data.
self.object_ids = np.append(self.object_ids, o_id)
return {
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": position,
"rotation": rotation,
"category": record.wcategory,
"id": o_id
}
def 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 _init_object(c: Controller, name: str, pos: Dict[str, float], rot: Dict[str, float]) -> List[dict]:
"""
:param c: The controller.
:param name: The name of the model.
:param pos: The initial position of the model.
:param rot: The initial rotation of the model.
:return: A list of commands to instantiate an object from ObjectInfo values.
"""
o_id = c.get_unique_id()
Scene.OBJECT_IDS.update({o_id: name})
info = Scene._OBJECT_INFO[name]
return [c.get_add_object(name,
object_id=o_id,
position=pos,
rotation=rot,
library=info.library),
{"$type": "set_mass",
"id": o_id,
"mass": info.mass},
{"$type": "set_physic_material",
"id": o_id,
"bounciness": info.bounciness,
"static_friction": 0.1,
"dynamic_friction": 0.8}]
def get_image(self, record: ModelRecord):
o_id = Controller.get_unique_id()
self.communicate({"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"rotation": record.canonical_rotation,
"id": o_id})
s = TDWUtils.get_unit_scale(record) * 2
# Scale the model and get an image.
# Look at the model's centroid.
resp = self.communicate([{"$type": "scale_object",
"id": o_id,
"scale_factor": {"x": s, "y": s, "z": s}},
{"$type": "look_at",
"avatar_id": "a",
"object_id": o_id,
"use_centroid": True}])
# Destroy the model and unload the asset bundle.
self.communicate([{"$type": "destroy_object",
"id": o_id},
{"$type": "unload_asset_bundles"}])
return Images(resp[0]), resp[-1]
def __init__(self, model_name: str, initial_position: Dict[str, float]):
"""
:param model_name: The name of the 3D model in the model library.
:param initial_position: The initial position of the object.
"""
self.object_id = Controller.get_unique_id()
self.initial_position = initial_position
self.model_name = model_name
def get_commands(self) -> Tuple[int, List[dict]]:
"""
:return: Tuple: The ID of the object; a list of commands to create the object: `[add_object, rotate_object_to, scale_object, set_kinematic_state, set_object_collision_detection_mode]`
"""
record = TransformInitData.LIBRARIES[self.library].get_record(
name=self.name)
object_id = Controller.get_unique_id()
commands = [{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": self.position,
"category": record.wcategory,
"id": object_id
}]
# The rotation is a quaternion.
if "w" in self.rotation:
commands.append({
"$type": "rotate_object_to",
"rotation": self.rotation,
"id": object_id
})
# The rotation is in Euler angles.
else:
commands.append({
"$type": "rotate_object_to_euler_angles",
"euler_angles": self.rotation,
"id": object_id
})
commands.extend([{
"$type": "scale_object",
"scale_factor": self.scale_factor,
"id": object_id
}, {
"$type": "set_kinematic_state",
"id": object_id,
"is_kinematic": self.kinematic,
"use_gravity": self.gravity
}])
# Kinematic objects must be continuous_speculative.
if self.kinematic:
commands.append({
"$type": "set_object_collision_detection_mode",
"id": object_id,
"mode": "continuous_speculative"
})
return object_id, commands
def _initialize_scene(self, c: Controller) -> List[dict]:
commands = super()._initialize_scene(c)
model_name = self._get_model_name()
o_id = c.get_unique_id()
Scene.OBJECT_IDS.update({o_id: model_name})
commands.extend([c.get_add_object(model_name, object_id=o_id, library=self._get_library()),
{"$type": "set_mass",
"id": o_id,
"mass": 1000},
{"$type": "set_physic_material",
"id": o_id,
"bounciness": Scene._OBJECT_INFO[model_name].bounciness,
"static_friction": 0.1,
"dynamic_friction": 0.8}])
return commands
def precalculate_movements() -> None:
"""
Precalculate random movements for the avatar before running this test.
This way the movements are "random" per frame but the same per trial.
"""
frames: List[List[dict]] = list()
for i in range(num_frames):
# Move and turn randomly.
commands = [{
"$type": "move_avatar_forward_by",
"magnitude": uniform(50, 100)
}, {
"$type": "rotate_avatar_by",
"angle": uniform(-45, 45),
"axis": "yaw",
"is_world": True,
"avatar_id": "a"
}]
# Bend arm joints randomly.
for j in Baby.JOINTS:
commands.append({
"$type": "bend_arm_joint_by",
"angle": uniform(0, 89),
"joint": j.joint,
"axis": j.axis
})
for is_left in [True, False]:
if random() > 0.5:
commands.append({
"$type": "pick_up_proximity",
"distance": 3,
"grip": 1000,
"is_left": is_left
})
else:
commands.append({"$type": "put_down", "is_left": is_left})
frames.append(commands)
movements = {"id": Controller.get_unique_id(), "frames": frames}
VarianceAvatar.FRAME_PATH.write_text(json.dumps(movements))
def _initialize_scene(self, c: Controller) -> List[dict]:
c.model_librarian = ModelLibrarian("models_full.json")
commands = super()._initialize_scene(c)
# Add a table, chair, and boxes.
commands.extend(self._init_object(c, "b05_table_new",
pos={"x": 0, "y": 0, "z": 4.33},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "chair_willisau_riale",
pos={"x": 0, "y": 0, "z": 3.7},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "iron_box",
pos={"x": 0.13, "y": 0.65, "z": 4.83},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "iron_box",
pos={"x": -0.285, "y": 1.342, "z": 4.79},
rot={"x": 90, "y": 0, "z": 0}))
# Add a shelf with a custom scale.
shelf_id = c.get_unique_id()
shelf_name = "metal_lab_shelf"
Scene.OBJECT_IDS.update({shelf_id: shelf_name})
commands.extend([c.get_add_object(shelf_name,
object_id=shelf_id,
rotation={"x": 0, "y": -90, "z": 0},
position={"x": 0, "y": 0, "z": 4.93}),
{"$type": "set_mass",
"id": shelf_id,
"mass": 400},
{"$type": "set_physic_material",
"id": shelf_id,
"bounciness": Scene._OBJECT_INFO[shelf_name].bounciness,
"static_friction": 0.1,
"dynamic_friction": 0.8},
{"$type": "scale_object",
"id": shelf_id,
"scale_factor": {"x": 1, "y": 1.5, "z": 1.8}}])
return commands
def get_trial_initialization_commands(self) -> List[dict]:
self._table_id = Controller.get_unique_id()
# Teleport the avatar.
commands = [{
"$type": "teleport_avatar_to",
"position": random.choice(_TableScripted._CAMERA_POSITIONS)
}, {
"$type": "look_at_position",
"position": {
"x": -10.8,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.5
}
}]
# Add the table.
commands.extend(
self.add_physics_object_default(name="quatre_dining_table",
position=self._TABLE_POSITION,
rotation={
"x": 0,
"y": -90,
"z": 0
},
o_id=self._table_id))
x0 = -9.35
x1 = -9.9
x2 = -10.15
x3 = -10.8
x4 = -11.55
x5 = -11.8
x6 = -12.275
z0 = -4.5
z1 = -5.05
z2 = -5.45
z3 = -5.95
z4 = -6.475
# Add chairs.
for x, z, rot in zip([x2, x3, x4, x6, x4, x3, x2, x0],
[z0, z0, z0, z2, z4, z4, z4, z2],
[180, 180, 180, 90, 0, 0, 0, -90]):
commands.extend(
self.add_physics_object_default(
name="brown_leather_dining_chair",
position={
"x": x,
"y": _TableScripted._FLOOR_HEIGHT,
"z": z
},
rotation={
"x": 0,
"y": rot,
"z": 0
}))
# Add plates.
plates_x = [x2, x3, x4, x5, x4, x3, x2, x1]
plates_z = [z1, z1, z1, z2, z3, z3, z3, z2]
for x, z in zip(plates_x, plates_z):
commands.extend(
self.add_physics_object_default(
name="plate05",
position={
"x": x,
"y": _TableScripted._TABLE_HEIGHT,
"z": z
},
rotation=TDWUtils.VECTOR3_ZERO))
# Cutlery offset from plate (can be + or -).
cutlery_off = 0.1
cutlery_rots = [180, 180, 180, 90, 0, 0, 0, -90]
# Forks.
fork_x_offsets = [
cutlery_off, cutlery_off, cutlery_off, 0, -cutlery_off,
-cutlery_off, -cutlery_off, 0
]
fork_z_offsets = [0, 0, 0, cutlery_off, 0, 0, 0, -cutlery_off]
forks_x = list(map(add, plates_x, fork_x_offsets))
forks_z = list(map(add, plates_z, fork_z_offsets))
for x, z, rot in zip(forks_x, forks_z, cutlery_rots):
commands.extend(
self.add_physics_object_default(
name="vk0010_dinner_fork_subd0",
position={
"x": x,
"y": _TableScripted._TABLE_HEIGHT,
"z": z
},
rotation={
"x": 0,
"y": rot,
"z": 0
}))
# Knives.
knife_x_offsets = [
-cutlery_off, -cutlery_off, -cutlery_off, 0, cutlery_off,
cutlery_off, cutlery_off, 0
]
knife_z_offsets = [0, 0, 0, -cutlery_off, 0, 0, 0, cutlery_off]
knives_x = list(map(add, plates_x, knife_x_offsets))
knives_z = list(map(add, plates_z, knife_z_offsets))
for x, z, rot in zip(knives_x, knives_z, cutlery_rots):
commands.extend(
self.add_physics_object_default(
name="vk0007_steak_knife",
position={
"x": x,
"y": _TableScripted._TABLE_HEIGHT,
"z": z
},
rotation={
"x": 0,
"y": rot,
"z": 0
}))
incidental_names = [
"moet_chandon_bottle_vray", "peppermill", "salt_shaker",
"coffeemug", "coffeecup004", "bowl_wood_a_01", "glass1", "glass2",
"glass3"
]
incidental_positions = [{
"x": -10.35,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.325
}, {
"x": -10.175,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.635
}, {
"x": -11.15,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.85
}, {
"x": -11.525,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.625
}, {
"x": -11.25,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.185
}, {
"x": -10.5,
"y": _TableScripted._TABLE_HEIGHT,
"z": -5.05
}]
random.shuffle(incidental_names)
random.shuffle(incidental_positions)
# Select 4 incidental objects.
for i in range(4):
name = incidental_names.pop(0)
o_id = Controller.get_unique_id()
commands.extend(
self.add_physics_object_default(
name=name,
position=incidental_positions.pop(0),
rotation=TDWUtils.VECTOR3_ZERO,
o_id=o_id))
# These objects need further scaling.
if name in ["salt_shaker", "peppermill"]:
commands.append({
"$type": "scale_object",
"id": o_id,
"scale_factor": {
"x": 0.254,
"y": 0.254,
"z": 0.254
}
})
# Select 2 bread objects.
bread_names = ["bread", "bread_01", "bread_02", "bread_03"]
bread_positions = [{
"x": x2,
"y": _TableScripted._TABLE_HEIGHT,
"z": z2
}, {
"x": x4,
"y": _TableScripted._TABLE_HEIGHT,
"z": z2
}, {
"x": x3,
"y": _TableScripted._TABLE_HEIGHT,
"z": z2
}]
random.shuffle(bread_names)
random.shuffle(bread_positions)
for i in range(2):
commands.extend(
self.add_physics_object_default(
name=bread_names.pop(0),
position=bread_positions.pop(0),
rotation=TDWUtils.VECTOR3_ZERO))
return commands
def get_trial_initialization_commands(self) -> List[dict]:
num_objects = random.choice([2, 3])
# Positions where objects will be placed (used to prevent interpenetration).
object_positions: List[ObjectPosition] = []
# Randomize the order of the records and pick the first one.
# This way, the objects are always different.
random.shuffle(self.records)
commands = []
# Add 2-3 objects.
for i in range(num_objects):
o_id = Controller.get_unique_id()
record = self.records[i]
# Set randomized physics values and update the physics info.
scale = TDWUtils.get_unit_scale(record) * random.uniform(0.8, 1.1)
# Get a random position.
o_pos = self._get_object_position(
object_positions=object_positions)
# Add the object and the radius, which is defined by its scale.
object_positions.append(
ObjectPosition(position=o_pos, radius=scale))
commands.extend(
self.add_physics_object(o_id=o_id,
record=self.records[i],
position=self._get_object_position(
object_positions=object_positions),
rotation={
"x": 0,
"y": random.uniform(-90, 90),
"z": 0
},
mass=random.uniform(1, 5),
dynamic_friction=random.uniform(
0, 0.9),
static_friction=random.uniform(0, 0.9),
bounciness=random.uniform(0, 1)))
# Scale the object.
commands.append({
"$type": "scale_object",
"id": o_id,
"scale_factor": {
"x": scale,
"y": scale,
"z": scale
}
})
# Point one object at the center, and then offset the rotation.
# Apply a force allow the forward directional vector.
# Teleport the avatar and look at the object that will be hit. Then slightly rotate the camera randomly.
# Listen for output data.
force_id = int(self.object_ids[0])
self._target_id = int(self.object_ids[1])
commands.extend([{
"$type": "object_look_at",
"other_object_id": self._target_id,
"id": force_id
}, {
"$type": "rotate_object_by",
"angle": random.uniform(-5, 5),
"id": force_id,
"axis": "yaw",
"is_world": True
}, {
"$type": "apply_force_magnitude_to_object",
"magnitude": random.uniform(20, 60),
"id": force_id
}, {
"$type":
"teleport_avatar_to",
"position":
self.get_random_avatar_position(radius_min=0.9,
radius_max=1.5,
y_min=0.5,
y_max=1.25,
center=TDWUtils.VECTOR3_ZERO)
}, {
"$type": "look_at",
"object_id": self._target_id,
"use_centroid": True
}, {
"$type": "rotate_sensor_container_by",
"axis": "pitch",
"angle": random.uniform(-5, 5)
}, {
"$type": "rotate_sensor_container_by",
"axis": "yaw",
"angle": random.uniform(-5, 5)
}, {
"$type": "focus_on_object",
"object_id": self._target_id
}])
return commands
def get_trial_initialization_commands(self) -> List[dict]:
self._table_id = Controller.get_unique_id()
self._a_pos = self.get_random_avatar_position(radius_min=1.7, radius_max=2.3, y_min=1.8, y_max=2.5,
center=TDWUtils.VECTOR3_ZERO)
# Teleport the avatar.
commands = []
# Add the table.
table_name = random.choice(self._TABLES)
commands.extend(self.add_physics_object_default(name=table_name,
position=TDWUtils.VECTOR3_ZERO,
rotation=TDWUtils.VECTOR3_ZERO,
o_id=self._table_id))
commands.extend([{"$type": "teleport_avatar_to",
"position": self._a_pos},
{"$type": "look_at",
"object_id": self._table_id,
"use_centroid": True},
{"$type": "focus_on_object",
"object_id": self._table_id,
"use_centroid": True}])
table_record = PHYSICS_INFO[table_name].record
top = table_record.bounds["top"]
# Select random model names.
chair_name = random.choice(self._CHAIRS)
plate_name = random.choice(self._PLATES)
fork_name = random.choice(self._FORKS)
spoon_name = random.choice(self._SPOONS)
knife_name = random.choice(self._KNIVES)
cup_name = random.choice(self._CUPS)
# Get the chair positions.
setting_positions = [table_record.bounds["left"],
table_record.bounds["right"],
table_record.bounds["front"],
table_record.bounds["back"]]
# Add 4 chairs around the table and their table settings.
for setting_pos, s_p in zip(setting_positions, self._SETTINGS):
chair_id = Controller.get_unique_id()
chair_pos = {"x": setting_pos["x"], "y": 0, "z": setting_pos["z"]}
commands.extend(self.add_physics_object_default(position=chair_pos,
rotation=TDWUtils.VECTOR3_ZERO,
o_id=chair_id,
name=chair_name))
# Move the chair back a bit.
chair_pos = get_move_along_direction(pos=chair_pos,
target=TDWUtils.VECTOR3_ZERO,
d=-random.uniform(0.4, 0.55),
noise=0.01)
commands.append({"$type": "teleport_object",
"id": chair_id,
"position": chair_pos})
# Look at the center.
commands.extend(get_object_look_at(o_id=chair_id,
pos=TDWUtils.VECTOR3_ZERO,
noise=5))
# Set the plates on top of the table and moved in a bit.
plate_pos = get_move_along_direction(pos={"x": setting_pos["x"], "y": top["y"], "z": setting_pos["z"]},
target={"x": 0, "y": top["y"], "z": 0},
d=random.uniform(0.1, 0.125),
noise=0.01)
# Add a plate.
commands.extend(self.add_physics_object_default(position=plate_pos,
rotation=TDWUtils.VECTOR3_ZERO,
name=plate_name))
# Maybe add food on the plate.
if random.random() > 0.33:
# Use the plate bounds to add the food on top of the plate.
plate_bounds = MODEL_LIBRARIES["models_full.json"].get_record(plate_name).bounds
food_id = Controller.get_unique_id()
food_pos = {"x": plate_pos["x"] + random.uniform(-0.02, 0.02),
"y": top["y"] + plate_bounds["top"]["y"] + 0.001,
"z": plate_pos["z"] + random.uniform(-0.02, 0.02)}
commands.extend(self.add_physics_object_default(position=food_pos,
rotation={"x": 0,
"y": random.uniform(-89, 89),
"z": 0},
name=random.choice(self._FOOD),
o_id=food_id))
# Make the food small.
c_s = random.uniform(0.2, 0.45)
commands.append({"$type": "scale_object",
"id": food_id,
"scale_factor": {"x": c_s, "y": c_s, "z": c_s}})
for cutlery, offset in zip([fork_name, knife_name, spoon_name, cup_name],
[s_p.fork_offset, s_p.knife_offset, s_p.spoon_offset, s_p.cup_offset]):
# Maybe add cutlery at this position.
if random.random() > 0.25:
c_id = Controller.get_unique_id()
# Add the object. Slide it offset from the plate.
commands.extend(self.add_physics_object_default(position={"x": plate_pos["x"] + offset["x"],
"y": top["y"],
"z": plate_pos["z"] + offset["z"]},
rotation={"x": 0,
"y": s_p.cutlery_rotation,
"z": 0},
name=cutlery,
o_id=c_id))
# Add a centerpiece.
if random.random() > 0.25:
centerpiece_id = Controller.get_unique_id()
commands.extend(self.add_physics_object_default(position={"x": 0, "y": top["y"], "z": 0},
rotation={"x": 0,
"y": random.uniform(-89, 89),
"z": 0},
name=random.choice(self._CENTERPIECES),
o_id=centerpiece_id))
return commands
def get_trial_initialization_commands(self) -> List[dict]:
super().get_trial_initialization_commands()
trial_commands = []
if self.pool_id is not None:
trial_commands.append({"$type": "destroy_flex_container", "id": 0})
# Load a pool container for the fluid.
self.pool_id = Controller.get_unique_id()
self.non_flex_objects.append(self.pool_id)
trial_commands.append(
self.add_transforms_object(record=self.receptacle_record,
position={
"x": 0,
"y": 0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
o_id=self.pool_id))
trial_commands.extend([{
"$type": "scale_object",
"id": self.pool_id,
"scale_factor": {
"x": 1.5,
"y": 2.5,
"z": 1.5
}
}, {
"$type": "set_kinematic_state",
"id": self.pool_id,
"is_kinematic": True,
"use_gravity": False
}])
# Randomly select a fluid type
fluid_type_selection = choice(self.ft.fluid_type_names)
# Create the container, set up for fluids.
# Slow down physics so the water can settle without splashing out of the container.
trial_commands.extend([{
"$type":
"create_flex_container",
"collision_distance":
0.04,
"static_friction":
0.1,
"dynamic_friction":
0.1,
"particle_friction":
0.1,
"viscocity":
self.ft.fluid_types[fluid_type_selection].viscosity,
"adhesion":
self.ft.fluid_types[fluid_type_selection].adhesion,
"cohesion":
self.ft.fluid_types[fluid_type_selection].cohesion,
"radius":
0.1,
"fluid_rest":
0.05,
"damping":
0.01,
"substep_count":
5,
"iteration_count":
8,
"buoyancy":
1.0
}, {
"$type": "set_time_step",
"time_step": 0.005
}])
# Recreate fluid.
# Add the fluid actor, using the FluidPrimitive. Allow 500 frames for the fluid to settle before continuing.
fluid_id = Controller.get_unique_id()
trial_commands.extend(
self.add_fluid_object(position={
"x": 0,
"y": 1.0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
o_id=fluid_id,
fluid_type=fluid_type_selection))
trial_commands.append({"$type": "set_time_step", "time_step": 0.03})
# Select an object at random.
model = choice(self.model_list)
model_record = self.full_lib.get_record(model)
info = PHYSICS_INFO[model]
# Randomly select an object, and randomly orient it.
# Set the solid actor and assign the container.
o_id = Controller.get_unique_id()
trial_commands.extend(
self.add_solid_object(record=model_record,
position={
"x": 0,
"y": 2,
"z": 0
},
rotation={
"x": uniform(-45.0, 45.0),
"y": uniform(-45.0, 45.0),
"z": uniform(-45.0, 45.0)
},
o_id=o_id,
scale={
"x": 0.5,
"y": 0.5,
"z": 0.5
},
mass_scale=info.mass,
particle_spacing=0.05))
# Reset physics time-step to a more normal value.
# Position and aim avatar.
trial_commands.extend([{
"$type": "set_kinematic_state",
"id": o_id
}, {
"$type": "teleport_avatar_to",
"position": {
"x": -2.675,
"y": 1.375,
"z": 0
}
}, {
"$type": "look_at",
"object_id": self.pool_id,
"use_centroid": True
}])
return trial_commands