def _get_mesh_occupancy()

in habitat_extensions/sensors.py [0:0]

• 51 lines of code
• 8 McCabe index (conditional complexity)

    def _get_mesh_occupancy(self, episode, agent_state):
        episode_id = (episode.episode_id, episode.scene_id)
        if self.current_episode_id != episode_id:
            self.current_episode_id = episode_id
            # Transpose to make x rightward and y downward
            self._top_down_map = self.get_original_map().T

        agent_position = agent_state.position
        agent_rotation = agent_state.rotation
        a_x, a_y = maps.to_grid(
            agent_position[0],
            agent_position[2],
            self._coordinate_min,
            self._coordinate_max,
            self._map_resolution,
        )

        # Crop region centered around the agent
        mrange = int(self.map_size * 1.5)
        ego_map = self._top_down_map[
            (a_y - mrange) : (a_y + mrange), (a_x - mrange) : (a_x + mrange)
        ]
        if ego_map.shape[0] == 0 or ego_map.shape[1] == 0:
            ego_map = np.zeros((2 * mrange + 1, 2 * mrange + 1), dtype=np.uint8)

        # Rotate to get egocentric map
        # Negative since the value returned is clockwise rotation about Y,
        # but we need anti-clockwise rotation
        agent_heading = -compute_heading_from_quaternion(agent_rotation)
        agent_heading = math.degrees(agent_heading)

        half_size = ego_map.shape[0] // 2
        center = (half_size, half_size)
        M = cv2.getRotationMatrix2D(center, agent_heading, scale=1.0)

        ego_map = (
            cv2.warpAffine(
                ego_map * 255,
                M,
                (ego_map.shape[1], ego_map.shape[0]),
                flags=cv2.INTER_LINEAR,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=(1,),
            ).astype(np.float32)
            / 255.0
        )

        mrange = int(self.map_size)
        ego_map = ego_map[
            (half_size - mrange) : (half_size + mrange),
            (half_size - mrange) : (half_size + mrange),
        ]
        ego_map[ego_map > 0.5] = 1.0
        ego_map[ego_map <= 0.5] = 0.0

        # This map is currently 0 if occupied and 1 if unoccupied. Flip it.
        ego_map = 1.0 - ego_map

        # Flip the x axis because to_grid() flips the conventions
        ego_map = np.flip(ego_map, axis=1)

        # Get forward region infront of the agent
        half_size = ego_map.shape[0] // 2
        quarter_size = ego_map.shape[0] // 4
        center = (half_size, half_size)

        ego_map = ego_map[0:half_size, quarter_size : (quarter_size + half_size)]

        # Append explored status in the 2nd channel
        ego_map = np.stack([ego_map, np.ones_like(ego_map)], axis=2)

        return ego_map