torchdrivesim.mesh ================== .. py:module:: torchdrivesim.mesh .. autoapi-nested-parse:: Custom classes for representing various triangular meshes as tensors. Attributes ---------- .. autoapisummary:: torchdrivesim.mesh.logger torchdrivesim.mesh.Color Exceptions ---------- .. autoapisummary:: torchdrivesim.mesh.BadMeshFormat Classes ------- .. autoapisummary:: torchdrivesim.mesh.BaseMesh torchdrivesim.mesh.AttributeMesh torchdrivesim.mesh.RGBMesh torchdrivesim.mesh.BirdviewMesh torchdrivesim.mesh.BirdviewRGBMeshGenerator Functions --------- .. autoapisummary:: torchdrivesim.mesh.tensor_color torchdrivesim.mesh.rendering_mesh torchdrivesim.mesh.set_colors_with_defaults torchdrivesim.mesh.generate_trajectory_mesh torchdrivesim.mesh.generate_annulus_polygon_mesh torchdrivesim.mesh.generate_disc_mesh torchdrivesim.mesh.build_verts_faces_from_bounding_box Module Contents --------------- .. py:data:: logger :value: None .. py:data:: Color .. py:function:: tensor_color(color: Color, device: torch.device = 'cpu', dtype=torch.float) -> torch.Tensor Converts all supported color representations to the tensor representation. :param color: RGB color as either an int 3-tuple in [0,255] range or tensor of shape (3,) in [0,1] range :param device: device to place the resulting tensor on, ignored if color is already a tensor :param dtype: element type for the resulting tensor, ignored if color is already a tensor :returns: RGB tensor of shape (3,) in [0,1] range .. py:exception:: BadMeshFormat Bases: :py:obj:`RuntimeError` The mesh received had the wrong format, usually indicating loading from old cache. .. py:class:: BaseMesh Basic triangle mesh in a space of arbitrary dimensions Dim. Only specifies triangles, with no additional properties. Always includes exactly one batch dimension. .. py:attribute:: verts :type: torch.Tensor .. py:attribute:: faces :type: torch.Tensor .. py:attribute:: _verts_fill :type: float :value: 0.0 .. py:attribute:: _faces_fill :type: int :value: 0 .. py:method:: __post_init__() .. py:method:: _validate_input() .. py:property:: dim :type: int Dimension of the space in which the mesh lives, usually 2 or 3. .. py:property:: batch_size :type: int .. py:property:: verts_count :type: int .. py:property:: faces_count :type: int .. py:property:: device :type: torch.device .. py:property:: center :type: torch.Tensor A BxDim center of the mesh, calculated as the average of min and max vertex coordinates for each dimension. Note that if the vertices are padded, the padding value may distort this result. .. py:method:: to(device: torch.device) Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object. .. py:method:: clone() Deep copy of the mesh. .. py:method:: expand(size: int) Expands batch dimension on the right by the given factor, returning a new mesh. .. py:method:: select_batch_elements(idx: torch.Tensor) Selects given indices from batch dimension, possibly with repetitions, returning a new mesh .. py:method:: translate(xy: torch.Tensor, inplace: bool = True) -> typing_extensions.Self Shifts the mesh by the given coordinate so that it is at the origin (0,0). :param xy: Bx2 tensor specifying the camera position :param inplace: whether to modify the mesh in place .. py:method:: __getitem__(item) .. py:method:: collate(meshes) :classmethod: Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension. .. py:method:: concat(meshes) :classmethod: Concatenates multiple meshes to form a single scene. .. py:method:: merge(other) Merges the current mesh with another to form a single scene, returning a new mesh. .. py:method:: offset(offset: torch.Tensor) Shifts the mesh by a given distance, returning a new mesh. Missing dimensions in the argument are padded with zeros if needed. .. py:method:: pad(pad_size: int) .. py:method:: pytorch3d(include_textures=True) -> pytorch3d.structures.Meshes Converts the mesh to a PyTorch3D one. For the base class there are no textures, but subclasses may include them. Empty meshes are augmented with a single degenerate face on conversion, since PyTorch3D does not handle empty meshes correctly. .. py:method:: pickle(mesh_file_path: str) Store this mesh to a given file. .. py:method:: unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module) :classmethod: Load a mesh of this type from the given file. .. py:method:: serialize() .. py:method:: save(file_save_path: str) Save the attributes of the mesh object in a json where tensors are converted to lists. .. py:method:: _deserialize_tensors(data: Dict) -> Dict :classmethod: Convert list attributes to tensor attributes if there is any tensor attribute .. py:method:: deserialize(data: Dict) -> typing_extensions.Self :classmethod: .. py:method:: load(filepath) :classmethod: .. py:method:: empty(dim: int = 2, batch_size: int = 1) :classmethod: Create empty mesh. .. py:method:: _trim_and_return_verts_and_faces(vertices_to_keep: torch.Tensor, trim_face_only=False) .. py:method:: trim(polygon: torch.Tensor, trim_face_only: bool = False) Crops the mesh to a given 2D convex polygon, returning a new mesh. Faces where all vertices are outside the polygon are removed, even if the triangle intersects with the polygon. Vertices are removed if they are not used by any face. :param polygon: BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion :param trim_face_only: whether to keep all vertices, including those unused by any faces .. py:class:: AttributeMesh Bases: :py:obj:`BaseMesh` Endows each vertex with an attribute. An attribute is a vector of arbitrary dimension Attr, usually a color or something similar. Typically, in any given face all vertices have the same attribute values. .. py:attribute:: attrs :type: torch.Tensor .. py:attribute:: _attrs_fill :type: float :value: 0.0 .. py:method:: _validate_input() .. py:property:: attr_dim :type: int Size of the attribute dimension. .. py:method:: set_attr(mesh: BaseMesh, attr: torch.Tensor) :classmethod: Sets a given attribute value for all vertices in a given mesh. .. py:method:: to(device: torch.device) Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object. .. py:method:: pad(pad_size: int) .. py:method:: expand(size: int) Expands batch dimension on the right by the given factor, returning a new mesh. .. py:method:: select_batch_elements(idx: int) Selects given indices from batch dimension, possibly with repetitions, returning a new mesh .. py:method:: concat(meshes) :classmethod: Concatenates multiple meshes to form a single scene. .. py:method:: collate(meshes) :classmethod: Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension. .. py:method:: pytorch3d(include_textures=True) -> pytorch3d.structures.Meshes PyTorch3D uses per-face textures, which are obtained by averaging attributes of the face. The resulting texture for each face is constant. .. py:method:: unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module) :classmethod: Load a mesh of this type from the given file. .. py:method:: serialize() .. py:method:: _deserialize_tensors(data: Dict) -> Dict :classmethod: Convert list attributes to tensor attributes if there is any tensor attribute .. py:method:: load(filepath) :classmethod: .. py:method:: empty(dim=2, batch_size=1, attr_dim=3) :classmethod: Create empty mesh. .. py:method:: trim(polygon: torch.Tensor, trim_face_only: bool = False) Crops the mesh to a given 2D convex polygon, returning a new mesh. Faces where all vertices are outside the polygon are removed, even if the triangle intersects with the polygon. Vertices are removed if they are not used by any face. :param polygon: BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion :param trim_face_only: whether to keep all vertices, including those unused by any faces .. py:class:: RGBMesh Bases: :py:obj:`AttributeMesh` AttributeMesh where the attribute is an RGB color in [0,1] range. .. py:method:: _validate_input() .. py:method:: set_color(mesh: BaseMesh, color: Color) :classmethod: Sets a constant color for all vertices in a given mesh. .. py:class:: BirdviewMesh Bases: :py:obj:`BaseMesh` 2D mesh with vertices and faces assigned to discrete categories, to facilitate rendering simple 2D worlds. Category assignment is stored per vertex, and faces should not mix vertices from different categories. For each category there is a color and a rendering priority z (lower renders on top), but those don't need to be specified until the mesh is converted to a different representation. .. py:attribute:: categories :type: List[str] .. py:attribute:: colors :type: Dict[str, torch.Tensor] .. py:attribute:: zs :type: Dict[str, float] .. py:attribute:: vert_category :type: torch.Tensor .. py:attribute:: _cat_fill :type: int :value: 0 .. py:method:: _validate_input() .. py:property:: num_categories :type: int .. py:method:: set_properties(mesh: BaseMesh, category: str, color: Optional[Color] = None, z: Optional[float] = None) :classmethod: Lifts a BaseMesh into a BirdviewMesh with a single category. .. py:method:: to(device) Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object. .. py:method:: pad(pad_size: int) .. py:method:: expand(size) Expands batch dimension on the right by the given factor, returning a new mesh. .. py:method:: select_batch_elements(idx) Selects given indices from batch dimension, possibly with repetitions, returning a new mesh .. py:method:: unify(meshes) :classmethod: Generates meshes equivalent to input meshes, only all sharing the same category definitions. .. py:method:: concat(meshes) :classmethod: Concatenates multiple meshes to form a single scene. .. py:method:: collate(meshes) :classmethod: Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension. .. py:method:: fill_attr() -> RGBMesh Computes explicit color for each vertex and augments vertices with z values corresponding to rendering priority. .. py:method:: pytorch3d(include_textures=True) -> pytorch3d.structures.Meshes Converts the mesh to a PyTorch3D one. For the base class there are no textures, but subclasses may include them. Empty meshes are augmented with a single degenerate face on conversion, since PyTorch3D does not handle empty meshes correctly. .. py:method:: unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module) :classmethod: Load a mesh of this type from the given file. .. py:method:: serialize() .. py:method:: _deserialize_tensors(data: Dict) -> Dict :classmethod: Convert list attributes to tensor attributes if there is any tensor attribute .. py:method:: empty(dim=2, batch_size=1) :classmethod: Create empty mesh. .. py:method:: trim(polygon: torch.Tensor, trim_face_only=False) Crops the mesh to a given 2D convex polygon, returning a new mesh. Faces where all vertices are outside the polygon are removed, even if the triangle intersects with the polygon. Vertices are removed if they are not used by any face. :param polygon: BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion :param trim_face_only: whether to keep all vertices, including those unused by any faces .. py:method:: separate_by_category() -> Dict[str, BaseMesh] Splits the mesh into meshes representing different categories. .. py:class:: BirdviewRGBMeshGenerator(background_mesh: BirdviewMesh, color_map: Dict[str, Tuple[int, int, int]], rendering_levels: Dict[str, float], world_center: Optional[torch.Tensor] = None, agent_attributes: Optional[torch.Tensor] = None, agent_types: Optional[torch.Tensor] = None, agent_type_names: Optional[List[str]] = None, render_agent_direction: bool = True, traffic_controls: Optional[Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl]] = None, waypoint_radius: float = 2.0, waypoint_num_triangles: int = 10) A generator for constructing simulation birdview representations. All meshes (backround mesh, agent mesh and traffic control mesh) are constructed as templates once at the creation of the generator. Calling the generate function will transform each mesh to a new position/state and return a single RGB mesh per camera for rendering. .. py:attribute:: color_map .. py:attribute:: rendering_levels .. py:attribute:: actor_mesh :value: None .. py:attribute:: static_traffic_controls_mesh :value: None .. py:attribute:: traffic_lights_mesh :value: None .. py:attribute:: traffic_light_colors :value: None .. py:method:: to(device: torch.device) Moves the renderer to another device in place. .. py:method:: copy() .. py:method:: expand(n: int) Adds another dimension with size n on the right of the batch dimension and flattens them. Returns a new renderer, without modifying the current one. .. py:method:: select_batch_elements(idx: torch.Tensor) Selects given elements from the batch, potentially with repetitions. Returns a new renderer, without modifying the current one. :param idx: one-dimensional integer tensor .. py:method:: initialize_background_mesh(background_mesh: BirdviewMesh, world_center: Optional[torch.Tensor] = None) .. py:method:: add_static_meshes(meshes: List[BirdviewMesh]) -> None Includes additional static elements to background mesh. .. py:method:: add_static_rgb_meshes(meshes: List[RGBMesh]) -> None Includes additional static rgb elements to background mesh. .. py:method:: initialize_waypoint_mesh(waypoint_radius: float = 2.0, waypoint_num_triangles: int = 10) .. py:method:: _make_waypoint_mesh(batch_size: int, radius: float = 2.0, num_triangles: int = 10, device: torch.device = torch.device('cpu')) -> BirdviewMesh :classmethod: Create a mesh of the given waypoints. :param batch_size: int number of batch waypoint meshes to create :param radius: float radius of the disc :param num_triangles: int number of triangles used for the disc :param device: torch.device to create the mesh .. py:method:: _make_direction_mesh(lenwid: torch.Tensor, size: float = 0.3, device: torch.device = torch.device('cpu')) -> BaseMesh :classmethod: Create a mesh indicating the direction of each agent. :param lenwid: BxAx2 tensor specifying length and width of the agents :param size: determines the size of the triangle indicating direction .. py:method:: _make_actors_mesh(agent_attributes: torch.Tensor, agent_types: torch.Tensor, agent_type_names: List[str], render_agent_direction: bool = True, device: torch.device = torch.device('cpu')) -> BirdviewMesh :classmethod: .. py:method:: initialize_actors_mesh(agent_attributes: torch.Tensor, agent_types: torch.Tensor, agent_type_names: List[str], render_agent_direction: bool = True) .. py:method:: _create_traffic_controls_mesh(traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl], selected_traffic_control_types: Optional[List[str]] = None) -> BirdviewMesh :classmethod: Create a mesh showing traffic controls. .. py:method:: initialize_traffic_controls_mesh(traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl]) .. py:method:: generate(num_cameras: int, agent_state: Optional[torch.Tensor] = None, present_mask: Optional[torch.Tensor] = None, traffic_lights: Optional[torchdrivesim.traffic_controls.TrafficLightControl] = None, waypoints: Optional[torch.Tensor] = None, waypoints_rendering_mask: Optional[torch.Tensor] = None, custom_agent_colors: Optional[torch.Tensor] = None) -> RGBMesh Create an RGB mesh updates given the provided states. :param num_cameras: int the number of cameras :param agent_state: BxNcxAx4 tensor specifying :param present_mask: BxNcxA tensor specifying :param traffic_lights: TrafficLightControl object extended for each camera :param waypoints: BxNcxMx2 tensor of `M` waypoints per camera (x,y) :param waypoints_rendering_mask: BxNcxM tensor of `M` waypoint masks per camera, indicating which waypoints should be rendered :param custom_agent_colors: a BxNcxAx3 tensor of specifying what color each agent is to what camera .. py:function:: rendering_mesh(mesh: BaseMesh, category: str) -> BirdviewMesh Assigns a category to a given mesh. .. py:function:: set_colors_with_defaults(mesh: BirdviewMesh, color_map: Dict[str, torch.Tensor], rendering_levels: Dict[str, float]) -> RGBMesh .. py:function:: generate_trajectory_mesh(points: torch.Tensor, category: Optional[str] = None, edge_length: float = 1) -> BaseMesh Create a triangle mesh used to visualize a given trajectory. Each point is converted to a triangle matching its position and orientation. :param points: Bx3 tensor of x-y coordinates and orientations in radians :param category: if specified, produces BirdviewMesh :param edge_length: specifies the size of the resulting triangle .. py:function:: generate_annulus_polygon_mesh(polygon: torch.Tensor, scaling_factor: float, origin: torch.Tensor, category: Optional[str] = None) -> BaseMesh For a given polygon, generates a mesh covering the space between the polygon and its scaled version. :param polygon: tensor of size Nx2 defining subsequent points of the polygon hull :param scaling_factor: determines the side of the annulus, should be larger than 1 :param origin: tensor of size (2,) defining the point around which scaling is performed :param category: if specified, BirdviewMesh will be returned .. py:function:: generate_disc_mesh(radius: float = 2, num_triangles: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, torch.Tensor] For a given radius, it will create a disc mesh using `num_triangles` triangles. :param radius: float defining the radius of the disc :param num_triangles: int defining the number of triangles to be used for creating the disc :param device: the device to be used for the generated PyTorch tensors .. py:function:: build_verts_faces_from_bounding_box(bbs: torch.Tensor, z: float = 2) -> Tuple[torch.Tensor, torch.Tensor] Triangulates bounding boxes for rendering. Input is a tensor of bounding boxes of shape ...xAx4x2, where A is the number of actors. Outputs are shaped ...x4*Ax2 and ...x2*Ax3 respectively.