aitviewer Documentation

aitviewer is a toolset for visualizing and interacting with sequences of 3D data. Key capabilities include:

• Parametric Model Support: Load and display SMPL[-H/-X], MANO, FLAME, STAR, and SUPR sequences.
• Rendering Modes:
  • Interactive Viewer: Built-in extensible GUI based on Dear ImGui.
  • Headless Mode: For server-side rendering of videos and images.
  • Remote Mode: For non-blocking integration of visualization code.
• Advanced Visualization:
  • Render 3D data on top of images using weak-perspective or OpenCV camera models.
  • Animatable camera paths.
  • Prebuilt renderable primitives (cylinders, spheres, point clouds, etc.).
• Editing & Export:
  • Manually edit SMPL sequences and poses.
  • Export screenshots, videos, and turntable views (mp4/gif).
• Performance: High-performance ModernGL-based rendering pipeline (100fps+ on most laptops).

The Viewer can act as a server to receive commands from remote clients. This is controlled by the C.server_enabled configuration.

Key Methods for Remote Interaction:

• get_node_by_remote_uid(remote_uid, client): Retrieves a local Node using a unique ID generated by a remote client.
• process_message(type, remote_uid, args, kwargs, client): The default handler for incoming messages. This can be overridden to intercept or modify messages before they reach the viewer.
• send_message(msg, client=None): Sends a pickled Python object to a specific client (tuple of (host, port)) or to all connected clients if client is None.

When overriding process_message, you can implement custom logic to respond to specific message types sent by your simulation or data generation scripts.

# Example of overriding message processing
class MyCustomViewer(Viewer):
    def process_message(self, type, remote_uid, args, kwargs, client):
        if type == MY_CUSTOM_MSG_TYPE:
            print(f"Received custom message: {args}")
        else:
            super().process_message(type, remote_uid, args, kwargs, client)

The ViewerServer class enables networked visualization by running a WebSocket server in a background thread. It allows remote clients to send messages (via pickle serialization) to control the viewer, such as adding nodes, updating frames, or deleting objects.

To use it, instantiate ViewerServer with an existing viewer instance and a target port. The server runs in a daemon thread, meaning it will be terminated when the main thread exits. You must periodically call process_messages() in your main loop to handle incoming requests from clients.

from aitviewer.server import ViewerServer
from aitviewer.viewer import Viewer

# Initialize the viewer
viewer = Viewer()

# Start the server on port 8080
server = ViewerServer(viewer, port=8080)

# In your main application loop:
try:
    while True:
        # Process any messages received from remote clients
        server.process_messages()
        # ... rest of your viewer loop ...
finally:
    server.close()

The Viewer class is the main entrypoint for the visualization interface. It inherits from moderngl_window.WindowConfig and manages the scene, renderer, and GUI (via ImGui).

Initialization Parameters:

• title (str): The window title.
• size (Tuple[int, int]): Window dimensions (width, height). If None, uses values from the configuration file.
• samples (int): Number of multisample anti-aliasing (MSAA) samples.
• **kwargs: Additional arguments passed to the window configuration.

Key Methods:

• run(*args, log=True): Enters the blocking visualization loop. This initializes the scene and starts the rendering/event loop.
• reset(): Releases the current scene and initializes a new one.
• toggle_animation(run: bool): Enables or disables animation playback.
• set_ortho_grid_viewports(): Sets up a 2x2 grid of viewports showing the scene from each main axis using orthographic projection.

from aitviewer.viewer import Viewer

viewer = Viewer(title="My Scene", size=(1280, 720))
# ... add nodes to viewer.scene ...
viewer.run()

You can generate color-coded or ID-based masks using the following methods:

• get_mask(color_map: Dict[int, Tuple[int, int, int]] = None, id_map: Dict[int, int] = None) -> Image: Returns a color mask as an RGB PIL image. Each object has a uniform color based on its Node UID.
• save_mask(file_path, color_map=None, id_map=None): Saves the color mask to a file.
• get_mask_ids(id_map: Dict[int, int] = None) -> np.ndarray: Returns a mask as a NumPy array of shape (height, width) and type np.uint32. Each element is the UID of the node covering that pixel (or zero if empty).

Mapping Parameters:

• id_map: A dictionary mapping Node UIDs to specific IDs. This is applied before color mapping or hashing.
• color_map: A dictionary mapping IDs (after id_map application) to (R, G, B) integer tuples (0-255). If None, colors are computed by hashing the UID.

# Define a custom mapping for specific objects
# Map Node UID 123 to ID 1, and Node UID 456 to ID 2
custom_id_map = {123: 1, 456: 2}

# Map ID 1 to Red and ID 2 to Blue
custom_color_map = {1: (255, 0, 0), 2: (0, 0, 255)}

# Save the color mask
renderer.save_mask("mask.png", color_map=custom_color_map, id_map=custom_id_map)

# Get raw IDs as a numpy array
ids = renderer.get_mask_ids(id_map=custom_id_map)