torchdrivesim.mesh

Custom classes for representing various triangular meshes as tensors.

Attributes

logger
Color

Exceptions

BadMeshFormat

The mesh received had the wrong format, usually indicating loading from old cache.

Classes

BaseMesh	Basic triangle mesh in a space of arbitrary dimensions Dim.
AttributeMesh	Endows each vertex with an attribute.
RGBMesh	AttributeMesh where the attribute is an RGB color in [0,1] range.
BirdviewMesh	2D mesh with vertices and faces assigned to discrete categories, to facilitate rendering simple 2D worlds.
BirdviewRGBMeshGenerator	A generator for constructing simulation birdview representations. All meshes (backround mesh, agent mesh and

Functions

tensor_color(→ torch.Tensor)	Converts all supported color representations to the tensor representation.
rendering_mesh(→ BirdviewMesh)	Assigns a category to a given mesh.
set_colors_with_defaults(→ RGBMesh)
generate_trajectory_mesh(→ BaseMesh)	Create a triangle mesh used to visualize a given trajectory.
generate_annulus_polygon_mesh(→ BaseMesh)	For a given polygon, generates a mesh covering the space between the polygon
generate_disc_mesh(→ Tuple[torch.Tensor, torch.Tensor])	For a given radius, it will create a disc mesh using num_triangles triangles.
build_verts_faces_from_bounding_box(...)	Triangulates bounding boxes for rendering. Input is a tensor of bounding boxes of shape ...xAx4x2,

Module Contents

torchdrivesim.mesh.logger = None

torchdrivesim.mesh.Color

torchdrivesim.mesh.tensor_color(color: Color, device: torch.device = 'cpu', dtype=torch.float) → torch.Tensor

Converts all supported color representations to the tensor representation.

Parameters:

• color – RGB color as either an int 3-tuple in [0,255] range or tensor of shape (3,) in [0,1] range

• device – device to place the resulting tensor on, ignored if color is already a tensor

• dtype – element type for the resulting tensor, ignored if color is already a tensor

Returns:

RGB tensor of shape (3,) in [0,1] range

exception torchdrivesim.mesh.BadMeshFormat

Bases: RuntimeError

The mesh received had the wrong format, usually indicating loading from old cache.

class torchdrivesim.mesh.BaseMesh

Basic triangle mesh in a space of arbitrary dimensions Dim. Only specifies triangles, with no additional properties. Always includes exactly one batch dimension.

verts: torch.Tensor

faces: torch.Tensor

_verts_fill: float = 0.0

_faces_fill: int = 0

__post_init__()

_validate_input()

property dim: int

Dimension of the space in which the mesh lives, usually 2 or 3.

property batch_size: int

property verts_count: int

property faces_count: int

property device: torch.device

property center: 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.

to(device: torch.device)

Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object.

clone()

Deep copy of the mesh.

expand(size: int)

Expands batch dimension on the right by the given factor, returning a new mesh.

select_batch_elements(idx: torch.Tensor)

Selects given indices from batch dimension, possibly with repetitions, returning a new mesh

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).

Parameters:

• xy – Bx2 tensor specifying the camera position

• inplace – whether to modify the mesh in place

__getitem__(item)

classmethod collate(meshes)

Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension.

classmethod concat(meshes)

Concatenates multiple meshes to form a single scene.

merge(other)

Merges the current mesh with another to form a single scene, returning a new mesh.

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.

pad(pad_size: int)

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.

pickle(mesh_file_path: str)

Store this mesh to a given file.

classmethod unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module)

Load a mesh of this type from the given file.

serialize()

save(file_save_path: str)

Save the attributes of the mesh object in a json where tensors are converted to lists.

classmethod _deserialize_tensors(data: Dict) → Dict

Convert list attributes to tensor attributes if there is any tensor attribute

classmethod deserialize(data: Dict) → typing_extensions.Self

classmethod load(filepath)

classmethod empty(dim: int = 2, batch_size: int = 1)

Create empty mesh.

_trim_and_return_verts_and_faces(vertices_to_keep: torch.Tensor, trim_face_only=False)

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.

