torchdrivesim.rendering.nvdiffrast

Nvdiffrast-based renderers, equivalent to those based on PyTorch3D but sometimes faster. This module imports correctly if nvdiffrast is missing, but the renderer will raise the NvdiffrastNotFound exception.

Attributes

Exceptions

NvdiffrastNotFound

Nvdiffrast is not installed.

Classes

BirdviewMesh

2D mesh with vertices and faces assigned to discrete categories, to facilitate rendering simple 2D worlds.

RendererConfig

Determines behavior of the renderer.

BirdviewRenderer

A renderer producing simple 2D birdview images based on static background meshes and rectangular agents.

Cameras

Lightweight version of pytorch3d.renderer.FoVOrthographicCameras.

NvdiffrastRendererConfig

Configuration of nvdiffrast-based renderer.

NvdiffrastRenderer

Similar to PyTorch3DRenderer, and producing indistinguishable images, but sometimes faster.

Functions

tensor_color(→ torch.Tensor)

Converts all supported color representations to the tensor representation.

get_glctx_session(device[, opengl])

Module Contents

torchdrivesim.rendering.nvdiffrast.is_available = True[source]
class torchdrivesim.rendering.nvdiffrast.BirdviewMesh[source]

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()[source]
property num_categories: int
classmethod set_properties(mesh: BaseMesh, category: str, color: Color | None = None, z: float | None = None)[source]

Lifts a BaseMesh into a BirdviewMesh with a single category.

to(device)[source]

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

expand(size)[source]

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

select_batch_elements(idx)[source]

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

classmethod unify(meshes)[source]

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

classmethod concat(meshes)[source]

Concatenates multiple meshes to form a single scene.

classmethod collate(meshes)[source]

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

fill_attr() RGBMesh[source]

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

pytorch3d(include_textures=True) pytorch3d.structures.Meshes[source]

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)[source]

Load a mesh of this type from the given file.

serialize()[source]
classmethod _deserialize_tensors(data: Dict) Dict[source]

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

classmethod empty(dim=2, batch_size=1)[source]

Create empty mesh.

trim(polygon: torch.Tensor, trim_face_only=False)[source]

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][source]

Splits the mesh into meshes representing different categories.

torchdrivesim.rendering.nvdiffrast.tensor_color(color: Color, device: torch.device = 'cpu', dtype=torch.float) torch.Tensor[source]

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

class torchdrivesim.rendering.nvdiffrast.RendererConfig[source]

Determines behavior of the renderer. Subclasses determine renderer class used.

backend: str = 'default'
render_agent_direction: bool = True
left_handed_coordinates: bool = False
highlight_ego_vehicle: bool = False
class torchdrivesim.rendering.nvdiffrast.BirdviewRenderer(cfg: RendererConfig, device: torch.device | None = None, batch_size: int | None = None, static_mesh: torchdrivesim.mesh.BirdviewMesh | None = None, world_center: torch.Tensor | None = None, color_map: Dict[str, Tuple[int, int, int]] | None = None, rendering_levels: Dict[str, float] | None = None, res: torchdrivesim.utils.Resolution = Resolution(64, 64), fov: float = 35)[source]

Bases: abc.ABC

A renderer producing simple 2D birdview images based on static background meshes and rectangular agents. Currently only square resolutions are supported. The renderer always operates in batch mode, with a single batch dimension on the left.

Parameters:
  • cfg – configuration object, usually subclassed

  • device – torch device used for rendering

  • batch_size – if road_mesh is not specified, this is used to determine batch size

  • static_mesh – BirdviewMesh object specifying drivable surface (empty mesh is used if not provided)

  • world_center – Bx2 float tensor, defaults to geometric centre of the road mesh

  • color_map – a dictionary of RGB tuples in 0-255 range specifying colors of different rendered elements

  • res – default resolution

get_color(element_type: str) Tuple[int, int, int][source]
to(device: torch.device)[source]

Moves the renderer to another device in place.

add_static_meshes(meshes: List[torchdrivesim.mesh.BirdviewMesh]) None[source]

Includes additional static elements to render.

copy()[source]
expand(n: int)[source]

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)[source]

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

Parameters:

idx – one-dimensional integer tensor

render_static_meshes(camera_xy: torch.Tensor | None = None, camera_sc: torch.Tensor | None = None, res: torchdrivesim.utils.Resolution = None, fov: float = None) torch.Tensor[source]

Render a single birdview of the static mesh only. Nc is the number of cameras. C=3 is the number of RGB channels.

Parameters:
  • camera_xy – Ncx2 tensor of camera positions

  • camera_sc – Ncx2 tensor of camera orientations (sine and cosine of yaw angle)

  • res – Resolution, currently only square resolutions are supported

  • fov – Field of view in meters

Returns:

birdview image tensor of shape NcxHxWxC

transform(points: torch.Tensor, pose: torch.Tensor) torch.Tensor[source]

Given points relative to a pose, produce absolute positions of the points. There can be zero or more batch dimensions.

Parameters:
  • points – BxNx2 tensor

  • pose – Bx3 tensor of position (x,y) and orientation (yaw angle in radians)

Returns:

Bx2 tensor of absolute positions

make_actor_mesh(agent_state: Dict[str, torch.Tensor], agent_attributes: Dict[str, torch.Tensor]) torchdrivesim.mesh.BirdviewMesh[source]

Creates a mesh representing given actors. Each vertex and each face corresponds to a unique agent and both vertices and faces for each agent are continuous in the resulting tensor to allow for subsequent masking. For each agent there are seven vertices and three faces, specifying its bounding box and direction. Direction vertices use the ‘direction’ category, while agent categories are copied from input dictionaries.

