Beispiel #1
0
    def run(self):
        # Load the streamed scene and add controller rig.
        self.load_streamed_scene(scene="tdw_room")

        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 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 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
Beispiel #4
0
    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
Beispiel #5
0
    # Add a camera to the scene.
    commands.extend(
        TDWUtils.create_avatar(position={
            "x": -2.49,
            "y": 4,
            "z": 0
        },
                               look_at={
                                   "x": 0,
                                   "y": 0,
                                   "z": 0
                               }))
    resp = c.communicate(commands)
    wheel_ids = []
    for i in range(len(resp) - 1):
        r_id = OutputData.get_data_type_id(resp[i])
        if r_id == "srob":
            sr = StaticRobot(resp[i])
            for j in range(sr.get_num_joints()):
                joint_id = sr.get_joint_id(j)
                joint_name = sr.get_joint_name(j)
                print(joint_name, joint_id)
                # Find all of the wheels.
                if "wheel" in joint_name:
                    wheel_ids.append(joint_id)

    # Move the wheels forward.
    commands = []
    for wheel_id in wheel_ids:
        commands.append({
            "$type": "set_revolute_target",
Beispiel #6
0
        probe_material=args.pmaterial,
        zone_material=args.zmaterial,
        zone_color=args.zcolor,
        zone_friction=args.zfriction,
        distractor_types=args.distractor,
        distractor_categories=args.distractor_categories,
        num_distractors=args.num_distractors,
        occluder_types=args.occluder,
        occluder_categories=args.occluder_categories,
        num_occluders=args.num_occluders,
        flex_only=args.only_use_flex_objects,
        no_moving_distractors=args.no_moving_distractors,
        use_test_mode_colors=args.use_test_mode_colors
    )

    if bool(args.run):
        DC.run(num=args.num,
               output_dir=args.dir,
               temp_path=args.temp,
               width=args.width,
               height=args.height,
               write_passes=args.write_passes.split(','),
                save_passes=args.save_passes.split(','),
                save_movies=args.save_movies,
                save_labels=args.save_labels,
                save_meshes=args.save_meshes,
                args_dict=vars(args))
    else:
        end = DC.communicate({"$type": "terminate"})
        print([OutputData.get_data_type_id(r) for r in end])
    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
Beispiel #8
0
    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
Beispiel #9
0
 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"
         ],
         "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 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([])
Beispiel #11
0
    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")
Beispiel #12
0
    def run(self):
        self.start()
        self.communicate(TDWUtils.create_empty_room(12, 12))

        # Find the local asset bundle for this platform.
        url = "file:///" + str(
            Path("composite_objects/" + platform.system() +
                 "/test_composite_object").resolve())

        # Add the local object.
        o_id = self.get_unique_id()
        self.communicate([{
            "$type": "add_object",
            "name": "test_composite_object",
            "url": url,
            "scale_factor": 1,
            "id": o_id
        }, {
            "$type": "set_mass",
            "id": o_id,
            "mass": 100
        }])

        self.communicate(
            TDWUtils.create_avatar(position={
                "x": 0,
                "y": 1.49,
                "z": -2.77
            }))

        # Test that each sub-object has a unique segmentation color.
        resp = self.communicate({
            "$type": "send_segmentation_colors",
            "ids": []
        })
        colors = SegmentationColors(resp[0])
        segmentation_colors = []
        # There are 4 objects: The parent object, the motor, the "base" of the motor, and the light.
        assert colors.get_num() == 4, colors.get_num()
        print("Segmentation colors:")
        for i in range(colors.get_num()):
            print(colors.get_object_id(i), colors.get_object_color(i))

            # Cache the color for the next test.
            segmentation_colors.append(colors.get_object_color(i))

        # Test that each sub-object is visible to the avatar.
        resp = self.communicate({
            "$type": "send_id_pass_segmentation_colors",
            "frequency": "once"
        })
        colors = IdPassSegmentationColors(resp[0])
        # There are three visible objects (the light won't be visible in the _id pass).
        assert colors.get_num_segmentation_colors() == 3

        print("\nObserved colors:")

        # Test that the colors observed by the avatar are in the cache.
        for i in range(colors.get_num_segmentation_colors()):
            observed_color = colors.get_segmentation_color(i)
            assert observed_color in segmentation_colors
            print(observed_color)

        # Get composite objects data.
        resp = self.communicate({"$type": "send_composite_objects"})
        assert len(resp) > 1
        assert OutputData.get_data_type_id(resp[0]) == "comp"

        o = CompositeObjects(resp[0])
        assert o.get_num() == 1
        # There are 3 sub-objects: The motor, the "base" of the motor, and the light.
        assert o.get_num_sub_objects(0) == 3

        print("\nCompositeObjects: ")

        commands = []
        lights = []

        # Iterate through each sub-object.
        for s in range(o.get_num_sub_objects(0)):
            sub_object_id = o.get_sub_object_id(0, s)
            s_type = o.get_sub_object_machine_type(0, s)
            print(sub_object_id, s_type)

            # Add commands depending on the type of sub-object.
            if s_type == "motor":
                # Start the motor.
                commands.append({
                    "$type": "set_motor",
                    "target_velocity": 500,
                    "force": 500,
                    "id": sub_object_id
                })
            elif s_type == "light":
                commands.append({
                    "$type": "set_sub_object_light",
                    "is_on": True,
                    "id": sub_object_id
                })
                lights.append(sub_object_id)

        self.communicate(commands)

        is_on = True
        for i in range(1000):
            commands = []
            # Every 50 frames, blink the lights on and off.
            if i % 50 == 0:
                is_on = not is_on
                for light in lights:
                    commands.append({
                        "$type": "set_sub_object_light",
                        "is_on": is_on,
                        "id": light
                    })
            else:
                commands.append({"$type": "do_nothing"})
            self.communicate(commands)
