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)