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 _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 _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