Parameters:

• polygon – BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion

• trim_face_only – whether to keep all vertices, including those unused by any faces

class torchdrivesim.mesh.AttributeMesh

Bases: 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.

attrs: torch.Tensor

_attrs_fill: float = 0.0

_validate_input()

property attr_dim: int

Size of the attribute dimension.

classmethod set_attr(mesh: BaseMesh, attr: torch.Tensor)

Sets a given attribute value for all vertices in a given mesh.

to(device: torch.device)

Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object.

pad(pad_size: int)

expand(size: int)

Expands batch dimension on the right by the given factor, returning a new mesh.

select_batch_elements(idx: int)

Selects given indices from batch dimension, possibly with repetitions, returning a new mesh

classmethod concat(meshes)

Concatenates multiple meshes to form a single scene.

classmethod collate(meshes)

Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension.

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.

classmethod unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module)

Load a mesh of this type from the given file.

serialize()

classmethod _deserialize_tensors(data: Dict) → Dict

Convert list attributes to tensor attributes if there is any tensor attribute

classmethod load(filepath)

classmethod empty(dim=2, batch_size=1, attr_dim=3)

Create empty mesh.

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.

Parameters:

• polygon – BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion

• trim_face_only – whether to keep all vertices, including those unused by any faces

class torchdrivesim.mesh.RGBMesh

Bases: AttributeMesh

AttributeMesh where the attribute is an RGB color in [0,1] range.

_validate_input()

classmethod set_color(mesh: BaseMesh, color: Color)

Sets a constant color for all vertices in a given mesh.

class torchdrivesim.mesh.BirdviewMesh

Bases: 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.

categories: List[str]

colors: Dict[str, torch.Tensor]

zs: Dict[str, float]

vert_category: torch.Tensor

_cat_fill: int = 0

_validate_input()

property num_categories: int

classmethod set_properties(mesh: BaseMesh, category: str, color: Color | None = None, z: float | None = None)

Lifts a BaseMesh into a BirdviewMesh with a single category.

to(device)

Moves all tensors to a given device, returning a new mesh object. Does not modify the existing mesh object.

pad(pad_size: int)

expand(size)

Expands batch dimension on the right by the given factor, returning a new mesh.

select_batch_elements(idx)

Selects given indices from batch dimension, possibly with repetitions, returning a new mesh

classmethod unify(meshes)

Generates meshes equivalent to input meshes, only all sharing the same category definitions.

classmethod concat(meshes)

Concatenates multiple meshes to form a single scene.

classmethod collate(meshes)

Batches a collection of meshes with appropriate padding. All input meshes must have a singleton batch dimension.

fill_attr() → RGBMesh

Computes explicit color for each vertex and augments vertices with z values corresponding to rendering priority.

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.

classmethod unpickle(mesh_file_path: str, pickle_module: Any = default_pickle_module)

Load a mesh of this type from the given file.

serialize()

classmethod _deserialize_tensors(data: Dict) → Dict

Convert list attributes to tensor attributes if there is any tensor attribute

classmethod empty(dim=2, batch_size=1)

Create empty mesh.

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.

Parameters:

• polygon – BxPx2 tensor specifying a convex polygon in either clockwise or counterclockwise fashion

• trim_face_only – whether to keep all vertices, including those unused by any faces

separate_by_category() → Dict[str, BaseMesh]

Splits the mesh into meshes representing different categories.

class torchdrivesim.mesh.BirdviewRGBMeshGenerator(background_mesh: BirdviewMesh, color_map: Dict[str, Tuple[int, int, int]], rendering_levels: Dict[str, float], world_center: torch.Tensor | None = None, agent_attributes: torch.Tensor | None = None, agent_types: torch.Tensor | None = None, agent_type_names: List[str] | None = None, render_agent_direction: bool = True, traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl] | None = 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.

color_map

rendering_levels

actor_mesh = None

static_traffic_controls_mesh = None

traffic_lights_mesh = None

traffic_light_colors = None

to(device: torch.device)

Moves the renderer to another device in place.

copy()

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.

select_batch_elements(idx: torch.Tensor)

