diff --git a/gamms/VisualizationEngine/builtin_artists.py b/gamms/VisualizationEngine/builtin_artists.py index 8ec9dea..5e16c55 100644 --- a/gamms/VisualizationEngine/builtin_artists.py +++ b/gamms/VisualizationEngine/builtin_artists.py @@ -67,4 +67,19 @@ class GraphData: node_size: float edge_color: ColorType draw_id: bool - edge_line_points: Dict[int, List[Tuple[float, float]]] = field(default_factory=dict) \ No newline at end of file + edge_line_points: Dict[int, List[Tuple[float, float]]] = field(default_factory=dict) + +@dataclass +class LabelData: + """Contains all necessary data for drawing a label. + Attributes: + text (str): The text of the label. + color (Optional[ColorType]): The color of the label text. + size (Optional[int]): The font size of the label text. + offset (Tuple[float, float]): The offset of the label. + visible (bool): Whether the label is visible.""" + text: str + color: Optional[ColorType] = None + size: Optional[int] = None + offset: Tuple[float, float] = (0.0, 0.0) + visible: bool = True diff --git a/gamms/VisualizationEngine/default_drawers.py b/gamms/VisualizationEngine/default_drawers.py index 5614fd5..f39ad83 100644 --- a/gamms/VisualizationEngine/default_drawers.py +++ b/gamms/VisualizationEngine/default_drawers.py @@ -1,6 +1,6 @@ from gamms.AgentEngine.agent_engine import AerialAgent from gamms.VisualizationEngine import Color -from gamms.VisualizationEngine.builtin_artists import AgentData, GraphData +from gamms.VisualizationEngine.builtin_artists import AgentData, GraphData, LabelData from gamms.typing import IContext, OSMEdge, Node, ColorType, AgentType from typing import Dict, Any, cast, List, Optional @@ -38,6 +38,26 @@ def render_rectangle(ctx: IContext, data: Dict[str, Any]): color = data.get('color', Color.Cyan) ctx.visual.render_rectangle(x, y, width, height, color) +def render_label(ctx: IContext, label_data: LabelData, anchor_x: float, anchor_y: float, anchor_color: Optional[ColorType], anchor_size: Optional[int]): + """ + Render a label with the specified text at the specified position. + + Args: + ctx (Context): The current simulation context. + label_data (LabelData): The data containing the label's text, color, size, and offset. + anchor_x (float): The x-coordinate of the anchor point for the label. + anchor_y (float): The y-coordinate of the anchor point for the label. + anchor_color (Optional[ColorType]): The color of the anchor, used as a fallback if label_data.color is not provided. + anchor_size (Optional[int]): The size of the anchor, used as a fallback if label_data.size is not provided. + """ + if label_data.visible is False: + return + color = cast(ColorType, label_data.color if label_data.color is not None else anchor_color if anchor_color is not None else Color.Black) + size = cast(Optional[int], label_data.size if label_data.size is not None else anchor_size) + x = anchor_x + label_data.offset[0] + y = anchor_y + label_data.offset[1] + ctx.visual.render_text(label_data.text, x, y, color, font_size=size) + def render_agent(ctx: IContext, data: Dict[str, Any]): """ Render an agent as a triangle at its current position on the screen. This is the default rendering method for agents. @@ -47,6 +67,7 @@ def render_agent(ctx: IContext, data: Dict[str, Any]): data (dict): The data containing the agent's information. """ agent_data = cast(AgentData, data.get('agent_data')) + label_data = cast(Optional[LabelData], data.get('label_data')) size = agent_data.size color = agent_data.color is_waiting = data.get('_is_waiting', False) @@ -76,11 +97,11 @@ def render_agent(ctx: IContext, data: Dict[str, Any]): else: position = (prev_position[0] + alpha * (target_position[0] - prev_position[0]), prev_position[1] + alpha * (target_position[1] - prev_position[1])) - - agent_data.current_position = position else: position = (target_node.x, target_node.y) + agent_data.current_position = position + # Draw each agent as a triangle at its current position angle = math.radians(45) @@ -101,6 +122,8 @@ def render_agent(ctx: IContext, data: Dict[str, Any]): else: position = aerial_agent.position + agent_data.current_position = (position[0], position[1]) + quat = aerial_agent.quat x = quat[1] y = quat[2] @@ -112,6 +135,9 @@ def render_agent(ctx: IContext, data: Dict[str, Any]): else: raise ValueError(f"Unsupported agent type: {agent.type}") + + if label_data is not None: + render_label(ctx, label_data, agent_data.current_position[0], agent_data.current_position[1], color, size) def render_aerial_agent(ctx: IContext, position: tuple[float, float], angle: float, size: float, color: ColorType): diff --git a/gamms/VisualizationEngine/pygame_engine.py b/gamms/VisualizationEngine/pygame_engine.py index 681056b..150e765 100644 --- a/gamms/VisualizationEngine/pygame_engine.py +++ b/gamms/VisualizationEngine/pygame_engine.py @@ -1,7 +1,7 @@ from gamms.AgentEngine.agent_engine import AerialAgent from gamms.VisualizationEngine import Color, Space, Shape, Artist, lazy from gamms.VisualizationEngine.render_manager import RenderManager -from gamms.VisualizationEngine.builtin_artists import AgentData, GraphData +from gamms.VisualizationEngine.builtin_artists import AgentData, GraphData, LabelData from gamms.VisualizationEngine.default_drawers import ( render_circle, render_rectangle, render_agent, render_graph, render_neighbor_sensor, @@ -131,8 +131,16 @@ def set_agent_visual(self, name: str, **kwargs: Dict[str, Any]) -> IArtist: artist.set_artist_type(ArtistType.DYNAMIC) artist.data['_alpha'] = 1.0 - self.add_artist(name, artist) + if 'label_data' in kwargs: + artist.data['label_data'] = LabelData( + text=cast(str, kwargs['label_data'].get('text', name)), + color=cast(Optional[ColorType], kwargs['label_data'].get('color', None)), + size=cast(Optional[int], kwargs['label_data'].get('size', None)), + offset=cast(Tuple[float, float], kwargs['label_data'].get('offset', (0.0, -15.0))), + visible=cast(bool, kwargs['label_data'].get('visible', True)), + ) + self.add_artist(name, artist) return artist def set_sensor_visual(self, name: str, **kwargs: Dict[str, Any]) -> IArtist: diff --git a/snippets/custom_sensors/config.py b/snippets/custom_sensors/config.py index c295733..cfa997e 100644 --- a/snippets/custom_sensors/config.py +++ b/snippets/custom_sensors/config.py @@ -58,12 +58,24 @@ agent_vis_config[f'agent_{i}'] = { 'color': (255, 0, 0), # color of the agent 'size': 10, # size of the agent + 'label_data': { # label data for the agent + 'text': f'Agent {i}', # text of the label + 'color': (255, 0, 0), # color of the label (optional as it will default to agent color if not provided) + 'size': 12, + 'offset': (0, -15), + }, } for i in range(RED_TEAM_AGENTS, RED_TEAM_AGENTS + BLUE_TEAM_AGENTS): agent_vis_config[f'agent_{i}'] = { 'color': (0, 0, 255), # color of the agent 'size': 10, # size of the agent + 'label_data': { # label data for the agent + 'text': f'Agent {i}', # text of the label + 'color': (0, 0, 255), # color of the label (optional as it will default to agent color if not provided) + 'size': 12, + 'offset': (0, -15), + }, } sensor_vis_config = {}