Source code for simulation.src.simulation_groundtruth.src.groundtruth.groundtruth

from collections import namedtuple
from dataclasses import dataclass
from typing import List, Tuple

from simulation_groundtruth.msg import LabeledPolygon as LabeledPolygonMsg
from simulation_groundtruth.msg import Lane as LaneMsg
from simulation_groundtruth.msg import Line as LineMsg
from simulation_groundtruth.msg import Parking as ParkingMsg
from simulation_groundtruth.msg import Section as SectionMsg

from simulation.utils.geometry import Line, Polygon
from simulation.utils.road.road import Road
from simulation.utils.road.sections import Intersection


[docs]def _lane_msg_from_lines(left: Line, middle: Line, right: Line) -> LaneMsg: """Create a lane message from three lines.""" msg = LaneMsg() msg.middle_line = middle.to_geometry_msgs() msg.left_line = left.to_geometry_msgs() msg.right_line = right.to_geometry_msgs() return msg
IntersectionTuple = namedtuple( "IntersectionTuple", ["turn", "rule", "south", "west", "east", "north"], ) """Container for information about an intersection.""" IntersectionTuple.turn.__doc__ = """int: Direction the road follows.""" IntersectionTuple.rule.__doc__ = """int: Priority rule.""" IntersectionTuple.south.__doc__ = """LaneMsg: Lane markings approaching the intersection.""" IntersectionTuple.west.__doc__ = ( """LaneMsg: Lane markings on the left side when coming from south.""" ) IntersectionTuple.east.__doc__ = ( """LaneMsg: Lane markings on the right side when coming from south.""" ) IntersectionTuple.north.__doc__ = ( """LaneMsg: Lane markings straight ahead when coming from south.""" )
[docs]@dataclass class Groundtruth: """Provide functionality to extract groundtruth information as ROS messages.""" road: Road """Store the current road."""
[docs] def get_section_msgs(self) -> List[SectionMsg]: """Section message for all sections.""" messages = [] for section in self.road.sections: msg = SectionMsg() msg.id = section.id msg.type = section.__class__.TYPE messages.append(msg) return messages
[docs] def get_lane_msg(self, id: int) -> LaneMsg: """Lane message for requested road section. Args: id: section id """ section = self.road.sections[id] return _lane_msg_from_lines( section.left_line, section.middle_line, section.right_line )
[docs] def get_intersection_msg(self, id: int) -> IntersectionTuple: """Intersection tuple for requested road section. Args: id: section id """ intersection: Intersection = self.road.sections[id] assert intersection.__class__.TYPE == Intersection.TYPE south = _lane_msg_from_lines( intersection.left_line_south, intersection.middle_line_south, intersection.right_line_south, ) west = _lane_msg_from_lines( intersection.left_line_west, intersection.middle_line_west, intersection.right_line_west, ) east = _lane_msg_from_lines( intersection.left_line_east, intersection.middle_line_east, intersection.right_line_east, ) north = _lane_msg_from_lines( intersection.left_line_north, intersection.middle_line_north, intersection.right_line_north, ) return IntersectionTuple( intersection.turn, intersection.rule, south, west, east, north )
[docs] def get_parking_msg(self, id: int) -> Tuple[ParkingMsg, ParkingMsg]: """Parking message for the left and right side of the requested road section. Args: id: section id """ def msg_from(lots): borders = [lot.border for lot in lots] spots = sum((lot.spots for lot in lots), []) msg = ParkingMsg() msg.borders = [LineMsg(border.to_geometry_msgs()) for border in borders] for spot in spots: spot_msg = LabeledPolygonMsg() spot_msg.frame = spot.frame.to_geometry_msg() spot_msg.type = spot.kind spot_msg.visible = True msg.spots.append(spot_msg) return msg left_msg = msg_from(self.road.sections[id].left_lots) right_msg = msg_from(self.road.sections[id].right_lots) return left_msg, right_msg
[docs] def get_obstacle_msgs(self, id: int) -> List[LabeledPolygonMsg]: """Labeled polygon msg for each obstacle in the requested road section. Args: id: section id """ obstacles = self.road.sections[id].obstacles if obstacles is None: return [] msgs = [] for obstacle in obstacles: assert isinstance(obstacle.frame, Polygon) assert isinstance(obstacle.height, float) assert isinstance(obstacle.id_, int) msg = LabeledPolygonMsg() msg.frame = obstacle.frame.to_geometry_msg() msg.height = obstacle.height msg.type = obstacle.id_ msg.desc = obstacle.desc msg.visible = True msgs.append(msg) return msgs
[docs] def get_surface_marking_msgs(self, id: int) -> List[LabeledPolygonMsg]: """Labeled polygon msg for each surface marking in the requested road section. Args: id: section id """ surface_markings = self.road.sections[id].surface_markings if surface_markings is None: return [] msgs = [] for surface_marking in surface_markings: msg = LabeledPolygonMsg() msg.frame = surface_marking.frame.to_geometry_msg() msg.type = surface_marking.kind[0] msg.desc = surface_marking.kind[1] msg.visible = True msgs.append(msg) return msgs
[docs] def get_traffic_sign_msgs(self, id: int) -> List[LabeledPolygonMsg]: """Labeled polygon msg for each traffic sign in the requested road section. Args: id: section id """ signs = self.road.sections[id].traffic_signs if signs is None: return [] msgs = [] for sign in signs: msg = LabeledPolygonMsg() msg.frame = sign.frame.to_geometry_msg() msg.type = sign.kind.id_ msg.desc = sign.kind.mesh msg.height = sign.kind.collision_box_size[2] msg.visible = sign.visible msgs.append(msg) return msgs