Source code for highway_env.road.graphics

from typing import List, Tuple, Union, TYPE_CHECKING

import numpy as np
import pygame

from highway_env.road.lane import LineType, AbstractLane
from highway_env.road.road import Road
from highway_env.types import Vector
from highway_env.vehicle.graphics import VehicleGraphics
from highway_env.road.objects import Obstacle, Landmark

if TYPE_CHECKING:
    from highway_env.road.objects import RoadObject

PositionType = Union[Tuple[float, float], np.ndarray]


[docs]class WorldSurface(pygame.Surface): """A pygame Surface implementing a local coordinate system so that we can move and zoom in the displayed area.""" BLACK = (0, 0, 0) GREY = (100, 100, 100) GREEN = (50, 200, 0) YELLOW = (200, 200, 0) WHITE = (255, 255, 255) INITIAL_SCALING = 5.5 INITIAL_CENTERING = [0.5, 0.5] SCALING_FACTOR = 1.3 MOVING_FACTOR = 0.1
[docs] def __init__(self, size: Tuple[int, int], flags: object, surf: pygame.SurfaceType) -> None: super().__init__(size, flags, surf) self.origin = np.array([0, 0]) self.scaling = self.INITIAL_SCALING self.centering_position = self.INITIAL_CENTERING
[docs] def pix(self, length: float) -> int: """ Convert a distance [m] to pixels [px]. :param length: the input distance [m] :return: the corresponding size [px] """ return int(length * self.scaling)
[docs] def pos2pix(self, x: float, y: float) -> Tuple[int, int]: """ Convert two world coordinates [m] into a position in the surface [px] :param x: x world coordinate [m] :param y: y world coordinate [m] :return: the coordinates of the corresponding pixel [px] """ return self.pix(x - self.origin[0]), self.pix(y - self.origin[1])
[docs] def vec2pix(self, vec: PositionType) -> Tuple[int, int]: """ Convert a world position [m] into a position in the surface [px]. :param vec: a world position [m] :return: the coordinates of the corresponding pixel [px] """ return self.pos2pix(vec[0], vec[1])
[docs] def move_display_window_to(self, position: PositionType) -> None: """ Set the origin of the displayed area to center on a given world position. :param position: a world position [m] """ self.origin = position - np.array( [self.centering_position[0] * self.get_width() / self.scaling, self.centering_position[1] * self.get_height() / self.scaling])
[docs] def handle_event(self, event: pygame.event.EventType) -> None: """ Handle pygame events for moving and zooming in the displayed area. :param event: a pygame event """ if event.type == pygame.KEYDOWN: if event.key == pygame.K_l: self.scaling *= 1 / self.SCALING_FACTOR if event.key == pygame.K_o: self.scaling *= self.SCALING_FACTOR if event.key == pygame.K_m: self.centering_position[0] -= self.MOVING_FACTOR if event.key == pygame.K_k: self.centering_position[0] += self.MOVING_FACTOR
[docs]class LaneGraphics(object): """A visualization of a lane.""" STRIPE_SPACING: float = 5 """ Offset between stripes [m]""" STRIPE_LENGTH: float = 3 """ Length of a stripe [m]""" STRIPE_WIDTH: float = 0.3 """ Width of a stripe [m]"""
[docs] @classmethod def display(cls, lane: AbstractLane, surface: WorldSurface) -> None: """ Display a lane on a surface. :param lane: the lane to be displayed :param surface: the pygame surface """ stripes_count = int(2 * (surface.get_height() + surface.get_width()) / (cls.STRIPE_SPACING * surface.scaling)) s_origin, _ = lane.local_coordinates(surface.origin) s0 = (int(s_origin) // cls.STRIPE_SPACING - stripes_count // 2) * cls.STRIPE_SPACING for side in range(2): if lane.line_types[side] == LineType.STRIPED: cls.striped_line(lane, surface, stripes_count, s0, side) elif lane.line_types[side] == LineType.CONTINUOUS: cls.continuous_curve(lane, surface, stripes_count, s0, side) elif lane.line_types[side] == LineType.CONTINUOUS_LINE: cls.continuous_line(lane, surface, stripes_count, s0, side)
[docs] @classmethod def striped_line(cls, lane: AbstractLane, surface: WorldSurface, stripes_count: int, longitudinal: float, side: int) -> None: """ Draw a striped line on one side of a lane, on a surface. :param lane: the lane :param surface: the pygame surface :param stripes_count: the number of stripes to draw :param longitudinal: the longitudinal position of the first stripe [m] :param side: which side of the road to draw [0:left, 1:right] """ starts = longitudinal + np.arange(stripes_count) * cls.STRIPE_SPACING ends = longitudinal + np.arange(stripes_count) * cls.STRIPE_SPACING + cls.STRIPE_LENGTH lats = [(side - 0.5) * lane.width_at(s) for s in starts] cls.draw_stripes(lane, surface, starts, ends, lats)
[docs] @classmethod def continuous_curve(cls, lane: AbstractLane, surface: WorldSurface, stripes_count: int, longitudinal: float, side: int) -> None: """ Draw a striped line on one side of a lane, on a surface. :param lane: the lane :param surface: the pygame surface :param stripes_count: the number of stripes to draw :param longitudinal: the longitudinal position of the first stripe [m] :param side: which side of the road to draw [0:left, 1:right] """ starts = longitudinal + np.arange(stripes_count) * cls.STRIPE_SPACING ends = longitudinal + np.arange(stripes_count) * cls.STRIPE_SPACING + cls.STRIPE_SPACING lats = [(side - 0.5) * lane.width_at(s) for s in starts] cls.draw_stripes(lane, surface, starts, ends, lats)
[docs] @classmethod def continuous_line(cls, lane: AbstractLane, surface: WorldSurface, stripes_count: int, longitudinal: float, side: int) -> None: """ Draw a continuous line on one side of a lane, on a surface. :param lane: the lane :param surface: the pygame surface :param stripes_count: the number of stripes that would be drawn if the line was striped :param longitudinal: the longitudinal position of the start of the line [m] :param side: which side of the road to draw [0:left, 1:right] """ starts = [longitudinal + 0 * cls.STRIPE_SPACING] ends = [longitudinal + stripes_count * cls.STRIPE_SPACING + cls.STRIPE_LENGTH] lats = [(side - 0.5) * lane.width_at(s) for s in starts] cls.draw_stripes(lane, surface, starts, ends, lats)
[docs] @classmethod def draw_stripes(cls, lane: AbstractLane, surface: WorldSurface, starts: List[float], ends: List[float], lats: List[float]) -> None: """ Draw a set of stripes along a lane. :param lane: the lane :param surface: the surface to draw on :param starts: a list of starting longitudinal positions for each stripe [m] :param ends: a list of ending longitudinal positions for each stripe [m] :param lats: a list of lateral positions for each stripe [m] """ starts = np.clip(starts, 0, lane.length) ends = np.clip(ends, 0, lane.length) for k, _ in enumerate(starts): if abs(starts[k] - ends[k]) > 0.5 * cls.STRIPE_LENGTH: pygame.draw.line(surface, surface.WHITE, (surface.vec2pix(lane.position(starts[k], lats[k]))), (surface.vec2pix(lane.position(ends[k], lats[k]))), max(surface.pix(cls.STRIPE_WIDTH), 1))
[docs] @classmethod def draw_ground(cls, lane: AbstractLane, surface: WorldSurface, color: Tuple[float], width: float, draw_surface: pygame.Surface = None) -> None: draw_surface = draw_surface or surface stripes_count = int(2 * (surface.get_height() + surface.get_width()) / (cls.STRIPE_SPACING * surface.scaling)) s_origin, _ = lane.local_coordinates(surface.origin) s0 = (int(s_origin) // cls.STRIPE_SPACING - stripes_count // 2) * cls.STRIPE_SPACING dots = [] for side in range(2): longis = np.clip(s0 + np.arange(stripes_count) * cls.STRIPE_SPACING, 0, lane.length) lats = [2 * (side - 0.5) * width for _ in longis] new_dots = [surface.vec2pix(lane.position(longi, lat)) for longi, lat in zip(longis, lats)] new_dots = reversed(new_dots) if side else new_dots dots.extend(new_dots) pygame.draw.polygon(draw_surface, color, dots, 0)
[docs]class RoadGraphics(object): """A visualization of a road lanes and vehicles."""
[docs] @staticmethod def display(road: Road, surface: WorldSurface) -> None: """ Display the road lanes on a surface. :param road: the road to be displayed :param surface: the pygame surface """ surface.fill(surface.GREY) for _from in road.network.graph.keys(): for _to in road.network.graph[_from].keys(): for l in road.network.graph[_from][_to]: LaneGraphics.display(l, surface)
[docs] @staticmethod def display_traffic(road: Road, surface: WorldSurface, simulation_frequency: int = 15, offscreen: bool = False) \ -> None: """ Display the road vehicles on a surface. :param road: the road to be displayed :param surface: the pygame surface :param simulation_frequency: simulation frequency :param offscreen: render without displaying on a screen """ if road.record_history: for v in road.vehicles: VehicleGraphics.display_history(v, surface, simulation=simulation_frequency, offscreen=offscreen) for v in road.vehicles: VehicleGraphics.display(v, surface, offscreen=offscreen)
[docs] @staticmethod def display_road_objects(road: Road, surface: WorldSurface, offscreen: bool = False) -> None: """ Display the road objects on a surface. :param road: the road to be displayed :param surface: the pygame surface :param offscreen: whether the rendering should be done offscreen or not """ for o in road.objects: RoadObjectGraphics.display(o, surface, offscreen=offscreen)
[docs]class RoadObjectGraphics: """A visualization of objects on the road.""" YELLOW = (200, 200, 0) BLUE = (100, 200, 255) RED = (255, 100, 100) GREEN = (50, 200, 0) BLACK = (60, 60, 60) DEFAULT_COLOR = YELLOW
[docs] @classmethod def display(cls, object_: 'RoadObject', surface: WorldSurface, transparent: bool = False, offscreen: bool = False): """ Display a road objects on a pygame surface. The objects is represented as a colored rotated rectangle :param object_: the vehicle to be drawn :param surface: the surface to draw the object on :param transparent: whether the object should be drawn slightly transparent :param offscreen: whether the rendering should be done offscreen or not """ o = object_ s = pygame.Surface((surface.pix(o.LENGTH), surface.pix(o.LENGTH)), pygame.SRCALPHA) # per-pixel alpha rect = (0, surface.pix(o.WIDTH) / 2 - surface.pix(o.WIDTH) / 2, surface.pix(o.LENGTH), surface.pix(o.WIDTH)) pygame.draw.rect(s, cls.get_color(o, transparent), rect, 0) pygame.draw.rect(s, cls.BLACK, rect, 1) if not offscreen: # convert_alpha throws errors in offscreen mode TODO() Explain why s = pygame.Surface.convert_alpha(s) h = o.heading if abs(o.heading) > 2 * np.pi / 180 else 0 # Centered rotation position = (surface.pos2pix(o.position[0], o.position[1])) cls.blit_rotate(surface, s, position, np.rad2deg(-h))
[docs] @staticmethod def blit_rotate(surf: pygame.SurfaceType, image: pygame.SurfaceType, pos: Vector, angle: float, origin_pos: Vector = None, show_rect: bool = False) -> None: """Many thanks to https://stackoverflow.com/a/54714144.""" # calculate the axis aligned bounding box of the rotated image w, h = image.get_size() box = [pygame.math.Vector2(p) for p in [(0, 0), (w, 0), (w, -h), (0, -h)]] box_rotate = [p.rotate(angle) for p in box] min_box = (min(box_rotate, key=lambda p: p[0])[0], min(box_rotate, key=lambda p: p[1])[1]) max_box = (max(box_rotate, key=lambda p: p[0])[0], max(box_rotate, key=lambda p: p[1])[1]) # calculate the translation of the pivot if origin_pos is None: origin_pos = w / 2, h / 2 pivot = pygame.math.Vector2(origin_pos[0], -origin_pos[1]) pivot_rotate = pivot.rotate(angle) pivot_move = pivot_rotate - pivot # calculate the upper left origin of the rotated image origin = (pos[0] - origin_pos[0] + min_box[0] - pivot_move[0], pos[1] - origin_pos[1] - max_box[1] + pivot_move[1]) # get a rotated image rotated_image = pygame.transform.rotate(image, angle) # rotate and blit the image surf.blit(rotated_image, origin) # draw rectangle around the image if show_rect: pygame.draw.rect(surf, (255, 0, 0), (*origin, *rotated_image.get_size()), 2)
[docs] @classmethod def get_color(cls, object_: 'RoadObject', transparent: bool = False): color = cls.DEFAULT_COLOR if isinstance(object_, Obstacle): if object_.hit: # indicates failure color = cls.RED else: color = cls.YELLOW elif isinstance(object_, Landmark): if object_.hit: # indicates success color = cls.GREEN else: color = cls.BLUE if transparent: color = (color[0], color[1], color[2], 30) return color