in environments/gym-avd/gym_avd/envs/avd_occ_base_env.py [0:0]
def get_local_maps(self) -> Tuple[np.array, np.array, np.array]:
r"""Generates egocentric crops of the global occupancy map.
Returns:
The occupancy images display free, occupied and unknown space.
The color conventions are:
free-space - (0, 255, 0)
occupied-space - (0, 0, 255)
unknown-space - (255, 255, 255)
The outputs are:
fine_ego_map_color - (H, W, 3) occupancy image
coarse_ego_map_color - (H, W, 3) occupancy image
highres_coarse_ego_map_color - (H, W, 3) occupancy image
"""
# ================ The global occupancy map ===========================
top_down_map = self.grids_mat.copy() # (map_size, map_size)
# =========== Obtain local crop around the agent ======================
# Agent's world and map positions.
xzt_world, xz_map = self.get_camera_grid_pos()
# Write out the relevant parts of top_down_map to ego_map.
# A central egocentric crop of the map will be obtained.
self.ego_map.fill(0)
half_size = (
max(
top_down_map.shape[0],
top_down_map.shape[1],
self.configs["large_map_range"],
)
* 3
)
x_start = int(half_size - xz_map[0])
y_start = int(half_size - xz_map[1])
x_end = x_start + top_down_map.shape[0]
y_end = y_start + top_down_map.shape[1]
assert (
x_start >= 0
and y_start >= 0
and x_end <= self.ego_map.shape[0]
and y_end <= self.ego_map.shape[1]
)
self.ego_map[x_start:x_end, y_start:y_end] = top_down_map
# Crop out only the essential parts of the global map.
# This saves computation cost for the subsequent operations.
large_map_range = self.configs["large_map_range"]
crop_start = half_size - int(1.5 * large_map_range)
crop_end = half_size + int(1.5 * large_map_range)
ego_map = self.ego_map[crop_start:crop_end, crop_start:crop_end]
# Rotate the global map to obtain egocentric top-down view.
half_size = ego_map.shape[0] // 2
center = (half_size, half_size)
rot_angle = math.degrees(xzt_world[2]) + 90
M = cv2.getRotationMatrix2D(center, rot_angle, 1.0)
ego_map = cv2.warpAffine(
ego_map,
M,
(ego_map.shape[1], ego_map.shape[0]),
flags=cv2.INTER_NEAREST,
borderMode=cv2.BORDER_CONSTANT,
borderValue=(255,),
)
# ============ Obtain final maps at different resolutions =============
# Obtain the fine occupancy map.
start = int(half_size - self.configs["small_map_range"])
end = int(half_size + self.configs["small_map_range"])
fine_ego_map = ego_map[start:end, start:end]
assert start >= 0
assert end <= ego_map.shape[0]
fine_ego_map = cv2.resize(
fine_ego_map,
(self.configs["small_map_size"], self.configs["small_map_size"]),
interpolation=cv2.INTER_NEAREST,
)
fine_ego_map = np.clip(fine_ego_map, 0, 2)
# Obtain the coarse occupancy map.
start = half_size - self.configs["large_map_range"]
end = half_size + self.configs["large_map_range"]
assert start >= 0
assert end <= ego_map.shape[0]
coarse_ego_map_orig = ego_map[start:end, start:end]
coarse_ego_map = cv2.resize(
coarse_ego_map_orig,
(self.configs["large_map_size"], self.configs["large_map_size"]),
interpolation=cv2.INTER_NEAREST,
)
# Obtain a high-resolution coarse occupancy map.
# This is primarily useful as an input to an A* path-planner.
highres_coarse_ego_map = cv2.resize(
coarse_ego_map_orig, (200, 200), interpolation=cv2.INTER_NEAREST,
)
# Convert to RGB maps.
# Fine occupancy map.
map_shape = (*fine_ego_map.shape, 3)
fine_ego_map_color = np.zeros(map_shape, dtype=np.uint8)
fine_ego_map_color[fine_ego_map == 0] = np.array([255, 255, 255])
fine_ego_map_color[fine_ego_map == 1] = np.array([0, 0, 255])
fine_ego_map_color[fine_ego_map == 2] = np.array([0, 255, 0])
# Coarse occupancy map.
map_shape = (*coarse_ego_map.shape, 3)
coarse_ego_map_color = np.zeros(map_shape, dtype=np.uint8)
coarse_ego_map_color[coarse_ego_map == 0] = np.array([255, 255, 255])
coarse_ego_map_color[coarse_ego_map == 1] = np.array([0, 0, 255])
coarse_ego_map_color[coarse_ego_map == 2] = np.array([0, 255, 0])
# High-resolution coarse occupancy map.
map_shape = (*highres_coarse_ego_map.shape, 3)
highres_coarse_ego_map_color = np.zeros(map_shape, dtype=np.uint8)
highres_coarse_ego_map_color[highres_coarse_ego_map == 0] = np.array(
[255, 255, 255]
)
highres_coarse_ego_map_color[highres_coarse_ego_map == 1] = np.array(
[0, 0, 255]
)
highres_coarse_ego_map_color[highres_coarse_ego_map == 2] = np.array(
[0, 255, 0]
)
return fine_ego_map_color, coarse_ego_map_color, highres_coarse_ego_map_color