Beispiel #13
0
    def run(self):
        o_id = 0
        self.start()
        # 1. Create a room.
        # 2. Create a Flex container.
        # 3. Add a cloth object and set the cloth actor.
        # 4. Request Flex particle output data (this frame only).
        commands = [
            TDWUtils.create_empty_room(12, 12), {
                "$type": "create_flex_container",
                "collision_distance": 0.001,
                "static_friction": 1.0,
                "dynamic_friction": 1.0,
                "iteration_count": 12,
                "substep_count": 12,
                "radius": 0.1875,
                "damping": 0,
                "drag": 0
            },
            self.get_add_object("cloth_square",
                                library="models_special.json",
                                position={
                                    "x": 0,
                                    "y": 2,
                                    "z": 0
                                },
                                rotation={
                                    "x": 0,
                                    "y": 0,
                                    "z": 0
                                },
                                object_id=o_id), {
                                    "$type": "set_flex_cloth_actor",
                                    "id": o_id,
                                    "mesh_tesselation": 1,
                                    "stretch_stiffness": 0.5620341548096974,
                                    "bend_stiffness": 0.6528988964052056,
                                    "tether_stiffness": 0.7984931184979334,
                                    "tether_give": 0,
                                    "pressure": 0,
                                    "mass_scale": 1
                                }, {
                                    "$type": "assign_flex_container",
                                    "container_id": 0,
                                    "id": o_id
                                }, {
                                    "$type": "send_flex_particles",
                                    "frequency": "once"
                                }
        ]
        # Add to the list of commands: create the avatar, teleport it to a position, look at a position.
        commands.extend(
            TDWUtils.create_avatar(position={
                "x": -2.75,
                "y": 2.3,
                "z": -2
            },
                                   look_at={
                                       "x": 0,
                                       "y": 0.25,
                                       "z": 0
                                   }))
        # Send the commands.
        resp = self.communicate(commands)

        for r in resp[:-1]:
            r_id = OutputData.get_data_type_id(r)
            # This is FlexParticles data.
            if r_id == "flex":
                fp = FlexParticles(resp[0])
                # The maximum distance between a particle and the center of the cloth.
                max_d = 0
                # The ID of the particle at the maximum distance.
                max_id = 0
                p_id = 0
                for p in fp.get_particles(0):
                    # Get the particle that is furthest from the center.
                    d = np.linalg.norm(p[:-1] - 1)
                    if d > max_d:
                        max_d = d
                        max_id = p_id
                    p_id += 1
                # Set the particle as "fixed".
                self.communicate({
                    "$type": "set_flex_particle_fixed",
                    "is_fixed": True,
                    "particle_id": max_id,
                    "id": o_id
                })

            for i in range(800):
                self.communicate([])
Beispiel #14
0
    def run(self):
        self.start()
        self.communicate(TDWUtils.create_empty_room(12, 12))

        # Create the objects.
        statue_id = self.add_object("satiro_sculpture",
                                    position={
                                        "x": 2,
                                        "y": 0,
                                        "z": 1
                                    })
        # Request segmentation colors.
        resp = self.communicate({"$type": "send_segmentation_colors"})
        segmentation_colors = SegmentationColors(resp[0])

        # Get the segmentation color of the sculpture.
        statue_color = None
        for i in range(segmentation_colors.get_num()):
            if segmentation_colors.get_object_id(i) == statue_id:
                statue_color = segmentation_colors.get_object_color(i)
                break

        # Create the VR Rig.
        self.communicate({"$type": "create_vr_rig"})
        # Attach an avatar to the VR rig.
        self.communicate({"$type": "attach_avatar_to_vr_rig", "id": "a"})
        # Request the colors of objects currently observed by the avatar per frame.
        # Request VR rig data per frame.
        # Reduce render quality in order to improve framerate.
        self.communicate([{
            "$type": "send_id_pass_segmentation_colors",
            "frequency": "always"
        }, {
            "$type": "send_vr_rig",
            "frequency": "always"
        }, {
            "$type": "set_post_process",
            "value": False
        }, {
            "$type": "set_render_quality",
            "render_quality": 0
        }])

        while True:
            resp = self.communicate({"$type": "do_nothing"})
            head_rotation = None
            can_see_statue = False
            for r in resp[:-1]:
                r_id = OutputData.get_data_type_id(r)
                # Get the head rotation.
                if r_id == "vrri":
                    head_rotation = VRRig(r).get_head_rotation()
                # Check if we can see the statue.
                elif r_id == "ipsc":
                    observed_objects = IdPassSegmentationColors(r)
                    for i in range(
                            observed_objects.get_num_segmentation_colors()):
                        if observed_objects.get_segmentation_color(
                                i) == statue_color:
                            can_see_statue = True
            if can_see_statue:
                print("You can see the object!\nHead rotation: " +
                      str(head_rotation) + "\n")
