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

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

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):

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

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:

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

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)]
     )

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
 )
 
 

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

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)

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

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)]]