def save_images(images: Images,
filename: str,
output_directory="dist",
resize_to=None,
append_pass: bool = True) -> None:
"""
Save each image in the Images object.
The name of the image will be: pass_filename.extension, e.g.: `"0000"` -> `depth_0000.png`
The images object includes the pass and extension information.
:param images: The Images object. Contains each capture pass plus metadata.
:param output_directory: The directory to write images to.
:param filename: The filename of each image, minus the extension. The image pass will be appended as a prefix.
:param resize_to: Specify a (width, height) tuple to resize the images to. This is slower than saving as-is.
:param append_pass: If false, the image pass will _not_ be appended to the filename as a prefix, e.g.: `"0000"`: -> "`0000.jpg"`
"""
if not os.path.isdir(output_directory):
os.makedirs(output_directory)
for i in range(images.get_num_passes()):
if append_pass:
fi = images.get_pass_mask(
i)[1:] + "_" + filename + "." + images.get_extension(i)
else:
fi = filename + "." + images.get_extension(i)
if resize_to:
TDWUtils.get_pil_image(images, i).resize((resize_to[0], resize_to[1]), Image.LANCZOS)\
.save(os.path.join(output_directory, fi))
else:
with open(os.path.join(output_directory, fi), "wb") as f:
f.write(images.get_image(i))
def run(self):
self.start()
depth_pass = "******"
# Create an empty room.
# Set the screen size.
commands = [TDWUtils.create_empty_room(12, 12),
{"$type": "set_screen_size",
"width": 512,
"height": 512}]
# Add the avatar.
commands.extend(TDWUtils.create_avatar(position={"x": 1.57, "y": 3, "z": 3.56}, look_at=TDWUtils.VECTOR3_ZERO))
# Add an object.
# Request images and camera matrices.
commands.extend([self.get_add_object("trunck", object_id=0),
{"$type": "set_pass_masks",
"pass_masks": [depth_pass]},
{"$type": "send_images"},
{"$type": "send_camera_matrices"}])
resp = self.communicate(commands)
depth_image = None
camera_matrix = None
images = None
for i in range(len(resp) - 1):
r_id = OutputData.get_data_type_id(resp[i])
# Get the image.
if r_id == "imag":
images = Images(resp[i])
for j in range(images.get_num_passes()):
if images.get_pass_mask(j) == depth_pass:
depth_image = images.get_image(j)
# Get the camera matrix.
elif r_id == "cama":
camera_matrix = CameraMatrices(resp[i]).get_camera_matrix()
# Save the image.
TDWUtils.save_images(images=images, output_directory="D:/depth_shader", filename="0", append_pass=True)
# Get the depth values of each pixel.
depth = TDWUtils.get_depth_values(image=depth_image, width=images.get_width(), height=images.get_height())
print(np.min(depth), np.max(depth))
print(depth)
np.save("depth", depth)
np.save("camera_matrix", camera_matrix)
# Get a point cloud and write it to disk.
point_cloud_filename = "point_cloud.txt"
print(f"Point cloud saved to: {Path(point_cloud_filename)}")
TDWUtils.get_point_cloud(depth=depth, filename=point_cloud_filename, camera_matrix=camera_matrix)
# Show the depth values.
plt.imshow(depth)
plt.show()
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]:
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 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
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)