Beispiel #15
0
    def run(self):
        self.start()
        init_setup_commands = [{
            "$type": "set_screen_size",
            "width": 600,
            "height": 480
        }, {
            "$type": "set_render_quality",
            "render_quality": 5
        }]
        self.communicate(init_setup_commands)

        self.communicate(self.get_add_hdri_skybox("industrial_sunset_4k"))

        # Create an empty room.
        self.communicate(TDWUtils.create_empty_room(8, 8))
        # Disable physics.
        self.communicate({"$type": "set_gravity", "value": False})

        # Add the avatar.
        self.communicate(
            TDWUtils.create_avatar(position={
                "x": 1,
                "y": 1.5,
                "z": 0.3
            },
                                   look_at=TDWUtils.array_to_vector3(
                                       [2.5, 0.5, 0.5]),
                                   avatar_id="avatar"))

        lib = MaterialLibrarian(library="materials_med.json")
        record = lib.get_record("concrete_049")
        self.communicate({
            "$type": "add_material",
            "name": "concrete_049",
            "url": record.get_url()
        })
        self.communicate({
            "$type": "set_proc_gen_floor_material",
            "name": "concrete_049"
        })
        # self.communicate({"$type": "set_proc_gen_walls_material", "name": "concrete"})

        # self.communicate({"$type": "set_field_of_view",
        #                   "field_of_view": 68.0,
        #                   "avatar_id": "avatar"})

        bench = self.add_object(model_name="cgaxis_models_51_19_01",
                                position={
                                    "x": 2.5,
                                    "y": 0,
                                    "z": 0.5
                                },
                                rotation={
                                    "x": 0,
                                    "y": 90,
                                    "z": 0
                                },
                                library="models_full.json")

        bench_bounds = self.get_bounds_data(bench)
        top = bench_bounds.get_top(0)

        self.add_object(model_name="b05_cat_model_3dmax2012",
                        position={
                            "x": 2.5,
                            "y": top[1],
                            "z": 0.2
                        },
                        rotation={
                            "x": 0,
                            "y": 0,
                            "z": 0
                        },
                        library="models_full.json")

        self.add_object(model_name="azor",
                        position={
                            "x": 3,
                            "y": 0,
                            "z": 0.5
                        },
                        rotation={
                            "x": 0,
                            "y": 90,
                            "z": 0
                        },
                        library="models_full.json")
        """ Post-scene construction """

        # Enable image capture
        self.communicate({
            "$type": "set_pass_masks",
            "avatar_id": "avatar",
            "pass_masks": ["_id"]
        })

        scene_data = self.communicate({
            "$type":
            "look_at_position",
            "avatar_id":
            "avatar",
            "position":
            TDWUtils.array_to_vector3([3, 0.5, 0.5])
        })

        # images = Images(scene_data[0])
        im = cv2.imread(
            "/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/"
            "compare_COCO_TDW/replicated_images/exterior/id_bench_bike.png")
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)

        # WRITE THIS TO JSON!
        segmentation_data = []
        COCO_to_TDW = pd.read_csv(
            '/Users/leonard/Desktop/coco/COCO_to_TDW.csv',
            header=None,
            index_col=0,
            squeeze=True).to_dict()

        segmentation_colors = self.communicate({
            "$type": "send_segmentation_colors",
            "frequency": "once"
        })
        for r in segmentation_colors[:-1]:
            r_id = OutputData.get_data_type_id(r)
            if r_id == "segm":
                s = SegmentationColors(r)
                for i in range(s.get_num()):
                    name = s.get_object_name(i)
                    category = get_COCO(name, COCO_to_TDW)

                    color = s.get_object_color(i)
                    indices = np.where(im != 0)  # Segmentation pixels
                    print(color)
                    print(indices)
                    print(im[indices[0][0]]
                          )  # some random RGB val -- MUSt match one of the 3
                    segmentation_data.append([indices, category])

        # TDWUtils.save_images(images, "bench_bike",
        #                      output_directory="/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/compare_COCO_TDW/replicated_images/exterior")

        # Consider using zip() to clean up coordinates

        return segmentation_data
    def run(self) -> None:
        self.start()

        # Get IDs for the object, the table, and the avatars.
        o_id = self.get_unique_id()
        table_id = self.get_unique_id()
        table_position = {"x": 0, "y": 0, "z": 6.8}
        a = "a"

        # 1. Create the room.
        commands = [TDWUtils.create_empty_room(20, 20)]
        # 2. Add the avatar.
        commands.extend(
            TDWUtils.create_avatar(avatar_type="A_StickyMitten_Adult",
                                   avatar_id="a"))
        # 3. Add the objects. Scale the objects. Set a high mass for the table.
        # 4. Disable the avatar's cameras.
        # 5. Set the stickiness of the avatar's left mitten. Set a high drag. Rotate the head.
        # 6. Request AvatarChildrenNames data and Bounds data (this frame only).
        # 7. Strengthen the left shoulder a bit.
        commands.extend([
            self.get_add_object("jug05",
                                position={
                                    "x": -0.417,
                                    "y": 0.197,
                                    "z": 0.139
                                },
                                rotation={
                                    "x": 90,
                                    "y": 0,
                                    "z": 0
                                },
                                object_id=o_id), {
                                    "$type": "scale_object",
                                    "scale_factor": {
                                        "x": 2,
                                        "y": 2,
                                        "z": 2
                                    },
                                    "id": o_id
                                },
            self.get_add_object("small_table_green_marble",
                                position={
                                    "x": 0,
                                    "y": 0,
                                    "z": 6.8
                                },
                                rotation={
                                    "x": 0,
                                    "y": 0,
                                    "z": 0
                                },
                                object_id=table_id), {
                                    "$type": "scale_object",
                                    "scale_factor": {
                                        "x": 2,
                                        "y": 0.5,
                                        "z": 2
                                    },
                                    "id": table_id
                                },
            {
                "$type": "set_object_collision_detection_mode",
                "id": table_id,
                "mode": "continuous_dynamic"
            }, {
                "$type": "set_object_collision_detection_mode",
                "id": o_id,
                "mode": "continuous_dynamic"
            }, {
                "$type": "set_mass",
                "mass": 100,
                "id": table_id
            }, {
                "$type": "toggle_image_sensor",
                "sensor_name": "SensorContainer",
                "avatar_id": a
            }, {
                "$type": "toggle_image_sensor",
                "sensor_name": "FollowCamera",
                "avatar_id": a
            }, {
                "$type": "set_stickiness",
                "sub_mitten": "back",
                "sticky": True,
                "is_left": True,
                "avatar_id": a
            }, {
                "$type": "set_avatar_drag",
                "drag": 1000,
                "angular_drag": 1000,
                "avatar_id": a
            }, {
                "$type": "rotate_head_by",
                "axis": "pitch",
                "angle": 5
            }, {
                "$type": "send_avatar_children_names",
                "ids": [a],
                "frequency": "once"
            }, {
                "$type": "send_bounds",
                "ids": [table_id],
                "frequency": "once"
            }, {
                "$type": "adjust_joint_force_by",
                "delta": 2,
                "joint": "shoulder_left",
                "axis": "pitch"
            }
        ])
        # 7. Add a 3rd-person camera.
        commands.extend(
            TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
                                   avatar_id="c",
                                   position={
                                       "x": -3.9,
                                       "y": 2.3,
                                       "z": 4.3
                                   }))
        # 8. Request StickyMittenAvatar and Images data per-frame.
        commands.extend([{
            "$type": "set_pass_masks",
            "pass_masks": ["_img"],
            "avatar_id": self.cam_id
        }, {
            "$type": "send_avatars",
            "ids": [a],
            "frequency": "always"
        }, {
            "$type": "send_images",
            "ids": [self.cam_id],
            "frequency": "always"
        }])

        resp = self.communicate(commands)
        # Get the object ID of the left mitten, the size of the table, and the avatar's position.
        mitten_left_id = None
        table_size: Tuple[float, float, float] = (0, 0, 0)

        for r in resp[:-1]:
            r_id = OutputData.get_data_type_id(r)
            # Get the mitten ID.
            if r_id == "avcn":
                avcn = AvatarChildrenNames(r)
                for i in range(avcn.get_num_children()):
                    if avcn.get_child_name(i) == "mitten_left":
                        mitten_left_id = avcn.get_child_id(i)
            # Get the table bounds.
            elif r_id == "boun":
                boun = Bounds(r)
                for i in range(boun.get_num()):
                    if boun.get_id(i) == table_id:
                        table_size = (boun.get_right(i)[0] -
                                      boun.get_left(i)[0], boun.get_top(i)[1] -
                                      boun.get_bottom(i)[1],
                                      boun.get_front(i)[2] -
                                      boun.get_back(i)[2])
            # Get the avatar's position.
            elif r_id == "avsm":
                self.avatar_position = AvatarStickyMitten(r).get_position()

        # Pick up the object.
        # Lift up the object.
        self._bend_arm_joints([{
            "$type": "pick_up_proximity",
            "distance": 20,
            "grip": 10000,
            "is_left": True,
            "avatar_id": a
        }, {
            "$type": "rotate_head_by",
            "axis": "pitch",
            "angle": 20,
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_by",
            "angle": 25,
            "joint": "shoulder_left",
            "axis": "pitch",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_by",
            "angle": -25,
            "joint": "shoulder_left",
            "axis": "yaw",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_by",
            "angle": 60,
            "joint": "shoulder_left",
            "axis": "roll",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_by",
            "angle": 100,
            "joint": "elbow_left",
            "axis": "pitch",
            "avatar_id": a
        }])
        # Allow the avatar to move again.
        self._do_frame({
            "$type": "set_avatar_drag",
            "drag": 0.125,
            "angular_drag": 1000,
            "avatar_id": a
        })

        # Move to the table.
        move_to_table = True
        # The position of the side of the table the avatar is aiming for.
        table_side_position = {
            "x": table_position["x"],
            "y": 0,
            "z": table_position["z"] - table_size[2]
        }
        while move_to_table:
            # Stop moving if we are close enough.
            if TDWUtils.get_distance(
                    table_side_position,
                    TDWUtils.array_to_vector3(self.avatar_position)) < 0.7:
                move_to_table = False
            # Keep moving forward.
            else:
                self._do_frame({
                    "$type": "move_avatar_forward_by",
                    "avatar_id": a,
                    "magnitude": 20
                })
        # Stop.
        # Allow the avatar to move again.
        self._do_frame({
            "$type": "set_avatar_drag",
            "drag": 1000,
            "angular_drag": 1000,
            "avatar_id": a
        })

        mitten_over_table = False
        target_z = table_side_position["z"] + 0.17
        # Keep lifting until the mitten is over the table.
        while not mitten_over_table:
            avsm = self._bend_arm_joints([{
                "$type": "bend_arm_joint_by",
                "angle": 15,
                "joint": "shoulder_left",
                "axis": "pitch",
                "avatar_id": a
            }, {
                "$type": "bend_arm_joint_by",
                "angle": -10,
                "joint": "elbow_left",
                "axis": "pitch",
                "avatar_id": a
            }])
            # Get the mitten.
            for i in range(avsm.get_num_body_parts()):
                if avsm.get_body_part_id(i) == mitten_left_id:
                    # Check if the mitten is over the table.
                    mitten_over_table = avsm.get_body_part_position(
                        i)[2] >= target_z

        # Drop the arm gently.
        mitten_near_table = False
        while not mitten_near_table:
            avsm = self._bend_arm_joints([{
                "$type": "bend_arm_joint_by",
                "angle": -25,
                "joint": "shoulder_left",
                "axis": "pitch",
                "avatar_id": a
            }])
            # Get the mitten.
            for i in range(avsm.get_num_body_parts()):
                if avsm.get_body_part_id(i) == mitten_left_id:
                    # Check if the mitten very close to the table surface.
                    mitten_near_table = avsm.get_body_part_position(
                        i)[1] - table_size[1] <= 0.1
        # Drop the object.
        # Allow the avatar to move again.
        self._do_frame([{
            "$type": "put_down",
            "is_left": True,
            "avatar_id": a
        }, {
            "$type": "set_avatar_drag",
            "drag": 0.125,
            "angular_drag": 1000,
            "avatar_id": a
        }])
        # Back away from the table.
        mitten_away_from_table = False
        while not mitten_away_from_table:
            avsm = self._do_frame({
                "$type": "move_avatar_forward_by",
                "avatar_id": a,
                "magnitude": -50
            })
            # Get the mitten.
            for i in range(avsm.get_num_body_parts()):
                if avsm.get_body_part_id(i) == mitten_left_id:
                    # Check if the mitten has cleared the table.
                    mitten_away_from_table = avsm.get_body_part_position(
                        i)[2] < target_z
        # Drop the arm.
        self._do_frame([{
            "$type": "rotate_head_by",
            "axis": "pitch",
            "angle": -25
        }, {
            "$type": "bend_arm_joint_to",
            "angle": 0,
            "joint": "shoulder_left",
            "axis": "pitch",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_to",
            "angle": 0,
            "joint": "shoulder_left",
            "axis": "yaw",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_to",
            "angle": 0,
            "joint": "shoulder_left",
            "axis": "roll",
            "avatar_id": a
        }, {
            "$type": "bend_arm_joint_to",
            "angle": 0,
            "joint": "elbow_left",
            "axis": "pitch",
            "avatar_id": a
        }])
        # Back away from the table.
        away_from_table = False
        while not away_from_table:
            self._do_frame({
                "$type": "move_avatar_forward_by",
                "avatar_id": a,
                "magnitude": -20
            })
            away_from_table = TDWUtils.get_distance(
                TDWUtils.array_to_vector3(self.avatar_position),
                table_position) > 1.8
        # Let the joints drop.
        self._bend_arm_joints([])
        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]:
        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))
        collision_states = np.empty(dtype=str, shape=(0, 1))
        # 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

        # rtypes = [OutputData.get_data_type_id(r) for r in resp[:-1]]
        # print(frame_num, "COLLISION" if 'coll' in rtypes else "")

        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)):
                    try:
                        velocities[i] = ri_dict[o_id]["vel"]
                        angular_velocities[i] = ri_dict[o_id]["ang"]
                    except KeyError:
                        print("Couldn't store velocity data for object %d" %
                              o_id)
                        print("frame num", frame_num)
                        print("ri_dict", ri_dict)
                        print([OutputData.get_data_type_id(r) for r in resp])
            elif r_id == "coll":
                co = Collision(r)
                collision_states = np.append(collision_states, co.get_state())
                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")
        collisions.create_dataset("states",
                                  data=collision_states.astype('S'),
                                  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 _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 trial(self, scene: Scene, record: ModelRecord, output_path: Path,
              scene_index: int) -> None:
        """
        Run a trial in a scene that has been initialized.

        :param scene: Data for the current scene.
        :param record: The model's metadata record.
        :param output_path: Write the .wav file to this path.
        :param scene_index: The scene identifier.
        """

        self.py_impact.reset(initial_amp=0.05)

        # Initialize the scene, positioning objects, furniture, etc.
        resp = self.communicate(scene.initialize_scene(self))
        center = scene.get_center(self)

        max_y = scene.get_max_y()

        # The object's initial position.
        o_x = RNG.uniform(center["x"] - 0.15, center["x"] + 0.15)
        o_y = RNG.uniform(max_y - 0.5, max_y)
        o_z = RNG.uniform(center["z"] - 0.15, center["z"] + 0.15)
        # Physics values.
        mass = self.object_info[record.name].mass + RNG.uniform(
            self.object_info[record.name].mass * -0.15,
            self.object_info[record.name].mass * 0.15)
        static_friction = RNG.uniform(0.1, 0.3)
        dynamic_friction = RNG.uniform(0.7, 0.9)
        # Angles of rotation.
        yaw = RNG.uniform(-30, 30)
        pitch = RNG.uniform(0, 45)
        roll = RNG.uniform(-45, 45)
        # The force applied to the object.
        force = RNG.uniform(0, 5)
        # The avatar's position.
        a_r = RNG.uniform(1.5, 2.2)
        a_x = center["x"] + a_r
        a_y = RNG.uniform(1.5, 3)
        a_z = center["z"] + a_r
        cam_angle_min, cam_angle_max = scene.get_camera_angles()
        theta = np.radians(RNG.uniform(cam_angle_min, cam_angle_max))
        a_x = np.cos(theta) * (a_x - center["x"]) - np.sin(theta) * (
            a_z - center["z"]) + center["x"]
        a_z = np.sin(theta) * (a_x - center["x"]) + np.cos(theta) * (
            a_z - center["z"]) + center["z"]

        o_id = 0
        # Create the object and apply a force.
        commands = [{
            "$type": "add_object",
            "name": record.name,
            "url": record.get_url(),
            "scale_factor": record.scale_factor,
            "position": {
                "x": o_x,
                "y": o_y,
                "z": o_z
            },
            "category": record.wcategory,
            "id": o_id
        }, {
            "$type": "set_mass",
            "id": o_id,
            "mass": mass
        }, {
            "$type": "set_physic_material",
            "id": o_id,
            "bounciness": self.object_info[record.name].bounciness,
            "static_friction": static_friction,
            "dynamic_friction": dynamic_friction
        }, {
            "$type": "rotate_object_by",
            "angle": yaw,
            "id": o_id,
            "axis": "yaw",
            "is_world": True
        }, {
            "$type": "rotate_object_by",
            "angle": pitch,
            "id": o_id,
            "axis": "pitch",
            "is_world": True
        }, {
            "$type": "rotate_object_by",
            "angle": roll,
            "id": o_id,
            "axis": "roll",
            "is_world": True
        }, {
            "$type": "apply_force_magnitude_to_object",
            "magnitude": force,
            "id": o_id
        }, {
            "$type": "send_rigidbodies",
            "frequency": "always"
        }, {
            "$type": "send_collisions",
            "enter": True,
            "exit": False,
            "stay": False,
            "collision_types": ["obj", "env"]
        }, {
            "$type": "send_transforms",
            "frequency": "always"
        }]
        # Parse bounds data to get the centroid of all objects currently in the scene.
        bounds = Bounds(resp[0])
        if bounds.get_num() == 0:
            look_at = {"x": center["x"], "y": 0.1, "z": center["z"]}
        else:
            centers = []
            for i in range(bounds.get_num()):
                centers.append(bounds.get_center(i))
            centers_x, centers_y, centers_z = zip(*centers)
            centers_len = len(centers_x)
            look_at = {
                "x": sum(centers_x) / centers_len,
                "y": sum(centers_y) / centers_len,
                "z": sum(centers_z) / centers_len
            }
        # Add the avatar.
        # Set the position at a given distance (r) from the center of the scene.
        # Rotate around that position to a random angle constrained by the scene's min and max angles.
        commands.extend([{
            "$type": "teleport_avatar_to",
            "position": {
                "x": a_x,
                "y": a_y,
                "z": a_z
            }
        }, {
            "$type": "look_at_position",
            "position": look_at
        }])

        # Send the commands.
        resp = self.communicate(commands)

        AudioUtils.start(output_path=output_path, until=(0, 10))

        # Loop until all objects are sleeping.
        done = False
        while not done and AudioUtils.is_recording():
            commands = []
            collisions, environment_collisions, rigidbodies = PyImpact.get_collisions(
                resp)
            # Create impact sounds from object-object collisions.
            for collision in collisions:
                if PyImpact.is_valid_collision(collision):
                    # Get the audio material and amp.
                    collider_id = collision.get_collider_id()
                    collider_material, collider_amp = self._get_object_info(
                        collider_id, Scene.OBJECT_IDS, record.name)
                    collidee_id = collision.get_collider_id()
                    collidee_material, collidee_amp = self._get_object_info(
                        collidee_id, Scene.OBJECT_IDS, record.name)
                    impact_sound_command = self.py_impact.get_impact_sound_command(
                        collision=collision,
                        rigidbodies=rigidbodies,
                        target_id=collidee_id,
                        target_amp=collidee_amp,
                        target_mat=collidee_material.name,
                        other_id=collider_id,
                        other_mat=collider_material.name,
                        other_amp=collider_amp,
                        play_audio_data=False)
                    commands.append(impact_sound_command)
            # Create impact sounds from object-environment collisions.
            for collision in environment_collisions:
                collider_id = collision.get_object_id()
                if self._get_velocity(rigidbodies, collider_id) > 0:
                    collider_material, collider_amp = self._get_object_info(
                        collider_id, Scene.OBJECT_IDS, record.name)
                    surface_material = scene.get_surface_material()
                    impact_sound_command = self.py_impact.get_impact_sound_command(
                        collision=collision,
                        rigidbodies=rigidbodies,
                        target_id=collider_id,
                        target_amp=collider_amp,
                        target_mat=collider_material.name,
                        other_id=-1,
                        other_amp=0.01,
                        other_mat=surface_material.name,
                        play_audio_data=False)
                    commands.append(impact_sound_command)
            # If there were no collisions, check for movement. If nothing is moving, the trial is done.
            if len(commands) == 0:
                transforms = AudioDataset._get_transforms(resp)
                done = True
                for i in range(rigidbodies.get_num()):
                    if self._is_moving(rigidbodies.get_id(i), transforms,
                                       rigidbodies):
                        done = False
                        break
            # Continue the trial.
            if not done:
                resp = self.communicate(commands)

        # Stop listening for anything except audio data..
        resp = self.communicate([{
            "$type": "send_rigidbodies",
            "frequency": "never"
        }, {
            "$type": "send_transforms",
            "frequency": "never"
        }, {
            "$type": "send_collisions",
            "enter": False,
            "exit": False,
            "stay": False,
            "collision_types": []
        }, {
            "$type": "send_audio_sources",
            "frequency": "always"
        }])
        # Wait for the audio to finish.
        done = False
        while not done and AudioUtils.is_recording():
            done = True
            for r in resp[:-1]:
                if OutputData.get_data_type_id(r) == "audi":
                    audio_sources = AudioSources(r)
                    for i in range(audio_sources.get_num()):
                        if audio_sources.get_is_playing(i):
                            done = False
            if not done:
                resp = self.communicate([])
        # Cleanup.
        commands = [{
            "$type": "send_audio_sources",
            "frequency": "never"
        }, {
            "$type": "destroy_object",
            "id": o_id
        }]
        for scene_object_id in Scene.OBJECT_IDS:
            commands.append({"$type": "destroy_object", "id": scene_object_id})
        self.communicate(commands)

        # Insert the trial's values into the database.
        self.db_c.execute(
            "INSERT INTO sound20k VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
            (output_path.name, scene_index, a_x, a_y, a_z, o_x, o_y, o_z, mass,
             static_friction, dynamic_friction, yaw, pitch, roll, force))
        self.conn.commit()
Beispiel #20
0
    def get_real_grayscale(
            self, o_id: int, a_id: str,
            e: Environment) -> (float, float, dict, dict, dict, dict):
        """
        Get the "real" grayscale value of an image we hope to capture.

        :param o_id: The ID of the object.
        :param a_id: The ID of the avatar.
        :param e: The environment.
        
        :return: (grayscale, distance, avatar_position, object_position, object_rotation, avatar_rotation)
        """

        # Get a random position for the avatar.
        a_p = np.array([
            RNG.uniform(-e.w / 2, e.w / 2) + e.x,
            RNG.uniform(0.4, e.h * self.max_height),
            RNG.uniform(-e.l / 2, e.l / 2) + e.z
        ])

        # Get a random distance from the avatar.
        d = RNG.uniform(0.8, 3)

        # Get a random position for the object constrained to the environment bounds.
        o_p = SingleObject.sample_spherical() * d
        # Clamp the y value to positive.
        o_p[1] = abs(o_p[1])
        o_p = a_p + o_p

        # Clamp the y value of the object.
        if o_p[1] > e.h * self.max_height:
            o_p[1] = e.h * self.max_height

        # Convert the avatar's position to a Vector3.
        a_p = TDWUtils.array_to_vector3(a_p)

        # Set random camera rotations.
        yaw = RNG.uniform(-15, 15)
        pitch = RNG.uniform(-15, 15)

        # Convert the object position to a Vector3.
        o_p = TDWUtils.array_to_vector3(o_p)

        # Add rotation commands.
        # If we're clamping the rotation, rotate the object within +/- 30 degrees on each axis.
        if self.clamp_rotation:
            o_rot = None
            commands = [{
                "$type": "rotate_object_to",
                "id": o_id,
                "rotation": {
                    "x": 0,
                    "y": 0,
                    "z": 0,
                    "w": 0
                }
            }, {
                "$type": "rotate_object_by",
                "id": o_id,
                "angle": RNG.uniform(-30, 30),
                "axis": "pitch"
            }, {
                "$type": "rotate_object_by",
                "id": o_id,
                "angle": RNG.uniform(-30, 30),
                "axis": "yaw"
            }, {
                "$type": "rotate_object_by",
                "id": o_id,
                "angle": RNG.uniform(-30, 30),
                "axis": "roll"
            }]
        # Set a totally random rotation.
        else:
            o_rot = {
                "x": RNG.uniform(-360, 360),
                "y": RNG.uniform(-360, 360),
                "z": RNG.uniform(-360, 360),
                "w": RNG.uniform(-360, 360),
            }
            commands = [{
                "$type": "rotate_object_to",
                "id": o_id,
                "rotation": o_rot
            }]

        # After rotating the object:
        # 1. Teleport the object.
        # 2. Teleport the avatar.
        # 3. Look at the object.
        # 4. Perturb the camera slightly.
        # 5. Send grayscale data and image sensor data.
        commands.extend([{
            "$type": "teleport_object",
            "id": o_id,
            "position": o_p
        }, {
            "$type": "teleport_avatar_to",
            "avatar_id": a_id,
            "position": a_p
        }, {
            "$type": "look_at",
            "avatar_id": a_id,
            "object_id": o_id,
            "use_centroid": True
        }, {
            "$type": "rotate_sensor_container_by",
            "angle": pitch,
            "avatar_id": "a",
            "axis": "pitch"
        }, {
            "$type": "rotate_sensor_container_by",
            "angle": yaw,
            "avatar_id": "a",
            "axis": "yaw"
        }, {
            "$type": "send_id_pass_grayscale",
            "frequency": "once"
        }, {
            "$type": "send_image_sensors",
            "frequency": "once"
        }])

        resp = self.communicate(commands)

        # Parse the output data:
        # 1. The grayscale value of the image.
        # 2. The camera rotation.
        grayscale = 0
        cam_rot = None
        for r in resp[:-1]:
            r_id = OutputData.get_data_type_id(r)
            if r_id == "idgs":
                grayscale = IdPassGrayscale(r).get_grayscale()
            elif r_id == "imse":
                cam_rot = ImageSensors(r).get_sensor_rotation(0)
                cam_rot = {
                    "x": cam_rot[0],
                    "y": cam_rot[1],
                    "z": cam_rot[2],
                    "w": cam_rot[3]
                }

        # If we clamped the rotation of the object, we need to know its quaternion.
        if self.clamp_rotation:
            resp = self.communicate({
                "$type": "send_transforms",
                "frequency": "once",
                "ids": [o_id]
            })
            t = Transforms(resp[0])
            o_rot = t.get_rotation(0)
            o_rot = {
                "x": o_rot[0],
                "y": o_rot[1],
                "z": o_rot[2],
                "w": o_rot[3],
            }

        return grayscale, d, a_p, o_p, o_rot, cam_rot
Beispiel #21
0
    def run(self) -> None:
        robot_id = 0
        self.start()

        # Create the scene. Add a robot.
        # Request static robot data for this frame only.
        # Request dynamic robot data per frame.
        commands = [
            TDWUtils.create_empty_room(12, 12),
            self.get_add_robot(name="ur3", robot_id=robot_id), {
                "$type": "send_static_robots",
                "frequency": "once"
            }, {
                "$type": "send_robots",
                "frequency": "always"
            }
        ]
        # Add an avatar to render the scene (just for demo purposes).
        commands.extend(
            TDWUtils.create_avatar(look_at=TDWUtils.VECTOR3_ZERO,
                                   position={
                                       "x": -0.881,
                                       "y": 0.836,
                                       "z": -1.396
                                   }))

        resp = self.communicate(commands)

        # Parse the output data for static robot data.
        static_robot: Optional[StaticRobot] = None
        for i in range(len(resp) - 1):
            r_id = OutputData.get_data_type_id(resp[i])
            if r_id == "srob":
                r = StaticRobot(resp[i])
                if r.get_id() == robot_id:
                    static_robot = r
                    break
        assert static_robot is not None, f"No static robot data: {resp}"

        # Get the IDs of the shoulder and the elbow.
        body_part_ids: Dict[str, int] = dict()
        for i in range(static_robot.get_num_joints()):
            b_id = static_robot.get_joint_id(i)
            b_name = static_robot.get_joint_name(i)
            body_part_ids[b_name] = b_id
        assert "Shoulder" in body_part_ids
        assert "Elbow" in body_part_ids

        # Rotate the shoulder and the elbow for two motions.
        # The values in this array are for the angle that the [shoulder, elbow] should rotate to per action.
        # For more complex actions, you will probably want to organize your commands differently.
        for angles in [[70, 90], [-30, -25]]:
            resp = self.communicate([{
                "$type": "set_revolute_target",
                "id": robot_id,
                "joint_id": body_part_ids["Shoulder"],
                "target": angles[0]
            }, {
                "$type": "set_revolute_target",
                "id": robot_id,
                "joint_id": body_part_ids["Elbow"],
                "target": angles[1]
            }])
            # Get the robot output data.
            robot = self.get_robot(resp=resp)
            angles_0 = self.get_joint_angles(robot=robot)
            # Wait for the joints to stop moving.
            moving = True
            while moving:
                robot = self.get_robot(resp=self.communicate([]))
                angles_1 = self.get_joint_angles(robot=robot)
                # Compare the current joint angles to the angles of the previous frame.
                moving = False
                for angle_0, angle_1 in zip(angles_0, angles_1):
                    if np.abs(angle_0 - angle_1) > 0.001:
                        moving = True
                        break
                angles_0 = angles_1
Beispiel #22
0
    def run(self):
        self.start()
        positions_list = []  # Stores current model locations and radii
        scene_dimensions = []  # Store scene/environment dimensions
        init_setup_commands = [{
            "$type": "set_screen_size",
            "width": 1280,
            "height": 962
        }, {
            "$type": "set_render_quality",
            "render_quality": 5
        }]
        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 = 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 < 30 and obj_overlaps < 5:
                    # 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)
                    }

                    # 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])
                    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
                    })

                    # Returns bound data for added object
                    bounds_data = self.communicate({
                        "$type": "send_bounds",
                        "frequency": "once"
                    })

                    # 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)
                                print("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.9 * scene_dimensions[1][0] / 2,
                                        0.9 * 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")
                                    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
                                    }])
                                    print("Object ready to reset")
                                    self.communicate([{
                                        "$type": "send_images",
                                        "frequency": "never"
                                    }, {
                                        "$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 reset")

                                    # 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
                                    })

                                    # 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
                                        })

                                    # 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

                # 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
                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)
