added stopsign

pygpudrive/visualize/core.py

@@ -252,6 +252,8 @@ def _plot_roadgraph(
                     or road_point_type == int(gpudrive.EntityType.RoadLine)
                     or road_point_type == int(gpudrive.EntityType.RoadLane)
                     or road_point_type == int(gpudrive.EntityType.SpeedBump)
+                    or road_point_type == int(gpudrive.EntityType.StopSign)
+
                 ):
                     # Get coordinates and metadata
                     x_coords = road_graph.x[env_idx, road_mask].tolist()
@@ -264,7 +266,11 @@ def _plot_roadgraph(
                         env_idx, road_mask
                     ].tolist()
 
-                    if( not road_point_type == int(gpudrive.EntityType.SpeedBump) ):
+                    if(
+                        road_point_type == int(gpudrive.EntityType.RoadEdge)
+                        or road_point_type == int(gpudrive.EntityType.RoadLine)
+                        or road_point_type == int(gpudrive.EntityType.RoadLane)
+                    ):
                         # Compute and draw road edges using start and end points
                         for x, y, length, orientation in zip(x_coords, y_coords, segment_lengths, segment_orientations):
                             start, end = self._get_endpoints(x, y, length, orientation)
@@ -277,40 +283,28 @@ def _plot_roadgraph(
                                 linewidth=0.75
                             )
 
-                    if (road_point_type == int(gpudrive.EntityType.SpeedBump)):
-                        utils.plot_speed_bump(
+                    elif (road_point_type == int(gpudrive.EntityType.SpeedBump)):
+                        utils.plot_speed_bumps(
                             x_coords, 
                             y_coords, 
                             segment_lengths, 
                             segment_widths, 
                             segment_orientations,
                             ax
-                        )          
-                    # # Compute and draw road edges using start and end points
-                    # for x, y, length, orientation in zip(
-                    #     x_coords,
-                    #     y_coords,
-                    #     segment_lengths,
-                    #     segment_orientations,
-                    # ):
-                    #     start, end = self._get_endpoints(
-                    #         x, y, length, orientation
-                    #     )
-
-                    #     if road_point_type == int(
-                    #         gpudrive.EntityType.RoadEdge
-                    #     ):
-                    #         line_width = 1.1 * line_width_scale
-
-                    #     else:
-                    #         line_width = 0.75 * line_width_scale
-
-                    #     ax.plot(
-                    #         [start[0], end[0]],
-                    #         [start[1], end[1]],
-                    #         color=ROAD_GRAPH_COLORS[road_point_type],
-                    #         linewidth=line_width,
-                    #     )
+                        )  
+
+                    elif (road_point_type == int(gpudrive.EntityType.StopSign)):
+                        for x, y in zip(x_coords, y_coords):
+                            point = np.array([x, y])
+                            utils.plot_stop_sign(
+                                point=point,
+                                ax=ax,
+                                radius=1.5,        
+                                facecolor='xkcd:red',      
+                                edgecolor="none",    
+                                linewidth=3.0,        
+                                alpha=1.0,            
+                            )
 
                 else:
                     # Dots for other road point types
@@ -559,7 +553,13 @@ def plot_agent_observation(
         if observation_ego.id[env_idx, agent_idx] == -1:
             return None, None
 
-        fig, ax = plt.subplots(figsize=figsize)
+        # fig, ax = plt.subplots(figsize=figsize)
+        viz_config = (
+            utils.VizConfig()
+            if viz_config is None
+            else utils.VizConfig(**viz_config)
+        )
+        fig, ax = utils.init_fig_ax(viz_config)
         ax.clear()  # Clear any previous plots
         ax.set_aspect("equal", adjustable="box")
         ax.set_title(f"obs agent: {agent_idx}")

pygpudrive/visualize/utils.py

@@ -5,7 +5,7 @@
 import matplotlib.pylab as plt
 import numpy as np
 from PIL import Image
-from matplotlib.patches import Circle, Polygon
+from matplotlib.patches import Circle, Polygon, RegularPolygon
 
 import os
 import torch
@@ -260,7 +260,7 @@ def plot_bounding_box(
 def get_corners_polygon(x, y, length, width, orientation):
     """Calculate the four corners of a speed bump (can be any) polygon."""
     # Compute the direction vectors based on orientation
-    print(length)
+    # print(length)
     c = np.cos(orientation)
     s = np.sin(orientation)
     u = np.array((c, s))  # Unit vector along the orientation
@@ -319,7 +319,7 @@ def get_stripe_polygon(
 
     return np.array(stripe_corners)
 
-def plot_speed_bump(
+def plot_speed_bumps(
     x_coords: Union[float, np.ndarray],
     y_coords: Union[float, np.ndarray],
     segment_lengths: Union[float, torch.Tensor],
@@ -373,4 +373,34 @@ def plot_speed_bump(
         #     )
         #     ax.add_patch(stripe_polygon)
 
-    pass
+    pass
+
+def plot_stop_sign(
+    point: np.ndarray,
+    ax: matplotlib.axes.Axes,
+    radius: float = None,
+    facecolor: str = None,
+    edgecolor: str = None,
+    linewidth: float = None,
+    alpha: float = None,
+) -> None:
+    # Default configurations for the stop sign
+    facecolor = "red" if facecolor is None else facecolor
+    edgecolor = "white" if edgecolor is None else edgecolor
+    linewidth = 1.5 if linewidth is None else linewidth
+    radius = 1.0 if radius is None else radius
+    alpha = 1.0 if alpha is None else alpha
+
+    point = np.array(point).reshape(-1)
+
+    p = RegularPolygon(
+        point,
+        numVertices=6,  # For hexagonal stop sign
+        radius=radius,
+        facecolor=facecolor,
+        edgecolor=edgecolor,
+        linewidth=linewidth,
+        alpha=alpha,
+        zorder=2,
+    )
+    ax.add_patch(p)