def get_collisions( resp: List[bytes] ) -> Tuple[List[Collision], List[EnvironmentCollision], Optional[Rigidbodies]]: """ Parse collision and rigibody data from the output data. :param resp: The response from the build. :return: A list of collisions on this frame (can be empty), a list of environment collisions on this frame (can be empty), and Rigidbodies data (can be `None`). """ if len(resp) == 1: return [], [], None collisions: List[Collision] = [] environment_collisions: List[EnvironmentCollision] = [] rigidbodies: Optional[Rigidbodies] = None for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == 'coll': collisions.append(Collision(r)) if r_id == 'rigi': rigidbodies = Rigidbodies(r) if r_id == 'enco': environment_collisions.append(EnvironmentCollision(r)) return collisions, environment_collisions, rigidbodies
def communicate(self, commands: Union[dict, List[dict]]) -> List[bytes]: if isinstance(commands, dict): commands = [commands] # Add commands from key presses. commands.extend(self.on_key_commands[:]) # Clear the on-key commands. self.on_key_commands.clear() # Send the commands. resp = super().communicate(commands) # Get keyboard input. for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "keyb": keys = Keyboard(resp[i]) # Listen for events where the key was first pressed on the previous frame. for j in range(keys.get_num_pressed()): self._do_event(k=keys.get_pressed(j), events=self._press) # Listen for keys currently held down. for j in range(keys.get_num_held()): self._do_event(k=keys.get_held(j), events=self._hold) # Listen for keys that were released. for j in range(keys.get_num_released()): self._do_event(k=keys.get_released(j), events=self._release) return resp
def communicate(self, commands: Union[dict, List[dict]]) -> list: """ Send commands and receive output data in response. :param commands: A list of JSON commands. :return The output data from the build. """ if isinstance(commands, list): msg = [json.dumps(commands).encode('utf-8')] else: msg = [json.dumps([commands]).encode('utf-8')] # Send the commands. self.socket.send_multipart(msg) # Receive output data. resp = self.socket.recv_multipart() # Occasionally, the build's socket will stop receiving messages. # If that happens, it will close the socket, create a new socket, and send a dummy output data object. # The ID of the dummy object is "ftre" (FailedToReceive). # If the controller receives the dummy object, it should re-send its commands. # The dummy object is always in an array: [ftre, 0] # This way, the controller can easily differentiate it from a response that just has the frame count. while len(resp) > 1 and OutputData.get_data_type_id(resp[0]) == "ftre": self.socket.send_multipart(msg) resp = self.socket.recv_multipart() # Return the output data from the build. return resp
def run(self): self.start() # Create an office. self.communicate(create_office(12, 24, 4, 4, 6)) # Get a random position on the NavMesh. x, y, z = TDWUtils.get_random_position_on_nav_mesh(self, 12, 24) # Index of the next destination. d = 0 # Create the avatar. self.communicate( TDWUtils.create_avatar(avatar_type="A_Nav_Mesh", position={ "x": x, "y": 0.5, "z": 0 })) # Teleport the avatar. # Set the speed of the avatar. # Set the destination of the avatar. self.communicate([{ "$type": "set_nav_mesh_avatar", "avatar_id": "a", "speed": 2 }, { "$type": "set_nav_mesh_avatar_destination", "avatar_id": "a", "destination": NavMesh.DESTINATIONS[d] }]) # Go to 3 destinations. i = 0 while i < 3: resp = self.communicate({"$type": "do_nothing"}) # If there is only 1 element in the response, it is an empty frame. if len(resp) == 1: continue # The avatar arrived at the destination! else: assert OutputData.get_data_type_id(resp[0]) == "anmd" # Parse the output data. data = ArrivedAtNavMeshDestination(resp[0]) print(data.get_avatar_id()) print("") i += 1 d += 1 if d >= len(NavMesh.DESTINATIONS): d = 0 # Set a new destination. self.communicate({ "$type": "set_nav_mesh_avatar_destination", "avatar_id": "a", "destination": NavMesh.DESTINATIONS[d] })
def run(self): # Load the streamed scene and add controller rig. self.load_streamed_scene(scene="tdw_room_2018") self.communicate({"$type": "create_vr_rig"}) # Add the table object and make it kinematic. table_id = self.add_object("small_table_green_marble", position={ "x": 0, "y": 0, "z": 0.5 }) self.communicate([{ "$type": "set_kinematic_state", "id": table_id, "is_kinematic": True }]) # Add a box object and make it graspable. graspable_id_box = self.add_object("woven_box", position={ "x": 0.2, "y": 1.0, "z": 0.5 }, library="models_core.json") self.communicate([{"$type": "set_graspable", "id": graspable_id_box}]) # Add the ball object and make it graspable. graspable_id = self.add_object("prim_sphere", position={ "x": 0.2, "y": 3.0, "z": 0.5 }, library="models_special.json") self.communicate([{ "$type": "set_graspable", "id": graspable_id }, { "$type": "scale_object", "scale_factor": { "x": 0.2, "y": 0.2, "z": 0.2 }, "id": graspable_id }]) # Receive VR data. resp = self.communicate({"$type": "send_vr_rig", "frequency": "once"}) assert len(resp) > 1 and OutputData.get_data_type_id(resp[0]) == "vrri" vr = VRRig(resp[0]) print(vr.get_head_rotation()) # Start an infinite loop to allow the build to simulate physics. while True: self.communicate({"$type": "do_nothing"})
def run(self): self.start() commands = [TDWUtils.create_empty_room(self.size, self.size)] commands.extend( TDWUtils.create_avatar(position={ "x": 0, "y": 1.5, "z": 0 }, avatar_id="a")) commands.extend([{ "$type": "set_pass_masks", "avatar_id": "a", "pass_masks": ["_img"] }, { "$type": "send_images", "frequency": "always" }]) self.communicate(commands) # Listen for keys. # Turn left. self.listen(key="left", commands={ "$type": "rotate_sensor_container_by", "axis": "yaw", "angle": self.angle * -1, "sensor_name": "SensorContainer", "avatar_id": "a" }) # Turn right. self.listen(key="right", commands={ "$type": "rotate_sensor_container_by", "axis": "yaw", "angle": self.angle, "sensor_name": "SensorContainer", "avatar_id": "a" }) # Quit the application. self.listen(key="esc", function=self.stop) # Continue until the quit key is pressed. i = 0 while not self.done: # Listen for keyboard input. Receive output data. resp = self.communicate([]) for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) # Save images. if r_id == "imag": TDWUtils.save_images( images=Images(r), output_directory=self.images_directory, filename=TDWUtils.zero_padding(i, width=4)) # Increment the image number. i += 1 self.communicate({"$type": "terminate"})
def _set_tower_height_now(self, resp: List[bytes]) -> None: top_obj_id = self.object_ids[-1] 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()): if tr.get_id(i) == top_obj_id: self.tower_height = tr.get_position(i)[1]
def _write_frame(self, frames_grp: h5py.Group, resp: List[bytes], frame_num: int) -> \ Tuple[h5py.Group, h5py.Group, dict, bool]: frame, objs, tr, done = super()._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame_num) num_objects = len(self.object_ids) # Physics data. velocities = np.empty(dtype=np.float32, shape=(num_objects, 3)) angular_velocities = np.empty(dtype=np.float32, shape=(num_objects, 3)) # Collision data. collision_ids = np.empty(dtype=np.int32, shape=(0, 2)) collision_relative_velocities = np.empty(dtype=np.float32, shape=(0, 3)) collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3)) # Environment Collision data. env_collision_ids = np.empty(dtype=np.int32, shape=(0, 1)) env_collision_contacts = np.empty(dtype=np.float32, shape=(0, 2, 3)) sleeping = True for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == "rigi": ri = Rigidbodies(r) ri_dict = dict() for i in range(ri.get_num()): ri_dict.update({ri.get_id(i): {"vel": ri.get_velocity(i), "ang": ri.get_angular_velocity(i)}}) # Check if any objects are sleeping that aren't in the abyss. if not ri.get_sleeping(i) and tr[ri.get_id(i)]["pos"][1] >= -1: sleeping = False # Add the Rigibodies data. for o_id, i in zip(self.object_ids, range(num_objects)): velocities[i] = ri_dict[o_id]["vel"] angular_velocities[i] = ri_dict[o_id]["ang"] elif r_id == "coll": co = Collision(r) collision_ids = np.append(collision_ids, [co.get_collider_id(), co.get_collidee_id()]) collision_relative_velocities = np.append(collision_relative_velocities, co.get_relative_velocity()) for i in range(co.get_num_contacts()): collision_contacts = np.append(collision_contacts, (co.get_contact_normal(i), co.get_contact_point(i))) elif r_id == "enco": en = EnvironmentCollision(r) env_collision_ids = np.append(env_collision_ids, en.get_object_id()) for i in range(en.get_num_contacts()): env_collision_contacts = np.append(env_collision_contacts, (en.get_contact_normal(i), en.get_contact_point(i))) objs.create_dataset("velocities", data=velocities.reshape(num_objects, 3), compression="gzip") objs.create_dataset("angular_velocities", data=angular_velocities.reshape(num_objects, 3), compression="gzip") collisions = frame.create_group("collisions") collisions.create_dataset("object_ids", data=collision_ids.reshape((-1, 2)), compression="gzip") collisions.create_dataset("relative_velocities", data=collision_relative_velocities.reshape((-1, 3)), compression="gzip") collisions.create_dataset("contacts", data=collision_contacts.reshape((-1, 2, 3)), compression="gzip") env_collisions = frame.create_group("env_collisions") env_collisions.create_dataset("object_ids", data=env_collision_ids, compression="gzip") env_collisions.create_dataset("contacts", data=env_collision_contacts.reshape((-1, 2, 3)), compression="gzip") return frame, objs, tr, sleeping
def is_done(self, resp: List[bytes], frame: int) -> bool: for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) # If the ball reaches or overshoots the destination, the trial is done. if r_id == "tran": t = Transforms(r) d0 = TDWUtils.get_distance( TDWUtils.array_to_vector3(t.get_position(0)), self._p0) d1 = TDWUtils.get_distance(self._p0, self._p1) return d0 > d1 * 1.5 return False
def _get_transforms(resp: List[bytes]) -> Transforms: """ :param resp: The output data response. :return: Transforms data. """ for r in resp[:-1]: if OutputData.get_data_type_id(r) == "tran": return Transforms(r) raise Exception("Transforms output data not found!")
def get_object_position(self, obj_id: int, resp: List[bytes]) -> None: position = None for r in resp: r_id = OutputData.get_data_type_id(r) if r_id == "tran": tr = Transforms(r) for i in range(tr.get_num()): if tr.get_id(i) == obj_id: position = tr.get_position(i) return position
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 _get_frame(self, resp: List[bytes]) -> AvatarStickyMitten: """ :param resp: The response from the build. :return: AvatarStickyMitten output data for this avatar on this frame. """ for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "avsm": avsm = AvatarStickyMitten(resp[i]) if avsm.get_avatar_id() == self.id: return avsm raise Exception(f"No avatar data found for {self.id}")
def get_object_environment_collision(self, obj_id: int, resp: List[bytes]): contact_points = [] contact_normals = [] for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == 'enco': en = EnvironmentCollision(r) if en.get_object_id() == obj_id: contact_points = [np.array(en.get_contact_point(i)) for i in range(en.get_num_contacts())] contact_normals = [np.array(en.get_contact_normal(i)) for i in range(en.get_num_contacts())] return (contact_points, contact_normals)
def get_bounds_data(self, object_id): resp = self.communicate({ "$type": "send_bounds", "frequency": "once", "ids": [object_id] }) for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == "imag": print("This is an image") elif r_id == "boun": return Bounds(r) raise Exception("No bounds data found!")
def run(self): # Create the scene. # Set image encoding globals. self.start() commands = [TDWUtils.create_empty_room(12, 12), {"$type": "set_screen_size", "width": 128, "height": 128}, {"$type": "set_img_pass_encoding", "value": False}] # Create the avatar. commands.extend(TDWUtils.create_avatar(position={"x": 2.478, "y": 1.602, "z": 1.412}, look_at=TDWUtils.VECTOR3_ZERO, avatar_id="a")) # Enable all pass masks. Request an image for this frame only. commands.extend([{"$type": "set_pass_masks", "pass_masks": ["_img", "_id", "_category", "_mask", "_depth", "_normals", "_flow", "_depth_simple", "_albedo"], "avatar_id": "a"}, {"$type": "send_images", "ids": ["a"], "frequency": "once"}]) # Add objects. commands.append(self.get_add_object("small_table_green_marble", position=TDWUtils.VECTOR3_ZERO, rotation=TDWUtils.VECTOR3_ZERO, object_id=0)) commands.append(self.get_add_object("rh10", position={"x": 0.7, "y": 0, "z": 0.4}, rotation={"x": 0, "y": 30, "z": 0}, object_id=1)) commands.append(self.get_add_object("jug01", position={"x": -0.3, "y": 0.9, "z": 0.2}, rotation=TDWUtils.VECTOR3_ZERO, object_id=3)) commands.append(self.get_add_object("jug05", position={"x": 0.3, "y": 0.9, "z": -0.2}, rotation=TDWUtils.VECTOR3_ZERO, object_id=4)) # Send the commands. resp = self.communicate(commands) # Save the images. for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == "imag": # Save the images. TDWUtils.save_images(output_directory="dist", images=Images(r), filename="0") print(f"Images saved to: {Path('dist').resolve()}") # Stop the build. self.communicate({"$type": "terminate"})
def get_object_target_collision(self, obj_id: int, target_id: int, resp: List[bytes]): contact_points = [] contact_normals = [] for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == "coll": co = Collision(r) coll_ids = [co.get_collider_id(), co.get_collidee_id()] if [obj_id, target_id] == coll_ids or [target_id, obj_id] == coll_ids: contact_points = [co.get_contact_point(i) for i in range(co.get_num_contacts())] contact_normals = [co.get_contact_normal(i) for i in range(co.get_num_contacts())] return (contact_points, contact_normals)
def is_done(self, resp: List[bytes], frame: int) -> bool: # The ball must have a positive x coordinate (moving away from occluder) and be out of frame. positive_x = False segmentation_color = False for r in resp: r_id = OutputData.get_data_type_id(r) if r_id == "tran": t = Transforms(r) for i in range(t.get_num()): if t.get_id(i) == self._ball_id: positive_x = t.get_position(i)[0] >= 2 elif r_id == "ipsc": ip = IdPassSegmentationColors(r) segmentation_color = ip.get_num_segmentation_colors() == 1 return positive_x and segmentation_color
def get_robot(self, resp: List[bytes]) -> Robot: """ :param resp: The response from the build. :return: `Robot` output data. """ robot: Optional[Robot] = None for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "robo": r = Robot(resp[i]) if r.get_id() == self.robot_id: robot = r break return robot
def _set_segmentation_colors(self, resp: List[bytes]) -> None: self.object_segmentation_colors = None for r in resp: if OutputData.get_data_type_id(r) == 'segm': seg = SegmentationColors(r) colors = {} for i in range(seg.get_num()): colors[seg.get_object_id(i)] = seg.get_object_color(i) self.object_segmentation_colors = [] for o_id in self.object_ids: if o_id in colors.keys(): self.object_segmentation_colors.append( np.array(colors[o_id], dtype=np.uint8).reshape(1,3)) self.object_segmentation_colors = np.concatenate(self.object_segmentation_colors, 0)
def get_collisions( resp: List[bytes] ) -> Tuple[List[Collision], List[EnvironmentCollision]]: """ :param resp: The response from the build (a byte array). :return: Tuple: A list of collisions; a list of environment collisions. """ collisions: List[Collision] = list() env_collisions: List[EnvironmentCollision] = list() for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "coll": collisions.append(Collision(resp[i])) elif r_id == "enco": env_collisions.append(EnvironmentCollision(resp[i])) return collisions, env_collisions
def get_data(resp: List[bytes], d_type: Type[T]) -> Optional[T]: """ Parse the output data list of byte arrays to get a single type output data object. :param resp: The response from the build (a list of byte arrays). :param d_type: The desired type of output data. :return: An object of type `d_type` from `resp`. If there is no object, returns None. """ if d_type not in _OUTPUT_IDS: raise Exception(f"Output data ID not defined: {d_type}") for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == _OUTPUT_IDS[d_type]: return d_type(resp[i]) return None
def _do_frame(self, commands: Union[List[dict], dict]) -> AvatarStickyMitten: """ Send commands to the build. Receive images and sticky mitten avatar data. Save the images. :param commands: The commands for this frame. :return: The avatar data. """ # Add to the list of commands: Look at the avatar. frame_commands = [{ "$type": "look_at_position", "position": { "x": self.avatar_position[0], "y": self.avatar_position[1] + 0.5, "z": self.avatar_position[2] }, "avatar_id": self.cam_id }] if isinstance(commands, dict): frame_commands.append(commands) else: frame_commands.extend(commands) resp = self.communicate(frame_commands) avsm: Optional[AvatarStickyMitten] = None for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) # Save images. if r_id == "imag": TDWUtils.save_images( images=Images(r), filename=TDWUtils.zero_padding(self.frame_number, 4), output_directory=PutObjectOnTable.OUTPUT_DIR) elif r_id == "avsm": avsm = AvatarStickyMitten(r) # Update the position of the avatar and the frame number. self.avatar_position = avsm.get_position() self.frame_number += 1 return avsm
def communicate(self, commands: Union[dict, List[dict]]) -> List[bytes]: """ See `Magnebot.communicate()`. Images are saved per-frame. """ resp = super().communicate(commands=commands) if not self._debug: # Save all images. got_images = False for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "imag": got_images = True images = Images(resp[i]) avatar_id = images.get_avatar_id() if avatar_id in self.image_directories: TDWUtils.save_images(filename=TDWUtils.zero_padding(self._image_count, 8), output_directory=self.image_directories[avatar_id], images=images) if got_images: self._image_count += 1 return resp
def on_frame(self, resp: List[bytes]) -> List[dict]: """ Update the avatar based on its current arm-bending goals and its state. If the avatar has achieved a goal (for example, picking up an object), it will stop moving that arm. Update the avatar's state as needed. :param resp: The response from the build. :return: A list of commands to pick up, stop moving, etc. """ def _get_mitten_position(a: Arm) -> np.array: """ :param a: The arm. :return: The position of a mitten. """ if a == Arm.left: return np.array(frame.get_mitten_center_left_position()) else: return np.array(frame.get_mitten_center_right_position()) # Update dynamic data. frame = self._get_frame(resp=resp) # Update dynamic collision data. self.collisions.clear() self.env_collisions.clear() # Get each collision. for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "coll": coll = Collision(resp[i]) collider_id = coll.get_collider_id() collidee_id = coll.get_collidee_id() # Check if this was a mitten, if we're supposed to stop if there's a collision, # and if the collision was not with the target. for arm in self._ik_goals: if self._ik_goals[arm] is not None: if (collider_id == self.mitten_ids[arm] or collidee_id == self.mitten_ids[arm]) and \ (collider_id not in frame.get_held_right() and collider_id not in frame.get_held_left() and collidee_id not in frame.get_held_right() and collidee_id not in frame.get_held_left()) and \ self._ik_goals[arm].stop_on_mitten_collision and \ (self._ik_goals[arm].target is None or (self._ik_goals[arm].pick_up_id != collidee_id and self._ik_goals[arm].pick_up_id != collider_id)) and \ (collidee_id not in self.body_parts_static or collider_id not in self.body_parts_static): self.status = TaskStatus.mitten_collision self._ik_goals[arm] = None if self._debug: print( "Stopping because the mitten collided with something." ) return self._stop_arm(arm=arm) # Check if the collision includes a body part. if collider_id in self.body_parts_static and collidee_id not in self.body_parts_static: if collider_id not in self.collisions: self.collisions[collider_id] = [] self.collisions[collider_id].append(collidee_id) elif collidee_id in self.body_parts_static and collider_id not in self.body_parts_static: if collidee_id not in self.collisions: self.collisions[collidee_id] = [] self.collisions[collidee_id].append(collider_id) elif r_id == "enco": coll = EnvironmentCollision(resp[i]) collider_id = coll.get_object_id() if collider_id in self.body_parts_static: self.env_collisions.append(collider_id) # Check if IK goals are done. temp_goals: Dict[Arm, Optional[_IKGoal]] = dict() # Get commands for the next frame. commands: List[dict] = [] for arm in self._ik_goals: # No IK goal on this arm. if self._ik_goals[arm] is None: temp_goals[arm] = None # This is a dummy IK goal. Let it run. elif self._ik_goals[arm].target is None: temp_goals[arm] = self._ik_goals[arm] else: # Is the arm at the target? mitten_position = _get_mitten_position(arm) # If we're not trying to pick something up, check if we are at the target position. if self._ik_goals[arm].pick_up_id is None: # If we're at the position, stop. d = np.linalg.norm(mitten_position - self._ik_goals[arm].target) if d < self._ik_goals[arm].precision: if self._debug: print( f"{arm.name} mitten is at target position {self._ik_goals[arm].target}. Stopping." ) commands.extend(self._stop_arm(arm=arm)) temp_goals[arm] = None self.status = TaskStatus.success # Keep bending the arm. else: temp_goals[arm] = self._ik_goals[arm] self._ik_goals[arm].previous_distance = d # If we're trying to pick something, check if it was picked up on the previous frame. else: if self._ik_goals[arm].pick_up_id in frame.get_held_left() or self._ik_goals[arm]. \ pick_up_id in frame.get_held_right(): if self._debug: print( f"{arm.name} mitten picked up {self._ik_goals[arm].pick_up_id}. Stopping." ) commands.extend(self._stop_arm(arm=arm)) temp_goals[arm] = None self.status = TaskStatus.success # Keep bending the arm and trying to pick up the object. else: commands.extend([{ "$type": "pick_up_proximity", "distance": 0.02, "radius": 0.05, "grip": 1000, "is_left": arm == Arm.left, "avatar_id": self.id, "object_ids": [self._ik_goals[arm].pick_up_id] }, { "$type": "pick_up", "grip": 1000, "is_left": arm == Arm.left, "object_ids": [self._ik_goals[arm].pick_up_id], "avatar_id": self.id }]) temp_goals[arm] = self._ik_goals[arm] self._ik_goals = temp_goals # Check if the arms are still moving. temp_goals: Dict[Arm, Optional[_IKGoal]] = dict() for arm in self._ik_goals: # No IK goal on this arm. if self._ik_goals[arm] is None: temp_goals[arm] = None else: # Get the past and present angles. if arm == Arm.left: angles_0 = self.frame.get_angles_left() angles_1 = frame.get_angles_left() else: angles_0 = self.frame.get_angles_right() angles_1 = frame.get_angles_right() # Try to stop any moving joints. if self._ik_goals[arm].rotations is not None and self._ik_goals[ arm].pick_up_id is not None: joint_profile = self._get_default_sticky_mitten_profile() for angle, joint_name in zip(angles_1, Avatar.ANGLE_ORDER): target_angle = self._ik_goals[arm].rotations[ joint_name] # Check if the joint stopped moving. Ignore if the joint already stopped. if target_angle > 0.01 and np.abs(angle - target_angle) < 0.01 and \ joint_name in self._ik_goals[arm].moving_joints: self._ik_goals[arm].moving_joints.remove( joint_name) j = joint_name.split("_") j_name = f"{j[0]}_{arm.name}" axis = j[1] # Set the name of the elbow to the expected profile key. if "elbow" in joint_name: profile_key = "elbow" else: profile_key = joint_name if self._debug: print( f"{joint_name} {arm.name} slowing down: {np.abs(angle - target_angle)}" ) # Stop the joint from moving any more. # Set the damper, force, and angular drag to "default" (non-moving) values. commands.extend([{ "$type": "set_joint_damper", "joint": j_name, "axis": axis, "damper": joint_profile[profile_key]["damper"], "avatar_id": self.id }, { "$type": "set_joint_force", "joint": j_name, "axis": axis, "force": joint_profile[profile_key]["force"], "avatar_id": self.id }, { "$type": "set_joint_angular_drag", "joint": j_name, "axis": axis, "angular_drag": joint_profile[profile_key]["angular_drag"], "avatar_id": self.id }]) # Is any joint still moving? moving = False for a0, a1 in zip(angles_0, angles_1): if np.abs(a0 - a1) > 0.03: moving = True break # Keep moving. if moving: temp_goals[arm] = self._ik_goals[arm] else: if self._ik_goals[arm].rotations is not None: # This is a reset arm action. if self._ik_goals[arm].target is None: mitten_position = _get_mitten_position( arm) - frame.get_position() d = np.linalg.norm( self._initial_mitten_positions[arm] - mitten_position) # The reset arm action ended with the mitten very close to the initial position. if d < self._ik_goals[arm].precision: self.status = TaskStatus.success else: self.status = TaskStatus.no_longer_bending # This is a regular action. # It ended with the arm no longer moving but having never reached the target. else: if self._debug: print( f"{arm.name} is no longer bending. Cancelling." ) self.status = TaskStatus.no_longer_bending commands.extend(self._stop_arm(arm=arm)) temp_goals[arm] = None self._ik_goals = temp_goals self.frame = frame return commands
def __init__(self, resp: List[bytes], avatar_id: str = "a", debug: bool = False): """ :param resp: The response from the build after creating the avatar. :param avatar_id: The ID of the avatar. :param debug: If True, print debug statements. """ self.id = avatar_id self._debug = debug # Set the arm chains. self._arms: Dict[Arm, Chain] = { Arm.left: self._get_left_arm(), Arm.right: self._get_right_arm() } self._initial_mitten_positions = self._get_initial_mitten_positions() # Any current IK goals. self._ik_goals: Dict[Arm, Optional[_IKGoal]] = { Arm.left: None, Arm.right: None } smsc: Optional[AvatarStickyMittenSegmentationColors] = None for i in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[i]) if r_id == "smsc": q = AvatarStickyMittenSegmentationColors(resp[i]) if q.get_id() == avatar_id: smsc = q break assert smsc is not None, f"No avatar segmentation colors found for {avatar_id}" # Get data for the current frame. self.frame = self._get_frame(resp) # Get the masses of each body part. body_part_masses: Dict[int, float] = dict() for i in range(self.frame.get_num_rigidbody_parts()): body_part_masses[self.frame.get_rigidbody_part_id( i)] = self.frame.get_rigidbody_part_mass(i) # Cache static data of body parts. self.body_parts_static: Dict[int, BodyPartStatic] = dict() self.base_id = 0 self.mitten_ids: Dict[Arm, int] = dict() for i in range(smsc.get_num_body_parts()): body_part_id = smsc.get_body_part_id(i) if body_part_id in body_part_masses: mass = body_part_masses[body_part_id] else: mass = 0.1 name = smsc.get_body_part_name(i) # Cache the base ID and the mitten IDs. if name.startswith("A_StickyMitten"): self.base_id = body_part_id elif name == "mitten_left": self.mitten_ids[Arm.left] = body_part_id elif name == "mitten_right": self.mitten_ids[Arm.right] = body_part_id bps = BodyPartStatic( object_id=body_part_id, segmentation_color=smsc.get_body_part_segmentation_color(i), name=name, mass=mass) self.body_parts_static[body_part_id] = bps # Start dynamic data. self.collisions: Dict[int, List[int]] = dict() self.env_collisions: List[int] = list() self.status = TaskStatus.idle self._mass = self._get_mass()
def run(self): self.start() object_id = self.get_unique_id() resp = self.communicate([ TDWUtils.create_empty_room(12, 12), { "$type": "create_flex_container", "particle_size": 0.1, "collision_distance": 0.025, "solid_rest": 0.1 }, # {"$type": "create_flex_container", # "collision_distance": 0.001, # "static_friction": 1.0, # "dynamic_friction": 1.0, # "radius": 0.1875, # "max_particles": 200000}, self.get_add_object(model_name="linbrazil_diz_armchair", position={ "x": 0.0, "y": 2.0, "z": 0.0 }, rotation={ "x": 25.0, "y": 45.0, "z": -40.0 }, object_id=object_id), { "$type": "set_kinematic_state", "id": object_id }, { "$type": "set_flex_solid_actor", "id": object_id, "mesh_expansion": 0, "particle_spacing": 0.025, "mass_scale": 1 }, # {"$type": "set_flex_soft_actor", # "id": object_id, # "skinning_falloff": 0.5, # "volume_sampling": 1.0, # "mass_scale": 1.0, # "cluster_stiffness": 0.2, # "cluster_spacing": 0.2, # "cluster_radius": 0.2, # "link_radius": 0, # "link_stiffness": 1.0, # "particle_spacing": 0.025}, { "$type": "assign_flex_container", "id": object_id, "container_id": 0 }, { '$type': 'load_primitive_from_resources', 'id': 2, 'primitive_type': 'Cube', 'position': { 'x': 0, 'y': 0.5, 'z': 0 }, 'orientation': { 'x': 0, 'y': 0, 'z': 0 } }, { "$type": "scale_object", "id": 2, "scale_factor": { "x": 0.5, "y": 0.5, "z": 0.5 } }, { '$type': 'set_kinematic_state', 'id': 2, 'is_kinematic': False }, { "$type": "set_flex_solid_actor", "id": 2, "mesh_expansion": 0, "particle_spacing": 0.025, "mass_scale": 1 }, # {'$type': 'set_flex_soft_actor', # 'id': 2, # 'draw_particles': False, # 'particle_spacing': 0.125, # 'cluster_stiffness': 0.22055267521432875}, { '$type': 'assign_flex_container', 'id': 2, 'container_id': 0 }, { '$type': 'set_flex_particles_mass', 'id': 2, 'mass': 15.625 }, { "$type": "send_flex_particles", "frequency": "always" }, { "$type": "send_collisions", "enter": True, "stay": True, "exit": True, "collision_types": ["obj", "env"] } ]) for i in range(100): for j in range(len(resp) - 1): r_id = OutputData.get_data_type_id(resp[j]) if r_id == "flex": flex = FlexParticles(resp[j]) print(i, "flex", flex.get_id(0)) elif r_id == "enco": enco = EnvironmentCollision(resp[j]) print(i, "enco", enco.get_state()) if r_id == "coll": coll = Collision(resp[j]) print(i, "coll", coll.get_state()) if coll.get_state() == 'enter': print("BAM!") resp = self.communicate([])
"z": 0 }, rotation={ "x": 0, "y": 0, "z": 0 }), { "$type": "send_bounds", "ids": [table_id], "frequency": "once" }]) # Get the top of the table. top_y = 0 for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) # Find the bounds data. if r_id == "boun": b = Bounds(r) # We only requested the table, so it is object 0: _, top_y, _ = b.get_top(0) box_record = lib.get_record("puzzle_box_composite") box_id = 1 c.communicate( c.get_add_object(model_name=box_record.name, object_id=box_id, position={ "x": 0, "y": top_y, "z": 0
def run(self): """ Generate room using COCO_TDW dataset """ objects_in_scene = 15 object_ids = [] # Get Category-Object mapping TDW_COCO_models = TDW_relationships.get_COCO_TDW_mapping() # print("TDWCOCO:", TDW_COCO_models) # print("KEYS:", TDW_COCO_models.keys()) # Gets COCO categories co-occurring in a scene # +5 is for dealing with failed object insertion attempts COCO_configurations = msCOCO_matrix.get_max_co_occurrence(5, int(objects_in_scene + 5)) configuration_1 = COCO_configurations[0] print("Config 1:", configuration_1) # TDW models/objects objects = [] for COCO_object in configuration_1: print(COCO_object) print(COCO_object.split()) if len(COCO_object.split()) > 1: COCO_object = COCO_object.split()[-1] print(COCO_object) # Check if COCO category is a key in the COCO-TDW mapping if COCO_object in TDW_COCO_models.keys(): # Gets random TDW model (from COCO-to-TDW map) based on COCO category key print(TDW_COCO_models[COCO_object]) model = TDW_COCO_models[COCO_object][random.randint(0, len(TDW_COCO_models[COCO_object]) - 1)] objects.append(model) print("COCO to TDW objects:", objects) # print(len(objects)) # Stores object categories that other objects can be placed upon (e.g. table, chair, couch, bed) surface_properties_list = TDW_COCO_models['table'] + TDW_COCO_models['chair'] + \ TDW_COCO_models['bed'] + TDW_COCO_models['couch'] + \ TDW_COCO_models['bench'] + TDW_COCO_models['refrigerator'] surface_categories = [] for surface_properties in surface_properties_list: surface_categories.append(surface_properties[0]) print("Surface Categories:", surface_categories) # Stores the actual surface object instances/ids alongside number of objects on the surface surface_object_ids = {} self.start() positions_list = [] # Stores current model locations and radii scene_dimensions = [] # Store scene/environment dimensions init_setup_commands = [{"$type": "set_screen_size", "width": 640, "height": 481}, {"$type": "set_render_quality", "render_quality": 1}] self.communicate(init_setup_commands) scene_lib = SceneLibrarian() # Disable physics when adding in new objects (objects float) self.communicate({"$type": "simulate_physics", "value": False}) for scene in scenes[1:]: # Load in scene # print("Scene", scene[0]) if scene[3] == "interior" and scene[0] == "box_room_2018": self.start() scene_record = scene_lib.get_record(scene[0]) self.communicate({"$type": "add_scene", "name": scene_record.name, "url": scene_record.get_url()}) # Gets dimensions of environments (e.g. inside, outside) in the scene # This command returns environment data in the form of a list of serialized byte arrays scene_bytes = self.communicate({"$type": "send_environments", "frequency": "once"}) # Iterating through data and parsing byte array # Ignoring the last element (the frame count) for b in scene_bytes[:-1]: e = Environments(b) for i in range(e.get_num()): center = e.get_center(i) bounds = e.get_bounds(i) env_id = e.get_id(i) scene_dimensions = [center, bounds, env_id] # Center, bounds are tuples # Must come before set_pass_masks avatar_position = TDWUtils.array_to_vector3([0.9 * scene_dimensions[1][0] / 2, scene_dimensions[1][1] / 2, 0]) # print("Avatar Position:", avatar_position) self.communicate(TDWUtils.create_avatar(avatar_id="avatar", position=avatar_position, look_at={"x": 0, "y": scene_dimensions[0][1] / 2, "z": 0})) # Set collision mode self.communicate({"$type": "set_avatar_collision_detection_mode", "mode": "continuous_speculative", "avatar_id": "avatar"}) # Alter FOV self.communicate({"$type": "set_field_of_view", "field_of_view": 80, "avatar_id": "avatar"}) # # Gets rid of header (Model: Category) # objects = TDW_COCO_models[1:] # random.shuffle(objects) obj_count = 0 obj_overlaps = 0 # Number of failed attempts to place object due to over-dense objects area while obj_count < objects_in_scene and obj_overlaps < 5: # Handles if object has been added to a flat surface added_to_surface = False print("Object COUNT:", obj_count) # Need to have random position for Bounds Data to return meaningful info valid_obj_pos = {"x": random.uniform(-1 * scene_dimensions[1][0] / 2, 0.5 * scene_dimensions[1][0] / 2), "y": scene_dimensions[1][1] / 4, "z": random.uniform(-0.9 * scene_dimensions[1][2] / 2, 0.9 * scene_dimensions[1][2] / 2)} print("First random position") # Add in the object at random position # Object will later be removed or updated accordingly after performing collision calculations record = ModelLibrarian(library="models_full.json").get_record(objects[obj_count][0]) print("Record gotten") print(objects[obj_count][0]) o_id = self.communicate({"$type": "add_object", "name": objects[obj_count][0], "url": record.get_url(), "scale_factor": record.scale_factor, "position": valid_obj_pos, "rotation": TDWUtils.VECTOR3_ZERO, "category": record.wcategory, "id": obj_count}) print("Random first add") # Returns bound data for added object bounds_data = self.communicate({"$type": "send_bounds", "frequency": "once"}) print("Bounds returned") # Appends object, with information on position and obj_radius, to positions_list # Length of buffer array should be 1 # print("Bounds Data:", bounds_data) for b in bounds_data[:-1]: # print("Buffer Loop:", b) b_id = OutputData.get_data_type_id(b) if b_id == "boun": # print("BOUNDS") o = Bounds(b) # print("# of Objects:", o.get_num()) # print("# of Failed Attempts:", obj_overlaps) # print("Buffer Array:", b) # print("Bounds Object:", o) for i in range(o.get_num()): print("Object ID:", o.get_id(i)) print("obj_count:", obj_count) print("Object:", objects[obj_count][0], "Category:", objects[obj_count][1]) # print("Object Center:", o.get_center(i)) # Only want to compute valid_position for object we are about to add # Skip any computation if this is not the case if o.get_id(i) != obj_count: continue # Useful for detecting if object fits in environment # print("Calculating if object fits in environment") width = distance.euclidean(o.get_left(i), o.get_right(i)) depth = distance.euclidean(o.get_front(i), o.get_back(i)) height = distance.euclidean(o.get_top(i), o.get_bottom(i)) # print("Width:", width) # print("Depth:", depth) # ("Height:", height) # Useful for avoiding object overlap # print("Calculating Object Bounds") center_to_top = distance.euclidean(o.get_center(i), o.get_top(i)) center_to_bottom = distance.euclidean(o.get_center(i), o.get_bottom(i)) center_to_left = distance.euclidean(o.get_center(i), o.get_left(i)) center_to_right = distance.euclidean(o.get_center(i), o.get_right(i)) center_to_front = distance.euclidean(o.get_center(i), o.get_front(i)) center_to_back = distance.euclidean(o.get_center(i), o.get_back(i)) # Max object radius (center to diagonal of bounding box) obj_radius = \ max(math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_front ** 2), math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_front ** 2), math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_back ** 2), math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_back ** 2), math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_front ** 2), math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_front ** 2), math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_back ** 2), math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_back ** 2)) # print("Obj_Radius:", obj_radius) # Set sweeping radius, based on scene plane dimensions l_radius = random.uniform(0, min(0.5 * scene_dimensions[1][0] / 2, 0.5 * scene_dimensions[1][2] / 2)) # Checking that object fits in scene viewing if (width > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or depth > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or height > 0.7 * scene_dimensions[1][1]): print("Object does not fit in scene") self.communicate([{"$type": "send_images", "frequency": "never"}, {"$type": "destroy_object", "id": obj_count}]) # Ensures next attempt to load in item is not the same item as before random.shuffle(objects) break # Not possible to find valid object position -- too many overlapping objects elif (not self._get_object_position(scene_dimensions=scene_dimensions, object_positions=positions_list, object_to_add_radius=obj_radius, max_tries=20, location_radius=l_radius)[0]): print("Could not calculate valid object location") self.communicate([{"$type": "send_images", "frequency": "never"}, {"$type": "destroy_object", "id": obj_count}]) obj_overlaps += 1 # Ensures next attempt to load in item is not the same item as before random.shuffle(objects) break # Find appropriate, non-overlapping object position # Reset object position to the valid position else: print("Object fits in scene") # Check if object fits on table, chair, couch, etc. # Add object if it fits, place it somewhere on top of the surface for surface_id in surface_object_ids.keys(): print("Surface ID:", surface_id) # Skip placement feasibility if the object is already a surface-type object # Ex. no chair on top of a table if objects[obj_count][0] in surface_categories: print("Object: %s is already a surface object" % objects[obj_count][0]) break # Check how many objects are on surface if surface_object_ids[surface_id] >= 3: print("Too many objects on surface") print("From surface objects dict:", surface_object_ids[surface_id]) continue surface_bounds = self.get_bounds_data(surface_id) surface_area = distance.euclidean(surface_bounds.get_left(0), surface_bounds.get_right(0)) * \ distance.euclidean(surface_bounds.get_front(0), surface_bounds.get_back(0)) obj_area = width * height if obj_area < surface_area: s_center_to_top = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_top(0)) s_center_to_bottom = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_bottom(0)) s_center_to_left = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_left(0)) s_center_to_right = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_right(0)) s_center_to_front = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_front(0)) s_center_to_back = distance.euclidean(surface_bounds.get_center(0), surface_bounds.get_back(0)) surface_radius = \ max(math.sqrt( s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2), math.sqrt( s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2), math.sqrt( s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2), math.sqrt( s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2), math.sqrt( s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2), math.sqrt( s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2), math.sqrt( s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2), math.sqrt( s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2)) print("Surface-type object") self.communicate({"$type": "destroy_object", "id": obj_count}) # Adding the object to the top of the surface on_pos = surface_bounds.get_top(0) on_y = on_pos[1] on_pos = TDWUtils.get_random_point_in_circle(np.array(on_pos), 0.7 * surface_radius) on_pos[1] = on_y on_pos = TDWUtils.array_to_vector3(on_pos) on_rot = {"x": 0, "y": random.uniform(-45, 45), "z": 0} # Add the object. print("Model Name on Surface:", objects[obj_count][0]) record = ModelLibrarian(library="models_full.json").get_record( objects[obj_count][0]) on_id = self.communicate({"$type": "add_object", "name": objects[obj_count][0], "url": record.get_url(), "scale_factor": record.scale_factor, "position": on_pos, "rotation": on_rot, "category": record.wcategory, "id": obj_count}) obj_count += 1 surface_object_ids[surface_id] += 1 object_ids.append(obj_count) print("Object added on top of surface") added_to_surface = True # Breaking out of surface objects loop break if added_to_surface: print("Breaking out of object loop") # Breaking out of object loop break print("Post-surface") valid_obj_pos = self._get_object_position(scene_dimensions=scene_dimensions, object_positions=positions_list, object_to_add_radius=obj_radius, max_tries=20, location_radius=l_radius)[1] print("Position calculated") positions_list.append(ObjectPosition(valid_obj_pos, obj_radius)) self.communicate([{"$type": "send_images", "frequency": "never"}, {"$type": "destroy_object", "id": obj_count}]) added_object_id = self.communicate({"$type": "add_object", "name": objects[obj_count][0], "url": record.get_url(), "scale_factor": record.scale_factor, "position": valid_obj_pos, "rotation": {"x": 0, "y": 0, "z": 0}, "category": record.wcategory, "id": obj_count}) # print("Object ID:", added_object_id) print("Regular object add") object_ids.append(added_object_id) # If TDW model belongs to surface categories, store id_information if objects[obj_count][0] in surface_categories: surface_object_ids[obj_count] = 0 # Rotate the object randomly print("Rotating object") self.communicate({"$type": "rotate_object_by", "angle": random.uniform(-45, 45), "axis": "yaw", "id": obj_count, "is_world": True}) # Minimal rotating for position differences # Don't rotate the object if doing so will result in overlap into scene if not (o.get_bottom(i)[1] < 0 or o.get_top(i)[1] > 0.9 * scene_dimensions[1][1]): pitch_angle = random.uniform(-45, 45) self.communicate({"$type": "rotate_object_by", "angle": pitch_angle, "axis": "pitch", "id": obj_count, "is_world": True}) roll_angle = random.uniform(-45, 45) self.communicate({"$type": "rotate_object_by", "angle": roll_angle, "axis": "roll", "id": obj_count, "is_world": True}) # Don't need this for just changing positions # Setting random materials/textures # Looping through sub-objects and materials sub_count = 0 for sub_object in record.substructure: # Loop through materials in sub-objects for j in range(len(sub_object)): # Get random material and load in material = random.choice(materials[1:]) self.load_material(material) print("Material loaded") # Set random material on material of sub-object self.communicate({"$type": "set_visual_material", "material_index": j, "material_name": material[0], "object_name": sub_object['name'], "id": obj_count}) print("Material set") sub_count += 1 if sub_count > 10: break break print("Updating count") obj_count += 1 print("Breaking out of object_id loop") break # Break out of buffer loop print("Breaking out of buffer loop") break # Move onto next iteration of while loop (next object to load in) print("Object added - next while loop iteration") continue # for i in range(200): # self.communicate({"$type": "simulate_physics", # "value": False}) # Enable image capture self.communicate({"$type": "set_pass_masks", "avatar_id": "avatar", "pass_masks": ["_img", "_id"]}) self.communicate({"$type": "send_images", "frequency": "always"}) # Capture scene # NOTE: THESE SCENES GET REPLACED IN THE TARGET DIRECTORY scene_data = self.communicate({"$type": "look_at_position", "avatar_id": "avatar", "position": {"x": 0, "y": scene_dimensions[0][1] / 2, "z": 0}}) images = Images(scene_data[0]) TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=path) print("Object ids:", object_ids)
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