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