def _is_moving(o_id: int, transforms: Transforms,
rigidbodies: Rigidbodies) -> bool:
"""
:param o_id: The ID of the object.
:param transforms: The Transforms output data.
:param rigidbodies: The Rigidbodies output data.
:return: True if the object is still moving.
"""
y: Optional[float] = None
sleeping: bool = False
for i in range(transforms.get_num()):
if transforms.get_id(i) == o_id:
y = transforms.get_position(i)[1]
break
assert y is not None, f"y value is none for {o_id}"
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == o_id:
sleeping = rigidbodies.get_sleeping(i)
break
# If the object isn't sleeping, it is still moving.
# If the object fell into the abyss, we don't count it as moving (to prevent an infinitely long simulation).
return not sleeping and y > -10
def _set_tower_height_now(self, resp: List[bytes]) -> None:
top_obj_id = self.object_ids[-1]
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()):
if tr.get_id(i) == top_obj_id:
self.tower_height = tr.get_position(i)[1]
def get_object_position(self, obj_id: int, resp: List[bytes]) -> None:
position = None
for r in resp:
r_id = OutputData.get_data_type_id(r)
if r_id == "tran":
tr = Transforms(r)
for i in range(tr.get_num()):
if tr.get_id(i) == obj_id:
position = tr.get_position(i)
return position
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 is_done(self, resp: List[bytes], frame: int) -> bool:
# The ball must have a positive x coordinate (moving away from occluder) and be out of frame.
positive_x = False
segmentation_color = False
for r in resp:
r_id = OutputData.get_data_type_id(r)
if r_id == "tran":
t = Transforms(r)
for i in range(t.get_num()):
if t.get_id(i) == self._ball_id:
positive_x = t.get_position(i)[0] >= 2
elif r_id == "ipsc":
ip = IdPassSegmentationColors(r)
segmentation_color = ip.get_num_segmentation_colors() == 1
return positive_x and segmentation_color
def move_objects_along_forward(self, transforms: Transforms, delta, y=0):
commands = []
for i in range(transforms.get_num()):
object_id = transforms.get_id(i)
position = transforms.get_position(i)
forward = transforms.get_forward(i)
deltas = [d * delta for d in forward]
position = {
"x": position[0] + deltas[0],
"y": y,
"z": position[2] + deltas[2]
}
commands.append({
"$type": "teleport_object",
"id": object_id,
"position": position
})
self.communicate(commands)
def _get_frame_state(t: Transforms, r: Rigidbodies, num: int) -> dict:
"""
Returns a dictionary representing the current frame state.
:param t: The transforms data.
:param r: The rigidbodies data.
:param num: The current frame.
"""
objects = dict()
assert t.get_num() == r.get_num()
for i in range(t.get_num()):
objects.update({
t.get_id(i): {
"position": t.get_position(i),
"rotation": t.get_rotation(i),
"forward": t.get_forward(i),
"velocity": r.get_velocity(i),
"angular_velocity": r.get_angular_velocity(i),
"mass": r.get_mass(i)
}
})
return {"frame": num, "objects": objects}
def run(self):
rng = np.random.RandomState(0)
self.model_librarian = ModelLibrarian("models_special.json")
self.start()
commands = [TDWUtils.create_empty_room(100, 100)]
# The starting height of the objects.
y = 10
# The radius of the circle of objects.
r = 7.0
# The mass of each object.
mass = 5
# Get all points within the circle defined by the radius.
p0 = np.array((0, 0))
o_id = 0
for x in np.arange(-r, r, 1):
for z in np.arange(-r, r, 1):
p1 = np.array((x, z))
dist = np.linalg.norm(p0 - p1)
if dist < r:
# Add an object.
# Set its mass, physics properties, and color.
commands.extend([self.get_add_object("prim_cone",
object_id=o_id,
position={"x": x, "y": y, "z": z},
rotation={"x": 0, "y": 0, "z": 180}),
{"$type": "set_mass",
"id": o_id,
"mass": mass},
{"$type": "set_physic_material",
"dynamic_friction": 0.8,
"static_friction": 0.7,
"bounciness": 0.5,
"id": o_id},
{"$type": "set_color",
"color": {"r": rng.random_sample(),
"g": rng.random_sample(),
"b": rng.random_sample(),
"a": 1.0},
"id": o_id}])
o_id += 1
# Request transforms per frame.
commands.extend([{"$type": "send_transforms",
"frequency": "always"}])
# Create an avatar to observe the grisly spectacle.
avatar_position = {"x": -20, "y": 8, "z": 18}
commands.extend(TDWUtils.create_avatar(position=avatar_position, look_at=TDWUtils.VECTOR3_ZERO))
commands.append({"$type": "set_focus_distance",
"focus_distance": TDWUtils.get_distance(avatar_position, TDWUtils.VECTOR3_ZERO)})
resp = self.communicate(commands)
# If an objects are this far away from (0, 0, 0) the forcefield "activates".
forcefield_radius = 5
# The forcefield will bounce objects away at this force.
forcefield_force = -10
zeros = np.array((0, 0, 0))
for i in range(1000):
transforms = Transforms(resp[0])
commands = []
for j in range(transforms.get_num()):
pos = transforms.get_position(j)
pos = np.array(pos)
# If the object is in the forcefield, apply a force.
if TDWUtils.get_distance(TDWUtils.array_to_vector3(pos), TDWUtils.VECTOR3_ZERO) <= forcefield_radius:
# Get the normalized directional vector and multiply it by the force magnitude.
d = zeros - pos
d = d / np.linalg.norm(d)
d = d * forcefield_force
commands.append({"$type": "apply_force_to_object",
"id": transforms.get_id(j),
"force": TDWUtils.array_to_vector3(d)})
resp = self.communicate(commands)
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 get_video(self,
a_and_s: _AnimationAndScene,
humanoids: List[HumanoidRecord],
skyboxes: List[HDRISkyboxRecord],
pbar: Optional[tqdm] = None,
num_camera_iterations: int = 1,
h_id: int = 0,
min_duration: float = 2.5,
overwrite: bool = False) -> None:
"""
Create a single animation video as a series of images that will be written to the disk.
:param a_and_s: The animation/scene combo.
:param humanoids: The humanoid records. Iterate through all of these.
:param skyboxes: The HDRI skyboxes. Iterate through some of these (equal to num_camera_iterations).
:param pbar: The progress bar (optional).
:param num_camera_iterations: Make this many videos of each specific animation/humanoid/scene/skybox combo, changing the camera angle each time.
:param h_id: The object ID of the humanoid.
:param min_duration: The minimum duration of the video in seconds.
:param overwrite: If True, overwrite existing images.
"""
# Skip some frames to maintain 30 FPS.
save_per = a_and_s.animation.framerate / 30
# Randomize the skyboxes without randomizing the original list.
skyboxes_temp = skyboxes[:]
HumanoidVideo.RNG.shuffle(skyboxes_temp)
assert num_camera_iterations < len(skyboxes_temp), f"Number of camera iterations ({num_camera_iterations})" \
f" exceeds number of skyboxes ({len(skyboxes_temp)}"
# Determine how many times to loop the animation.
total_duration = 0
num_loops = 0
while total_duration < min_duration:
num_loops += 1
total_duration += a_and_s.animation.duration
for humanoid in humanoids:
for i in range(num_camera_iterations):
name = a_and_s.animation.name + "-" + a_and_s.scene.name + "-" + humanoid.name + "-" + str(
i)
# Prepare the output directory for the images.
output_dir = self.root_dir.joinpath(name)
metadata_path = output_dir.joinpath("metadata.json")
# Skip a completed video.
if metadata_path.exists() and not overwrite:
if pbar is not None:
pbar.update(1)
continue
if not output_dir.exists():
output_dir.mkdir(parents=True)
output_dir = str(output_dir.resolve())
# Add a skybox if this scene is HDRI-enabled and exterior.
add_skybox = a_and_s.scene.hdri and a_and_s.scene.location == "exterior"
skybox = skyboxes_temp.pop(0)
if pbar is not None:
pbar.set_description(name)
# Per-frame, focus on the object and look at the object.
# Pitch the camera slightly upward so that it isn't pointing at the stomach.
per_frame_commands = [{
"$type": "focus_on_object",
"object_id": h_id,
"use_centroid": True
}, {
"$type": "look_at",
"object_id": h_id,
"use_centroid": True
}, {
"$type": "rotate_sensor_container_by",
"axis": "pitch",
"angle": -5
}, {
"$type": "send_images",
"frequency": "once"
}]
# Write metadata to disk.
metadata = {
"humanoid": humanoid.name,
"animation": a_and_s.animation.name,
"scene": a_and_s.scene.name,
"location": a_and_s.scene.location,
"skybox": skybox.name if add_skybox else "",
"datetime": str(datetime.now()),
"tdw_version": __version__
}
# Load the scene, avatar, and humanoid.
# Play the animation and begin image capture.
commands = [{
"$type": "add_scene",
"name": a_and_s.scene.name,
"url": a_and_s.scene.get_url()
}, {
"$type": "add_humanoid_animation",
"name": a_and_s.animation.name,
"url": a_and_s.animation.get_url()
}, {
"$type": "create_avatar",
"type": "A_Img_Caps_Kinematic",
"id": "a"
}, {
"$type": "add_humanoid",
"name": humanoid.name,
"url": humanoid.get_url(),
"position": a_and_s.position,
"rotation": a_and_s.rotation,
"id": h_id
}, {
"$type": "set_pass_masks",
"pass_masks": ["_img"]
}]
if add_skybox:
commands.extend([
self.get_add_hdri_skybox(skybox.name), {
"$type": "set_post_exposure",
"post_exposure": 0.4
}, {
"$type": "set_aperture",
"aperture": 1.6
}, {
"$type": "set_contrast",
"contrast": -20
}
])
else:
# Set default post-processing values.
commands.extend([{
"$type": "set_post_exposure"
}, {
"$type": "set_aperture"
}, {
"$type": "set_contrast"
}])
commands.extend(per_frame_commands)
self.communicate(commands)
frame = 0
for loop in range(num_loops):
# Start to play the animation.
# Try to get all of the humanoid in the viewport.
init_loop_commands = per_frame_commands[:-1]
init_loop_commands.extend([{
"$type": "play_humanoid_animation",
"name": a_and_s.animation.name,
"id": h_id
}])
self.communicate(init_loop_commands)
if loop == 0:
resp = self.communicate({
"$type": "send_humanoids",
"frequency": "once",
"ids": [h_id]
})
tr = Transforms(resp[0])
o_pos = tr.get_position(0)
# Camera distance from the humanoid.
avatar_d = HumanoidVideo.RNG.uniform(-2.2, -4.5)
# A reasonable rotation.
avatar_theta = np.radians(
HumanoidVideo.RNG.uniform(-85, 85))
# Get the avatar position from the angle and distance.
avatar_x = o_pos[0] + (avatar_d * np.cos(avatar_theta))
avatar_z = o_pos[2] + (avatar_d * np.sin(avatar_theta))
avatar_y = HumanoidVideo.RNG.uniform(
1.8, 2 + (avatar_d * 0.12))
per_frame_commands.append({
"$type": "teleport_avatar_to",
"position": {
"x": avatar_x,
"y": avatar_y,
"z": avatar_z
}
})
# Play the animation.
count = 0
while count < a_and_s.animation.get_num_frames():
# Drop frames to stay at 30 FPS.
if count % save_per == 0:
resp = self.communicate(per_frame_commands)
TDWUtils.save_images(Images(resp[0]),
TDWUtils.zero_padding(frame),
output_dir,
append_pass=False)
frame += 1
else:
self.communicate([])
count += 1
# Destroy the humanoid.
self.communicate({"$type": "destroy_humanoid", "id": h_id})
# Write the metadata file (signalling that the video was completed).
metadata_path.write_text(json.dumps(metadata, indent=4),
encoding="utf-8")
if pbar is not None:
pbar.update(1)
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library="models_full.json")
c.communicate({"$type": "send_transforms", "frequency": "always"})
char = getch()
if char == b'w':
resp = c.communicate({
"$type": "teleport_object",
"position": {
"x": 1,
"y": 0,
"z": 0
},
"id": o_id
})
transform = Transforms(resp[0])
# Get the position and rotation of the iron_box object.
position = transform.get_position(0)
rotation = transform.get_rotation(0)
print("Position: " + position)
print("Rotation: " + rotation)
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 _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