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