import gfxdraw

metadriverse · Sep 18, 2023 · b572cac · b572cac
1 parent aa1a2f6
commit b572cac
Show file tree

Hide file tree

Showing 8 changed files with 10 additions and 8 deletions.
diff --git a/metadrive/constants.py b/metadrive/constants.py
@@ -411,4 +411,5 @@ def get_color(type, pygame=False):
 
         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
@@ -13,7 +13,7 @@
 
 from metadrive.utils import import_pygame
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 
 
 class Controller:

diff --git a/metadrive/obs/top_down_obs.py b/metadrive/obs/top_down_obs.py
@@ -16,7 +16,7 @@
     VehicleGraphics, 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)

diff --git a/metadrive/obs/top_down_obs_impl.py b/metadrive/obs/top_down_obs_impl.py
@@ -9,7 +9,7 @@
 from metadrive.type import MetaDriveType
 
 PositionType = Union[Tuple[float, float], np.ndarray]
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 COLOR_BLACK = pygame.Color("black")
 
 

diff --git a/metadrive/obs/top_down_obs_multi_channel.py b/metadrive/obs/top_down_obs_multi_channel.py
@@ -11,7 +11,7 @@
     ObservationWindowMultiChannel
 from metadrive.utils import import_pygame, clip
 
-pygame = import_pygame()
+pygame, gfxdraw = import_pygame()
 COLOR_WHITE = pygame.Color("white")
 
 

diff --git a/metadrive/obs/top_down_renderer.py b/metadrive/obs/top_down_renderer.py
@@ -16,7 +16,7 @@
 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 type")

diff --git 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
@@ -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():