From d82de0d9a350e49021ea875599116fa5db6b67c0 Mon Sep 17 00:00:00 2001 From: Quanyi Li Date: Wed, 20 Sep 2023 15:39:10 +0100 Subject: [PATCH] Top down semantic map (#499) * add script * add scaling * import gfxdraw * lane size * format * add test * add top_down_width * draw top down right size * format * fix bug * fix bug --- .../traffic_participants/pedestrian.py | 4 +- metadrive/component/vehicle/base_vehicle.py | 4 +- metadrive/constants.py | 30 ++++++ metadrive/engine/core/manual_controller.py | 2 +- metadrive/envs/base_env.py | 4 + metadrive/envs/marl_envs/marl_intersection.py | 2 +- metadrive/envs/scenario_env.py | 1 + metadrive/manager/scenario_traffic_manager.py | 19 ++-- metadrive/obs/top_down_obs.py | 12 +-- metadrive/obs/top_down_obs_impl.py | 30 ++++-- metadrive/obs/top_down_obs_multi_channel.py | 10 +- metadrive/obs/top_down_renderer.py | 91 +++++++++++++------ .../test_top_down_semantics.py | 56 ++++++++++++ metadrive/utils/draw_top_down_map.py | 2 +- metadrive/utils/utils.py | 3 +- 15 files changed, 205 insertions(+), 65 deletions(-) create mode 100644 metadrive/tests/test_functionality/test_top_down_semantics.py diff --git a/metadrive/component/traffic_participants/pedestrian.py b/metadrive/component/traffic_participants/pedestrian.py index f09ddcde4..2291c0630 100644 --- a/metadrive/component/traffic_participants/pedestrian.py +++ b/metadrive/component/traffic_participants/pedestrian.py @@ -110,8 +110,8 @@ def WIDTH(self): @property def top_down_width(self): - return self.RADIUS + return self.RADIUS * 2 @property def top_down_length(self): - return self.RADIUS + return self.RADIUS * 2 diff --git a/metadrive/component/vehicle/base_vehicle.py b/metadrive/component/vehicle/base_vehicle.py index e90327c1d..f53515d08 100644 --- a/metadrive/component/vehicle/base_vehicle.py +++ b/metadrive/component/vehicle/base_vehicle.py @@ -922,11 +922,11 @@ def max_speed_m_s(self): @property def top_down_length(self): - return self.LENGTH + return self.config["top_down_length"] if self.config["top_down_length"] else self.LENGTH @property def top_down_width(self): - return self.WIDTH + return self.config["top_down_width"] if self.config["top_down_width"] else self.WIDTH @property def lane(self): diff --git a/metadrive/constants.py b/metadrive/constants.py index ebef04660..bb90a81ed 100644 --- a/metadrive/constants.py +++ b/metadrive/constants.py @@ -2,6 +2,7 @@ from collections import namedtuple from typing import List, Tuple +import numpy as np from panda3d.bullet import BulletWorld from panda3d.core import Vec3 from panda3d.core import Vec4, BitMask32 @@ -381,3 +382,32 @@ def get_color(type): return (0, 0, 0, 1) else: raise ValueError("Unsupported type: {}".format(type)) + + +class TopDownSemanticColor: + """ + Do not modify this as it is for terrain generation. If you want your own palette, just add a new one or modify + class lMapSemanticColor + """ + @staticmethod + def get_color(type, pygame=False): + if MetaDriveType.is_yellow_line(type): + # return (255, 0, 0, 0) + ret = np.array([1, 0, 0, 0]) + elif MetaDriveType.is_lane(type): + ret = np.array([0, 1, 0, 0]) + elif type == MetaDriveType.GROUND: + ret = np.array([0, 0, 1, 0]) + elif MetaDriveType.is_white_line(type) or MetaDriveType.is_road_edge(type): + ret = np.array([0, 0, 0, 1]) + elif MetaDriveType.is_vehicle(type): + ret = np.array([0, 0, 1, 0]) + elif type == MetaDriveType.OTHER: + ret = np.array([0, 1, 1, 0]) + else: + raise ValueError("Unsupported type: {}".format(type)) + + if pygame: + ret *= 255 + ret.astype(np.int32) + return ret diff --git a/metadrive/engine/core/manual_controller.py b/metadrive/engine/core/manual_controller.py index 6b212e5a4..4367ae739 100644 --- a/metadrive/engine/core/manual_controller.py +++ b/metadrive/engine/core/manual_controller.py @@ -13,7 +13,7 @@ from metadrive.utils import import_pygame -pygame = import_pygame() +pygame, gfxdraw = import_pygame() class Controller: diff --git a/metadrive/envs/base_env.py b/metadrive/envs/base_env.py index ba1d023ba..d4ecd4e5e 100644 --- a/metadrive/envs/base_env.py +++ b/metadrive/envs/base_env.py @@ -115,6 +115,10 @@ height=None, mass=None, + # only for top-down drawing, the physics size wouldn't be changed + top_down_width=None, + top_down_length=None, + # ===== vehicle module config ===== lidar=dict( num_lasers=240, distance=50, num_others=0, gaussian_noise=0.0, dropout_prob=0.0, add_others_navi=False diff --git a/metadrive/envs/marl_envs/marl_intersection.py b/metadrive/envs/marl_envs/marl_intersection.py index 2a5062834..39b4b8139 100644 --- a/metadrive/envs/marl_envs/marl_intersection.py +++ b/metadrive/envs/marl_envs/marl_intersection.py @@ -359,7 +359,7 @@ def show_map_and_traj(): env.reset() with open("metasvodist_inter_best.json", "r") as f: traj = json.load(f) - m = draw_top_down_map(env.current_map, draw_drivable_area=False, return_surface=True, reverse_color=True) + m = draw_top_down_map(env.current_map, semantic_map=False, return_surface=True, reverse_color=True) m = draw_top_down_trajectory( m, traj, entry_differ_color=True, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)] ) diff --git a/metadrive/envs/scenario_env.py b/metadrive/envs/scenario_env.py index ef3ae541a..84ed9ef50 100644 --- a/metadrive/envs/scenario_env.py +++ b/metadrive/envs/scenario_env.py @@ -91,6 +91,7 @@ interface_panel=["dashboard"], # for boosting efficiency horizon=None, allowed_more_steps=None, # None=infinite + top_down_show_real_size=False ) diff --git a/metadrive/manager/scenario_traffic_manager.py b/metadrive/manager/scenario_traffic_manager.py index f123e9e7f..5977809e0 100644 --- a/metadrive/manager/scenario_traffic_manager.py +++ b/metadrive/manager/scenario_traffic_manager.py @@ -1,5 +1,5 @@ import copy -import logging +from metadrive.engine.logger import get_logger from metadrive.utils.math import norm import numpy as np @@ -17,7 +17,7 @@ from metadrive.type import MetaDriveType from metadrive.utils.math import wrap_to_pi -logger = logging.getLogger(__name__) +logger = get_logger() class ScenarioTrafficManager(BaseManager): @@ -192,12 +192,17 @@ def spawn_vehicle(self, v_id, track): float(state["length"]), None if self.even_sample_v else self.np_random, self.need_default_vehicle ) obj_name = v_id if self.engine.global_config["force_reuse_object_name"] else None + v_cfg = copy.copy(self._traffic_v_config) + if self.engine.global_config["top_down_show_real_size"]: + v_cfg["top_down_length"] = track["state"]["length"][self.episode_step] + v_cfg["top_down_width"] = track["state"]["width"][self.episode_step] + if v_cfg["top_down_length"] < 1 or v_cfg["top_down_width"] < 0.5: + logger.warning( + "Scenario ID: {}. The top_down size of vehicle {} is weird: " + "{}".format(self.engine.current_seed, v_id, [v_cfg["length"], v_cfg["width"]]) + ) v = self.spawn_object( - vehicle_class, - position=state["position"], - heading=state["heading"], - vehicle_config=self._traffic_v_config, - name=obj_name + vehicle_class, position=state["position"], heading=state["heading"], vehicle_config=v_cfg, name=obj_name ) self._scenario_id_to_obj_id[v_id] = v.name self._obj_id_to_scenario_id[v.name] = v_id diff --git a/metadrive/obs/top_down_obs.py b/metadrive/obs/top_down_obs.py index 5c073d940..1ac539d50 100644 --- a/metadrive/obs/top_down_obs.py +++ b/metadrive/obs/top_down_obs.py @@ -13,10 +13,10 @@ from metadrive.constants import Decoration, DEFAULT_AGENT, EDITION from metadrive.obs.observation_base import ObservationBase from metadrive.obs.top_down_obs_impl import WorldSurface, ObservationWindow, COLOR_BLACK, \ - VehicleGraphics, LaneGraphics + ObjectGraphics, LaneGraphics from metadrive.utils import import_pygame -pygame = import_pygame() +pygame, gfxdraw = import_pygame() class TopDownObservation(ObservationBase): @@ -38,7 +38,7 @@ def __init__(self, vehicle_config, clip_rgb: bool, onscreen, resolution=None, ma # self.obs_shape = (64, 64) self.obs_shape = self.resolution - self.pygame = import_pygame() + self.pygame, _ = import_pygame() self.onscreen = onscreen main_window_position = (0, 0) @@ -160,15 +160,15 @@ def draw_scene(self): ego_heading = vehicle.heading_theta ego_heading = ego_heading if abs(ego_heading) > 2 * np.pi / 180 else 0 - VehicleGraphics.display( - vehicle=vehicle, surface=self.canvas_runtime, heading=ego_heading, color=VehicleGraphics.GREEN + ObjectGraphics.display( + object=vehicle, surface=self.canvas_runtime, heading=ego_heading, color=ObjectGraphics.GREEN ) for v in self.engine.traffic_manager.vehicles: if v is vehicle: continue h = v.heading_theta h = h if abs(h) > 2 * np.pi / 180 else 0 - VehicleGraphics.display(vehicle=v, surface=self.canvas_runtime, heading=h, color=VehicleGraphics.BLUE) + ObjectGraphics.display(object=v, surface=self.canvas_runtime, heading=h, color=ObjectGraphics.BLUE) # Prepare a runtime canvas for rotation return self.obs_window.render(canvas=self.canvas_runtime, position=pos, heading=vehicle.heading_theta) diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py index 334ea14b0..2a13a2ac1 100644 --- a/metadrive/obs/top_down_obs_impl.py +++ b/metadrive/obs/top_down_obs_impl.py @@ -7,10 +7,13 @@ from metadrive.constants import PGLineType, PGLineColor from metadrive.utils.utils import import_pygame from metadrive.type import MetaDriveType +from metadrive.constants import DrivableAreaProperty +from collections import namedtuple PositionType = Union[Tuple[float, float], np.ndarray] -pygame = import_pygame() +pygame, gfxdraw = import_pygame() COLOR_BLACK = pygame.Color("black") +history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done type") class ObservationWindow: @@ -181,7 +184,7 @@ def copy(self): return ret -class VehicleGraphics: +class ObjectGraphics: RED = (255, 100, 100) GREEN = (50, 200, 0) BLUE = (100, 200, 255) @@ -194,22 +197,29 @@ class VehicleGraphics: @classmethod def display( - cls, vehicle, surface, color, heading, label: bool = False, draw_countour=False, contour_width=1 + cls, + object: history_object, + surface, + color, + heading, + label: bool = False, + draw_countour=False, + contour_width=1 ) -> None: """ Display a vehicle on a pygame surface. The vehicle is represented as a colored rotated rectangle. - :param vehicle: the vehicle to be drawn + :param object: the vehicle to be drawn :param surface: the surface to draw the vehicle on :param label: whether a text label should be rendered """ - if not surface.is_visible(vehicle.position): + if not surface.is_visible(object.position): return - w = surface.pix(vehicle.WIDTH) - h = surface.pix(vehicle.LENGTH) - position = [*surface.pos2pix(vehicle.position[0], vehicle.position[1])] + w = surface.pix(object.WIDTH) + h = surface.pix(object.LENGTH) + position = [*surface.pos2pix(object.position[0], object.position[1])] angle = np.rad2deg(heading) box = [pygame.math.Vector2(p) for p in [(-h / 2, -w / 2), (-h / 2, w / 2), (h / 2, w / 2), (h / 2, -w / 2)]] box_rotate = [p.rotate(angle) + position for p in box] @@ -222,7 +232,7 @@ def display( if label: if cls.font is None: cls.font = pygame.font.Font(None, 15) - text = "#{}".format(id(vehicle) % 1000) + text = "#{}".format(id(object) % 1000) text = cls.font.render(text, 1, (10, 10, 10), (255, 255, 255)) surface.blit(text, position) @@ -419,7 +429,7 @@ def draw_stripes(cls, lane, surface, starts: List[float], ends: List[float], lat pygame.draw.line( surface, color, (surface.vec2pix(lane.position(starts[k], lats[k]))), (surface.vec2pix(lane.position(ends[k], lats[k]))), - max(surface.pix(cls.STRIPE_WIDTH), surface.pix(cls.LANE_LINE_WIDTH)) + surface.pix(2 * DrivableAreaProperty.LANE_LINE_WIDTH) ) @classmethod diff --git a/metadrive/obs/top_down_obs_multi_channel.py b/metadrive/obs/top_down_obs_multi_channel.py index 2bbd209e9..11d7919c0 100644 --- a/metadrive/obs/top_down_obs_multi_channel.py +++ b/metadrive/obs/top_down_obs_multi_channel.py @@ -7,11 +7,11 @@ from metadrive.component.vehicle.base_vehicle import BaseVehicle from metadrive.constants import Decoration, DEFAULT_AGENT from metadrive.obs.top_down_obs import TopDownObservation -from metadrive.obs.top_down_obs_impl import WorldSurface, COLOR_BLACK, VehicleGraphics, LaneGraphics, \ +from metadrive.obs.top_down_obs_impl import WorldSurface, COLOR_BLACK, ObjectGraphics, LaneGraphics, \ ObservationWindowMultiChannel from metadrive.utils import import_pygame, clip -pygame = import_pygame() +pygame, gfxdraw = import_pygame() COLOR_WHITE = pygame.Color("white") @@ -147,7 +147,7 @@ def draw_scene(self): continue h = v.heading_theta h = h if abs(h) > 2 * np.pi / 180 else 0 - VehicleGraphics.display(vehicle=v, surface=self.canvas_runtime, heading=h, color=VehicleGraphics.BLUE) + ObjectGraphics.display(object=v, surface=self.canvas_runtime, heading=h, color=ObjectGraphics.BLUE) raw_pos = vehicle.position self.stack_past_pos.append(raw_pos) @@ -184,8 +184,8 @@ def draw_scene(self): def _draw_ego_vehicle(self): vehicle = self.engine.agents[DEFAULT_AGENT] - w = vehicle.WIDTH * self.scaling - h = vehicle.LENGTH * self.scaling + w = vehicle.top_down_width * self.scaling + h = vehicle.top_down_length * self.scaling position = (self.resolution[0] / 2, self.resolution[1] / 2) angle = 90 box = [pygame.math.Vector2(p) for p in [(-h / 2, -w / 2), (-h / 2, w / 2), (h / 2, w / 2), (h / 2, -w / 2)]] diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py index 7a0ab57a6..b3e4c09c6 100644 --- a/metadrive/obs/top_down_renderer.py +++ b/metadrive/obs/top_down_renderer.py @@ -1,5 +1,7 @@ import copy -from collections import deque, namedtuple +from metadrive.component.vehicle.base_vehicle import BaseVehicle +import math +from collections import deque from typing import Optional, Union, Iterable import numpy as np @@ -7,24 +9,25 @@ from metadrive.component.map.nuplan_map import NuPlanMap from metadrive.component.map.scenario_map import ScenarioMap from metadrive.constants import Decoration, TARGET_VEHICLES -from metadrive.obs.top_down_obs_impl import WorldSurface, VehicleGraphics, LaneGraphics +from metadrive.constants import TopDownSemanticColor, MetaDriveType, DrivableAreaProperty +from metadrive.obs.top_down_obs_impl import WorldSurface, ObjectGraphics, LaneGraphics, history_object from metadrive.scenario.scenario_description import ScenarioDescription from metadrive.utils.interpolating_line import InterpolatingLine from metadrive.utils.utils import import_pygame from metadrive.utils.utils import is_map_related_instance -pygame = import_pygame() +pygame, gfxdraw = import_pygame() color_white = (255, 255, 255) -history_object = namedtuple("history_object", "name position heading_theta WIDTH LENGTH color done") def draw_top_down_map( map, resolution: Iterable = (512, 512), - draw_drivable_area=True, + semantic_map=True, return_surface=False, film_size=None, + scaling=None, reverse_color=False, road_color=color_white, ) -> Optional[Union[np.ndarray, pygame.Surface]]: @@ -39,15 +42,40 @@ def draw_top_down_map( y_len = b_box[3] - b_box[2] max_len = max(x_len, y_len) # scaling and center can be easily found by bounding box - scaling = film_size[1] / max_len - 0.1 + scaling = scaling if scaling else film_size[1] / max_len - 0.1 surface.scaling = scaling centering_pos = ((b_box[0] + b_box[1]) / 2, (b_box[2] + b_box[3]) / 2) surface.move_display_window_to(centering_pos) if isinstance(map, ScenarioMap): - if draw_drivable_area: - for lane_info in map.road_network.graph.values(): - LaneGraphics.draw_drivable_area(lane_info.lane, surface, color=road_color) + if semantic_map: + line_sample_interval = 2 + all_lanes = map.get_map_features(line_sample_interval) + + for obj in all_lanes.values(): + if MetaDriveType.is_lane(obj["type"]): + pygame.draw.polygon( + surface, TopDownSemanticColor.get_color(obj["type"], True), + [surface.pos2pix(p[0], p[1]) for p in obj["polygon"]] + ) + + elif (MetaDriveType.is_road_line(obj["type"]) or MetaDriveType.is_sidewalk(obj["type"])): + if MetaDriveType.is_broken_line(obj["type"]): + points_to_skip = math.floor(DrivableAreaProperty.STRIPE_LENGTH * 2 / line_sample_interval) * 2 + else: + points_to_skip = 1 + for index in range(0, len(obj["polyline"]) - 1, points_to_skip): + if index + points_to_skip < len(obj["polyline"]): + s_p = obj["polyline"][index] + e_p = obj["polyline"][index + points_to_skip] + pygame.draw.line( + surface, + TopDownSemanticColor.get_color(obj["type"], True), + surface.vec2pix([s_p[0], s_p[1]]), + surface.vec2pix([e_p[0], e_p[1]]), + # max(surface.pix(LaneGraphics.STRIPE_WIDTH), + surface.pix(DrivableAreaProperty.LANE_LINE_WIDTH) * 2 + ) else: for id, data in map.blocks[-1].map_data.items(): if ScenarioDescription.POLYLINE not in data: @@ -70,7 +98,7 @@ def draw_top_down_map( LaneGraphics.display_scenario(waymo_line, type, surface) elif isinstance(map, NuPlanMap): - if draw_drivable_area: + if semantic_map: for lane_info in map.road_network.graph.values(): LaneGraphics.draw_drivable_area(lane_info.lane, surface, color=road_color) else: @@ -84,7 +112,7 @@ def draw_top_down_map( decoration = True if _from == Decoration.start else False for _to in map.road_network.graph[_from].keys(): for l in map.road_network.graph[_from][_to]: - if draw_drivable_area: + if semantic_map: LaneGraphics.draw_drivable_area(l, surface, color=road_color) else: two_side = True if l is map.road_network.graph[_from][_to][-1] or decoration else False @@ -166,6 +194,8 @@ def __init__( camera_position=None, target_vehicle_heading_up=False, draw_target_vehicle_trajectory=False, + semantic_map=False, + scaling=None, # auto-scale **kwargs # current_track_vehicle=None ): @@ -194,12 +224,15 @@ def __init__( self._text_render_pos = [50, 50] self._font_size = 25 self._text_render_interval = 20 + self.semantic_map = semantic_map + self.scaling = scaling # Setup the canvas # (1) background is the underlying layer. It is fixed and will never change unless the map changes. self._background_canvas = draw_top_down_map( self.map, - draw_drivable_area=False, + scaling=self.scaling, + semantic_map=self.semantic_map, return_surface=True, film_size=film_size, road_color=road_color, @@ -267,6 +300,7 @@ def render(self, text, *args, **kwargs): if self.draw_target_vehicle_trajectory: self.history_target_vehicle.append( history_object( + type=MetaDriveType.VEHICLE, name=self.current_track_vehicle.name, heading_theta=self.current_track_vehicle.heading_theta, WIDTH=self.current_track_vehicle.top_down_width, @@ -313,7 +347,8 @@ def reset(self, map): # Reset the super large background self._background_canvas = draw_top_down_map( map, - draw_drivable_area=False, + scaling=self.scaling, + semantic_map=self.semantic_map, return_surface=True, film_size=self._background_size, road_color=self.road_color, @@ -347,6 +382,7 @@ def _append_frame_objects(self, objects): frame_objects.append( history_object( name=name, + type=MetaDriveType.VEHICLE if isinstance(obj, BaseVehicle) else MetaDriveType.OTHER, heading_theta=obj.heading_theta, WIDTH=obj.top_down_width, LENGTH=obj.top_down_length, @@ -374,13 +410,11 @@ def _draw(self, *args, **kwargs): h = h if abs(h) > 2 * np.pi / 180 else 0 x = abs(int(i)) alpha_f = x / len(self.history_objects) - VehicleGraphics.display( - vehicle=v, - surface=self._runtime_canvas, - heading=h, - color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])), - draw_countour=False - ) + if self.semantic_map: + c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255 + else: + c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])) + ObjectGraphics.display(object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=False) # Draw the whole trajectory of ego vehicle with no gradient colors: if self.draw_target_vehicle_trajectory: @@ -394,8 +428,8 @@ def _draw(self, *args, **kwargs): x = abs(int(i)) alpha_f = min(x / len(self.history_target_vehicle), 0.5) # alpha_f = 0 - VehicleGraphics.display( - vehicle=v, + ObjectGraphics.display( + object=v, surface=self._runtime_canvas, heading=h, color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])), @@ -412,13 +446,12 @@ def _draw(self, *args, **kwargs): c = v.color h = h if abs(h) > 2 * np.pi / 180 else 0 alpha_f = 0 - VehicleGraphics.display( - vehicle=v, - surface=self._runtime_canvas, - heading=h, - color=(c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])), - draw_countour=True, - contour_width=2 + if self.semantic_map: + c = TopDownSemanticColor.get_color(v.type, True) * (1 - alpha_f) + alpha_f * 255 + else: + c = (c[0] + alpha_f * (255 - c[0]), c[1] + alpha_f * (255 - c[1]), c[2] + alpha_f * (255 - c[2])) + ObjectGraphics.display( + object=v, surface=self._runtime_canvas, heading=h, color=c, draw_countour=True, contour_width=2 ) if not hasattr(self, "_deads"): diff --git a/metadrive/tests/test_functionality/test_top_down_semantics.py b/metadrive/tests/test_functionality/test_top_down_semantics.py new file mode 100644 index 000000000..302565650 --- /dev/null +++ b/metadrive/tests/test_functionality/test_top_down_semantics.py @@ -0,0 +1,56 @@ +""" +This script demonstrates how to use the environment where traffic and road map are loaded from Waymo dataset. +""" +import argparse +import random + +import pygame + +from metadrive.constants import HELP_MESSAGE +from metadrive.engine.asset_loader import AssetLoader +from metadrive.envs.real_data_envs.waymo_env import WaymoEnv + + +class DemoWaymoEnv(WaymoEnv): + def reset(self, seed=None): + if self.engine is not None: + seeds = [i for i in range(self.config["num_scenarios"])] + seeds.remove(self.current_seed) + seed = random.choice(seeds) + return super(DemoWaymoEnv, self).reset(seed=seed) + + +def test_top_down_semantics(render=False): + asset_path = AssetLoader.asset_path + try: + env = DemoWaymoEnv( + { + "manual_control": True, + "reactive_traffic": False, + "use_render": False, + "data_directory": AssetLoader.file_path(asset_path, "waymo", return_raw_style=False), + "num_scenarios": 3 + } + ) + o, _ = env.reset() + + for i in range(1, 1000): + o, r, tm, tc, info = env.step([1.0, 0.]) + if render: + this_frame_fig = env.render( + mode="top_down", + semantic_map=True, + film_size=(2000, 2000), + num_stack=1, + scaling=10, + ) + # save + # pygame.image.save(this_frame_fig, "{}.png".format(i)) + if tm or tc: + env.reset() + finally: + env.close() + + +if __name__ == '__main__': + test_top_down_semantics(True) diff --git a/metadrive/utils/draw_top_down_map.py b/metadrive/utils/draw_top_down_map.py index 79fc4d8ad..e71b48ba6 100644 --- a/metadrive/utils/draw_top_down_map.py +++ b/metadrive/utils/draw_top_down_map.py @@ -5,7 +5,7 @@ from metadrive.obs.top_down_renderer import draw_top_down_map as native_draw from metadrive.utils.utils import import_pygame -pygame = import_pygame() +pygame, gfxdraw = import_pygame() def draw_top_down_map(map, resolution: Iterable = (512, 512)) -> Optional[Union[np.ndarray, pygame.Surface]]: diff --git a/metadrive/utils/utils.py b/metadrive/utils/utils.py index d1f3f256f..20192927d 100644 --- a/metadrive/utils/utils.py +++ b/metadrive/utils/utils.py @@ -14,7 +14,8 @@ def import_pygame(): os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide" import pygame - return pygame + from pygame import gfxdraw + return pygame, gfxdraw def get_time_str():