Beispiel #23
0
                                           "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
Beispiel #24
0
    def run(self,
            num: int,
            output_dir: str,
            temp_path: str,
            width: int,
            height: int,
            write_passes: List[str] = PASSES,
            save_passes: List[str] = [],
            save_movies: bool = False,
            save_labels: bool = False,
            args_dict: dict={}) -> None:
        """
        Create the dataset.

        :param num: The number of trials in the dataset.
        :param output_dir: The root output directory.
        :param temp_path: Temporary path to a file being written.
        :param width: Screen width in pixels.
        :param height: Screen height in pixels.
        :param save_passes: a list of which passes to save out as PNGs (or convert to MP4)
        :param save_movies: whether to save out a movie of each trial
        :param save_labels: whether to save out JSON labels for the full trial set.
        """

        self._height, self._width = height, width

        # which passes to write to the HDF5
        self.write_passes = write_passes
        if isinstance(self.write_passes, str):
            self.write_passes = self.write_passes.split(',')
        self.write_passes = [p for p in self.write_passes if (p in PASSES)]

        # which passes to save as an MP4
        self.save_passes = save_passes
        if isinstance(self.save_passes, str):
            self.save_passes = self.save_passes.split(',')
        self.save_passes = [p for p in self.save_passes if (p in self.write_passes)]
        self.save_movies = save_movies

        print("write passes", self.write_passes)
        print("save passes", self.save_passes)
        print("save movies", self.save_movies)
        if self.save_movies:
            assert len(self.save_passes),\
                "You need to pass \'--save_passes [PASSES]\' to save out movies, where [PASSES] is a comma-separated list of items from %s" % PASSES
        
        # whether to save a JSON of trial-level labels
        self.save_labels = save_labels
        if self.save_labels:
            self.meta_file = Path(output_dir).joinpath('metadata.json')
            if self.meta_file.exists():
                self.trial_metadata = json.loads(self.meta_file.read_text())
            else:
                self.trial_metadata = []

        initialization_commands = self.get_initialization_commands(width=width, height=height)

        # Initialize the scene.
        self.communicate(initialization_commands)

        # Run trials
        self.trial_loop(num, output_dir, temp_path)

        # Terminate TDW
        end = self.communicate({"$type": "terminate"})
        end_types = [OutputData.get_data_type_id(r) for r in end]
        end_count = 0
        print("TOLD BUILD TO TERMINATE, TRY %d" % end_count)
        print(end_types)
        while ('tre\x04' in end_types) and end_count <= 10:
            end_count += 1
            print([OutputData.get_data_type_id(r) for r in end])
            print("TOLD BUILD TO TERMINATE, TRY %d" % end_count)
            end = self.communicate({"$type": "terminate"})

        # Save the command line args
        if self.save_args:
            self.args_dict = copy.deepcopy(args_dict)
        self.save_command_line_args(output_dir)

        # Save the across-trial stats
        if self.save_labels:
            hdf5_paths = glob.glob(str(output_dir) + '/*.hdf5')
            stats = get_across_trial_stats_from(
                hdf5_paths, funcs=self.get_controller_label_funcs(classname=type(self).__name__))
            stats["num_trials"] = int(len(hdf5_paths))
            stats_str = json.dumps(stats, indent=4)
            stats_file = Path(output_dir).joinpath('trial_stats.json')
            stats_file.write_text(stats_str, encoding='utf-8')
            print("ACROSS TRIAL STATS")
            print(stats_str)
