From 05dbf7220e8386eb2337502c4d4851fc8dce30cd Mon Sep 17 00:00:00 2001
From: danielgordon10 <danielgordon10@gmail.com>
Date: Fri, 26 Apr 2019 16:03:27 -0700
Subject: [PATCH] Small changes to the topdown map interface to make plotting
the agent easier (#76)
* added conversion of matrix map to color picture
* added example
---
examples/shortest_path_follower_example.py | 69 ++++++++++++++++++----
habitat/tasks/nav/nav_task.py | 39 ++++++------
habitat/utils/visualizations/maps.py | 29 +++++++--
3 files changed, 103 insertions(+), 34 deletions(-)
diff --git a/examples/shortest_path_follower_example.py b/examples/shortest_path_follower_example.py
index e6cd147d3..b3f8ca7b9 100644
--- a/examples/shortest_path_follower_example.py
+++ b/examples/shortest_path_follower_example.py
@@ -7,28 +7,69 @@
import os
import shutil
-import imageio
-
+import cv2
import habitat
+import numpy as np
from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
+from habitat.utils.visualizations import maps
+from habitat.utils.visualizations.utils import images_to_video
IMAGE_DIR = os.path.join("examples", "images")
if not os.path.exists(IMAGE_DIR):
os.makedirs(IMAGE_DIR)
+class SimpleRLEnv(habitat.RLEnv):
+ def get_reward_range(self):
+ return [-1, 1]
+
+ def get_reward(self, observations):
+ return 0
+
+ def get_done(self, observations):
+ return self.habitat_env.episode_over
+
+ def get_info(self, observations):
+ return self.habitat_env.get_metrics()
+
+
+def draw_top_down_map(info, heading, output_size):
+ top_down_map = maps.colorize_topdown_map(info["top_down_map"]["map"])
+ original_map_size = top_down_map.shape[:2]
+ map_scale = np.array(
+ (1, original_map_size[1] * 1.0 / original_map_size[0])
+ )
+ new_map_size = np.round(output_size * map_scale).astype(np.int32)
+ # OpenCV expects w, h but map size is in h, w
+ top_down_map = cv2.resize(top_down_map, (new_map_size[1], new_map_size[0]))
+
+ map_agent_pos = info["top_down_map"]["agent_map_coord"]
+ map_agent_pos = np.round(
+ map_agent_pos * new_map_size / original_map_size
+ ).astype(np.int32)
+ top_down_map = maps.draw_agent(
+ top_down_map,
+ map_agent_pos,
+ heading - np.pi / 2,
+ agent_radius_px=top_down_map.shape[0] / 40,
+ )
+ return top_down_map
+
+
def shortest_path_example(mode):
config = habitat.get_config(config_file="tasks/pointnav.yaml")
- env = habitat.Env(config=config)
+ config.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
+ config.TASK.SENSORS.append("HEADING_SENSOR")
+ env = SimpleRLEnv(config=config)
goal_radius = env.episodes[0].goals[0].radius
if goal_radius is None:
goal_radius = config.SIMULATOR.FORWARD_STEP_SIZE
- follower = ShortestPathFollower(env.sim, goal_radius, False)
+ follower = ShortestPathFollower(env.habitat_env.sim, goal_radius, False)
follower.mode = mode
print("Environment creation successful")
for episode in range(3):
- observations = env.reset()
+ env.reset()
dirname = os.path.join(
IMAGE_DIR, "shortest_path_example", mode, "%02d" % episode
)
@@ -36,16 +77,20 @@ def shortest_path_example(mode):
shutil.rmtree(dirname)
os.makedirs(dirname)
print("Agent stepping around inside environment.")
- count_steps = 0
- while not env.episode_over:
+ images = []
+ while not env.habitat_env.episode_over:
best_action = follower.get_next_action(
- env.current_episode.goals[0].position
+ env.habitat_env.current_episode.goals[0].position
)
- observations = env.step(best_action.value)
- count_steps += 1
+ observations, reward, done, info = env.step(best_action.value)
im = observations["rgb"]
- imageio.imsave(os.path.join(dirname, "%03d.jpg" % count_steps), im)
- print("Episode finished after {} steps.".format(count_steps))
+ top_down_map = draw_top_down_map(
+ info, observations["heading"], im.shape[0]
+ )
+ output_im = np.concatenate((im, top_down_map), axis=1)
+ images.append(output_im)
+ images_to_video(images, dirname, "trajectory")
+ print("Episode finished")
def main():
diff --git a/habitat/tasks/nav/nav_task.py b/habitat/tasks/nav/nav_task.py
index e1b19731f..95fbda44e 100644
--- a/habitat/tasks/nav/nav_task.py
+++ b/habitat/tasks/nav/nav_task.py
@@ -547,7 +547,13 @@ class TopDownMap(habitat.Measure):
+ self._grid_delta,
]
- self._metric = house_map
+ self._metric = {
+ "map": house_map,
+ "agent_map_coord": (
+ map_agent_x - (self._ind_x_min - self._grid_delta),
+ map_agent_y - (self._ind_y_min - self._grid_delta),
+ ),
+ }
def update_map(self, agent_position):
a_x, a_y = maps.to_grid(
@@ -558,24 +564,21 @@ class TopDownMap(habitat.Measure):
self._map_resolution,
)
# Don't draw over the source point
- color = (
- min(self._step_count * 245 / self._config.MAX_EPISODE_STEPS, 245)
- + 10
- if self._top_down_map[a_x, a_y] != maps.MAP_SOURCE_POINT_INDICATOR
- else maps.MAP_SOURCE_POINT_INDICATOR
- )
- color = int(color)
+ if self._top_down_map[a_x, a_y] != maps.MAP_SOURCE_POINT_INDICATOR:
+ color = 10 + min(
+ self._step_count * 245 // self._config.MAX_EPISODE_STEPS, 245
+ )
- thickness = int(
- np.round(self._map_resolution[0] * 2 / MAP_THICKNESS_SCALAR)
- )
- cv2.line(
- self._top_down_map,
- self._previous_xy_location,
- (a_y, a_x),
- color,
- thickness=thickness,
- )
+ thickness = int(
+ np.round(self._map_resolution[0] * 2 / MAP_THICKNESS_SCALAR)
+ )
+ cv2.line(
+ self._top_down_map,
+ self._previous_xy_location,
+ (a_y, a_x),
+ color,
+ thickness=thickness,
+ )
self._previous_xy_location = (a_y, a_x)
return self._top_down_map, a_x, a_y
diff --git a/habitat/utils/visualizations/maps.py b/habitat/utils/visualizations/maps.py
index 9dca6e248..c44302acd 100644
--- a/habitat/utils/visualizations/maps.py
+++ b/habitat/utils/visualizations/maps.py
@@ -4,12 +4,13 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
-import numpy as np
+import os
+from typing import List, Tuple, Optional
+
import cv2
import imageio
+import numpy as np
import scipy.ndimage
-import os
-from typing import List, Tuple, Optional
from habitat.core.simulator import Simulator
from habitat.utils.visualizations import utils
@@ -32,6 +33,16 @@ MAP_BORDER_INDICATOR = 2
MAP_SOURCE_POINT_INDICATOR = 4
MAP_TARGET_POINT_INDICATOR = 6
+TOP_DOWN_MAP_COLORS = np.full((256, 3), 150, dtype=np.uint8)
+TOP_DOWN_MAP_COLORS[10:] = cv2.applyColorMap(
+ np.arange(246, dtype=np.uint8), cv2.COLORMAP_JET
+).squeeze(1)[:, ::-1]
+TOP_DOWN_MAP_COLORS[MAP_INVALID_POINT] = [255, 255, 255]
+TOP_DOWN_MAP_COLORS[MAP_VALID_POINT] = [150, 150, 150]
+TOP_DOWN_MAP_COLORS[MAP_BORDER_INDICATOR] = [50, 50, 50]
+TOP_DOWN_MAP_COLORS[MAP_SOURCE_POINT_INDICATOR] = [0, 0, 200]
+TOP_DOWN_MAP_COLORS[MAP_TARGET_POINT_INDICATOR] = [200, 0, 0]
+
def draw_agent(
image: np.ndarray,
@@ -51,7 +62,7 @@ def draw_agent(
# Rotate before resize to keep good resolution.
rotated_agent = scipy.ndimage.interpolation.rotate(
- AGENT_SPRITE, agent_rotation * -180 / np.pi
+ AGENT_SPRITE, agent_rotation * 180 / np.pi
)
# Rescale because rotation may result in larger image than original, but
# the agent sprite size should stay the same.
@@ -312,3 +323,13 @@ def get_topdown_map(
top_down_map[range_x[0] : range_x[1], range_y[0] : range_y[1]]
)
return top_down_map
+
+
+def colorize_topdown_map(top_down_map: np.ndarray) -> np.ndarray:
+ """Convert the top down map to RGB based on the indicator values.
+ Args:
+ top_down_map: A non-colored version of the map.
+ Returns:
+ A colored version of the top-down map.
+ """
+ return TOP_DOWN_MAP_COLORS[top_down_map]
--
GitLab