render_frame(agent_state: Dict[str, torch.Tensor], agent_attributes: Dict[str, torch.Tensor], camera_xy: torch.Tensor | None = None, camera_sc: torch.Tensor | None = None, rendering_mask: Dict[str, torch.Tensor] = None, res: torchdrivesim.utils.Resolution | None = None, traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl] | None = None, fov: float | None = None, waypoints: torch.Tensor | None = None, waypoints_rendering_mask: torch.Tensor | None = None) torch.Tensor[source]

Renders the agents and traffic controls on top of the static mesh. Cameras batch size is (B*Nc), which corresponds to using Nc cameras per batch element. This extra dimension is added on the right of batch dimension and flattened, to match the semantics of extend from pytorch3d. If cameras is None, one egocentric camera per agent is used, that is Nc = A.

Parameters:
  • agent_state – maps agent types to state tensors of shape BxAxSt, where St >= 3 and the first three components are x coordinate, y coordinate, and orientation in radians

  • agent_attributes – maps agent types to static attributes tensors of shape BxAxAttr, where Attr >= 2 and the first two components are length and width of the agent

  • camera_xy – BxNcx2 tensor of camera positions, by default one camera placed on each agent

  • camera_sc – BxNcx2 tensor of camera orientations (sine and cosine), by default matching agent orientations

  • rendering_mask – BxNcxA tensor per agent type, indicating which cameras see which agents

  • res – resolution HxW of the resulting image, currently only square resolutions are supported

  • traffic_controls – traffic controls by type (traffic_light, yield, etc.)

  • fov – Field of view in meters

  • 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

Returns:

tensor image of float RGB values in [0,255] range with shape shape (B*Nc)xAxCxHxW

abstract render_mesh(mesh: torchdrivesim.mesh.BirdviewMesh, res: torchdrivesim.utils.Resolution, cameras: Cameras) torch.Tensor[source]

Renders a given mesh, producing BxHxWxC tensor image of float RGB values in [0,255] range.

construct_cameras(xy: torch.Tensor, sc: torch.Tensor, scale: float | None = None) Cameras[source]

Create PyTorch3D cameras object for given positions and orientations. Input tensor dimensions should be Bx2.

build_verts_faces_from_bounding_box(bbs: torch.Tensor, z: float = 2) Tuple[torch.Tensor, torch.Tensor][source]

Triangulates actors 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.

make_traffic_controls_mesh(traffic_controls: Dict[str, torchdrivesim.traffic_controls.BaseTrafficControl]) torchdrivesim.mesh.BirdviewMesh[source]

Create a mesh showing traffic controls.

make_direction_mesh(lenwid: torch.Tensor, pose: torch.Tensor, size: float = 0.3) torchdrivesim.mesh.BaseMesh[source]

Create a mesh indicating the direction of each agent.

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

  • pose – Bx3 tensor of position (x,y) and orientation (yaw angle in radians)

  • size – determines the size of the triangle indicating direction

make_waypoint_mesh(waypoints: torch.Tensor, radius: float = 2.0, num_triangles: int = 10) torchdrivesim.mesh.BirdviewMesh[source]

Create a mesh of the given waypoints.

Parameters:
  • waypoints – BxNcxMx3 tensor of M waypoints per camera (x,y,psi)

  • radius – float radius of the disc

  • num_triangles – int number of triangles used for the disc

class torchdrivesim.rendering.nvdiffrast.Cameras(xy: torch.Tensor, sc: torch.Tensor, scale: float)[source]

Lightweight version of pytorch3d.renderer.FoVOrthographicCameras. Used to define an interface that works without pytorch3d.

get_camera_center() torch.Tensor[source]
get_world_to_view_transform() torch.Tensor[source]
get_view_to_proj_transform() torch.Tensor[source]
project_world_to_clip_space(points: torch.Tensor) torch.Tensor[source]
transform_points_screen(points: torch.Tensor, res: torchdrivesim.utils.Resolution) torch.Tensor[source]
reverse_transform_points_screen(points: torch.Tensor, res: torchdrivesim.utils.Resolution) torch.Tensor[source]
torchdrivesim.rendering.nvdiffrast.Resolution[source]
torchdrivesim.rendering.nvdiffrast.logger[source]
torchdrivesim.rendering.nvdiffrast.glctx_sessions[source]
exception torchdrivesim.rendering.nvdiffrast.NvdiffrastNotFound[source]

Bases: ImportError

Nvdiffrast is not installed.

torchdrivesim.rendering.nvdiffrast.get_glctx_session(device, opengl=True)[source]
class torchdrivesim.rendering.nvdiffrast.NvdiffrastRendererConfig[source]

Bases: torchdrivesim.rendering.base.RendererConfig

Configuration of nvdiffrast-based renderer.

backend: str = 'nvdiffrast'[source]
antialias: bool = False[source]
opengl: bool = False[source]
max_minibatch_size: int | None = None[source]
class torchdrivesim.rendering.nvdiffrast.NvdiffrastRenderer(cfg: NvdiffrastRendererConfig, *args, **kwargs)[source]

Bases: torchdrivesim.rendering.base.BirdviewRenderer

Similar to PyTorch3DRenderer, and producing indistinguishable images, but sometimes faster. Note that nvdiffrast requires separate installation and is subject to its own license terms.

render_mesh(mesh: torchdrivesim.mesh.BirdviewMesh, res: torchdrivesim.utils.Resolution, cameras: torchdrivesim.rendering.base.Cameras) torch.Tensor[source]

Renders a given mesh, producing BxHxWxC tensor image of float RGB values in [0,255] range.