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
Esempio n. 2
0
 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([])
    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)