Selects given elements from the batch, potentially with repetitions. Returns a new renderer, without modifying the current one.

Parameters:

idx – one-dimensional integer tensor

initialize_background_mesh(background_mesh: BirdviewMesh, world_center: torch.Tensor | None = None)

add_static_meshes(meshes: List[BirdviewMesh]) → None

Includes additional static elements to background mesh.

add_static_rgb_meshes(meshes: List[RGBMesh]) → None

Includes additional static rgb elements to background mesh.

initialize_waypoint_mesh(waypoint_radius: float = 2.0, waypoint_num_triangles: int = 10)

classmethod _make_waypoint_mesh(batch_size: int, radius: float = 2.0, num_triangles: int = 10, device: torch.device = torch.device('cpu')) → BirdviewMesh

Create a mesh of the given waypoints.

Parameters:

• batch_size – int number of batch waypoint meshes to create

• radius – float radius of the disc

• num_triangles – int number of triangles used for the disc

• device – torch.device to create the mesh

classmethod _make_direction_mesh(lenwid: torch.Tensor, size: float = 0.3, device: torch.device = torch.device('cpu')) → BaseMesh

Create a mesh indicating the direction of each agent.

Parameters:

• lenwid – BxAx2 tensor specifying length and width of the agents

• size – determines the size of the triangle indicating direction

classmethod _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

initialize_actors_mesh(agent_attributes: torch.Tensor, agent_types: torch.Tensor, agent_type_names: List[str], render_agent_direction: bool = True)

classmethod _create_traffic_controls_mesh(traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl], selected_traffic_control_types: List[str] | None = None) → BirdviewMesh

Create a mesh showing traffic controls.

initialize_traffic_controls_mesh(traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl])

generate(num_cameras: int, agent_state: torch.Tensor | None = None, present_mask: torch.Tensor | None = None, traffic_lights: torchdrivesim.traffic_controls.TrafficLightControl | None = None, waypoints: torch.Tensor | None = None, waypoints_rendering_mask: torch.Tensor | None = None, custom_agent_colors: torch.Tensor | None = None) → RGBMesh

Create an RGB mesh updates given the provided states.

Parameters:

• num_cameras – int the number of cameras

• agent_state – BxNcxAx4 tensor specifying

• present_mask – BxNcxA tensor specifying

• traffic_lights – TrafficLightControl object extended for each camera

• waypoints – BxNcxMx2 tensor of M waypoints per camera (x,y)

• waypoints_rendering_mask – BxNcxM tensor of M waypoint masks per camera, indicating which waypoints should be rendered

• custom_agent_colors – a BxNcxAx3 tensor of specifying what color each agent is to what camera

torchdrivesim.mesh.rendering_mesh(mesh: BaseMesh, category: str) → BirdviewMesh

Assigns a category to a given mesh.

torchdrivesim.mesh.set_colors_with_defaults(mesh: BirdviewMesh, color_map: Dict[str, torch.Tensor], rendering_levels: Dict[str, float]) → RGBMesh

torchdrivesim.mesh.generate_trajectory_mesh(points: torch.Tensor, category: str | None = 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.

Parameters:

• points – Bx3 tensor of x-y coordinates and orientations in radians

• category – if specified, produces BirdviewMesh

• edge_length – specifies the size of the resulting triangle

torchdrivesim.mesh.generate_annulus_polygon_mesh(polygon: torch.Tensor, scaling_factor: float, origin: torch.Tensor, category: str | None = None) → BaseMesh

For a given polygon, generates a mesh covering the space between the polygon and its scaled version.

Parameters:

• polygon – tensor of size Nx2 defining subsequent points of the polygon hull

• scaling_factor – determines the side of the annulus, should be larger than 1

• origin – tensor of size (2,) defining the point around which scaling is performed

• category – if specified, BirdviewMesh will be returned

torchdrivesim.mesh.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.

Parameters:

• radius – float defining the radius of the disc

• num_triangles – int defining the number of triangles to be used for creating the disc

• device – the device to be used for the generated PyTorch tensors

torchdrivesim.mesh.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.