Beispiel #25
0
    def trial(self,
              filepath: Path,
              temp_path: Path,
              trial_num: int) -> None:
        """
        Run a trial. Write static and per-frame data to disk until the trial is done.

        :param filepath: The path to this trial's hdf5 file.
        :param temp_path: The path to the temporary file.
        :param trial_num: The number of the current trial.
        """

        # Clear the object IDs and other static data
        self.clear_static_data()
        self._trial_num = trial_num

        # Create the .hdf5 file.
        f = h5py.File(str(temp_path.resolve()), "a")

        commands = []
        # Remove asset bundles (to prevent a memory leak).
        if trial_num % 100 == 0:
            commands.append({"$type": "unload_asset_bundles"})

        # Add commands to start the trial.
        commands.extend(self.get_trial_initialization_commands())
        # Add commands to request output data.
        commands.extend(self._get_send_data_commands())

        # Send the commands and start the trial.
        r_types = ['']
        count = 0
        while 'imag' not in r_types:
            resp = self.communicate(commands)
            r_types = [OutputData.get_data_type_id(r) for r in resp]
            count += 1
            print("INITIAL RESPONSE FROM BUILD ON TRY %d" % count)
            print(r_types)

        self._set_segmentation_colors(resp)
        frame = 0

        # Write static data to disk.
        static_group = f.create_group("static")
        self._write_static_data(static_group)

        # Add the first frame.
        done = False
        frames_grp = f.create_group("frames")
        frame_grp, _, _, _ = self._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame)
        self._write_frame_labels(frame_grp, resp, -1, False)

        # Continue the trial. Send commands, and parse output data.
        while not done:
            frame += 1
            # print('frame %d' % frame)
            resp = self.communicate(self.get_per_frame_commands(resp, frame))
            r_ids = [OutputData.get_data_type_id(r) for r in resp[:-1]]

            # Sometimes the build freezes and has to reopen the socket.
            # This prevents such errors from throwing off the frame numbering
            if ('imag' not in r_ids) or ('tran' not in r_ids):
                print("retrying frame %d, response only had %s" % (frame, r_ids))
                frame -= 1
                continue

            frame_grp, objs_grp, tr_dict, done = self._write_frame(frames_grp=frames_grp, resp=resp, frame_num=frame)

            # Write whether this frame completed the trial and any other trial-level data
            labels_grp, _, _, done = self._write_frame_labels(frame_grp, resp, frame, done)

        # Cleanup.
        commands = []
        for o_id in self.object_ids:
            commands.append({"$type": self._get_destroy_object_command_name(o_id),
                             "id": int(o_id)})
        self.communicate(commands)

        # Compute the trial-level metadata. Save it per trial in case of failure mid-trial loop
        if self.save_labels:
            meta = OrderedDict()
            meta = get_labels_from(f, label_funcs=self.get_controller_label_funcs(type(self).__name__), res=meta)
            self.trial_metadata.append(meta)

            # Save the trial-level metadata
            json_str =json.dumps(self.trial_metadata, indent=4)
            self.meta_file.write_text(json_str, encoding='utf-8')
            print("TRIAL %d LABELS" % self._trial_num)
            print(json.dumps(self.trial_metadata[-1], indent=4))

        # Close the file.
        f.close()
        # Move the file.
        try:
            temp_path.replace(filepath)
        except OSError:
            shutil.move(temp_path, filepath)
    def trial(self, d_f, s_f, b_f, d_c, s_c, b_c, collision_types):
        """
        Collide a chair with a fridge.

        :param d_f: The dynamic friction of the fridge.
        :param s_f: The static friction of the fridge.
        :param b_f: The bounciness of the fridge.
        :param d_c: The dynamic friction of the chair.
        :param s_c: The static friction of the chair.
        :param b_c: The bounciness of the chair.
        :param collision_types: The types of collisions to listen for.
        """

        print("###\n\nNew Trial\n\n###\n")

        fridge_id = 0
        chair_id = 1

        # Destroy all objects currently in the scene.
        init_commands = [{"$type": "destroy_all_objects"}]
        # Create the avatar.
        init_commands.extend(
            TDWUtils.create_avatar(position={
                "x": 1,
                "y": 2.5,
                "z": 5
            },
                                   look_at=TDWUtils.VECTOR3_ZERO))
        # Add the objects.
        # Set the masses and physic materials.
        # Apply a force to the chair.
        # Receive collision data (note that by setting "stay" to True, you will receive a LOT of data;
        # see "Performance Optimizations" documentation.)
        init_commands.extend([
            self.get_add_object("fridge_large", object_id=fridge_id),
            self.get_add_object("chair_billiani_doll",
                                object_id=chair_id,
                                position={
                                    "x": 4,
                                    "y": 0,
                                    "z": 0
                                }), {
                                    "$type": "set_mass",
                                    "id": fridge_id,
                                    "mass": 40
                                }, {
                                    "$type": "set_mass",
                                    "id": chair_id,
                                    "mass": 20
                                }, {
                                    "$type": "set_physic_material",
                                    "id": fridge_id,
                                    "dynamic_friction": d_f,
                                    "static_friction": s_f,
                                    "bounciness": b_f
                                }, {
                                    "$type": "set_physic_material",
                                    "id": chair_id,
                                    "dynamic_friction": d_c,
                                    "static_friction": s_c,
                                    "bounciness": b_c
                                }, {
                                    "$type": "apply_force_to_object",
                                    "force": {
                                        "x": -200,
                                        "y": 0,
                                        "z": 0
                                    },
                                    "id": chair_id
                                }, {
                                    "$type": "send_collisions",
                                    "collision_types": collision_types,
                                    "enter": True,
                                    "exit": True,
                                    "stay": True
                                }
        ])
        self.communicate(init_commands)

        # Iterate through 500 frames.
        # Every frame, listen for collisions, and parse the output data.
        for i in range(500):
            resp = self.communicate([])
            if len(resp) > 1:
                for r in resp[:-1]:
                    r_id = OutputData.get_data_type_id(r)
                    # There was a collision between two objects.
                    if r_id == "coll":
                        collision = Collision(r)
                        print("Collision between two objects:")
                        print("\tEvent: " + collision.get_state())
                        print("\tCollider: " +
                              str(collision.get_collider_id()))
                        print("\tCollidee: " +
                              str(collision.get_collidee_id()))
                        print("\tRelative velocity: " +
                              str(collision.get_relative_velocity()))
                        print("\tContacts:")
                        for j in range(collision.get_num_contacts()):
                            print(
                                str(collision.get_contact_normal(j)) + "\t" +
                                str(collision.get_contact_point(j)))
                    # There was a collision between an object and the environment.
                    elif r_id == "enco":
                        collision = EnvironmentCollision(r)
                        print(
                            "Collision between an object and the environment:")
                        print("\tEvent: " + collision.get_state())
                        print("\tCollider: " + str(collision.get_object_id()))
                        print("\tContacts:")
                        for j in range(collision.get_num_contacts()):
                            print(
                                str(collision.get_contact_normal(j)) + "\t" +
                                str(collision.get_contact_point(j)))
                    else:
                        raise Exception(r_id)
Beispiel #27
0
    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()
Beispiel #28
0
import json
from pathlib import Path
from tdw.controller import Controller
from tdw.output_data import OutputData

c = Controller(launch_build=True)
p = Path.home().joinpath("Downloads/tdw_commands (1).json")
txt = p.read_text()
for i, line in enumerate(txt.split("\n")):
    try:
        if not line.startswith("["):
            continue
        commands = json.loads(line.split(" trial ")[0])
        if not isinstance(commands, list):
            continue
        resp = c.communicate(commands)
        r_ids = [OutputData.get_data_type_id(r) for r in resp[:-1]]
        assert "imag" in r_ids, (i, line, r_ids)
    except json.JSONDecodeError:
        continue
    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 __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)