How to Extend LingBot-Map Functionality: A Developer's Guide to Custom Heads, Aggregators, and Pipelines

Extend LingBot-Map functionality by subclassing modular components—add prediction heads via _build_<head>_head methods, customize temporal aggregation by overriding AggregatorBase, or integrate new sensors through the dataset interface—without modifying the core streaming pipeline.

LingBot-Map is a modular streaming 3-D reconstruction framework designed for extensibility. The codebase in the Robbyant/lingbot-map repository separates concerns into distinct components—base models, streaming inference engines, aggregators, and prediction heads—making it straightforward to extend lingbot-map functionality for custom research needs. Whether you need semantic segmentation outputs, alternative temporal attention mechanisms, or support for additional sensor modalities, you can plug in new code without disrupting the high-performance KV-cache and sliding-window pipeline.

Add a Custom Prediction Head

To extend lingbot-map functionality with new output modalities—such as semantic segmentation, occupancy grids, or lighting estimation—you create a head module following the interface established by DPTHead and CameraHead in lingbot_map/heads/.

Implement the Head Class

Create a new Python file under lingbot_map/heads/ (e.g., semantic_head.py). The class must inherit from nn.Module and implement a forward method that accepts a list of aggregated token tensors and returns a dictionary of predictions.

# lingbot_map/heads/semantic_head.py

import torch.nn as nn
from lingbot_map.layers import Mlp

class SemanticHead(nn.Module):
    """Predict per-pixel semantic scores."""
    def __init__(self, dim_in: int, patch_size: int, num_classes: int = 20):
        super().__init__()
        self.head = Mlp(
            in_features=dim_in,
            hidden_features=dim_in // 2,
            out_features=num_classes,
            drop=0.0,
        )
        self.patch_size = patch_size

    def forward(self, aggregated_tokens_list, images):
        tokens = aggregated_tokens_list[-1]
        B, N, C = tokens.shape
        H = W = int((N - 1) ** 0.5)
        feat = tokens[:, 1:, :].permute(0, 2, 1).view(B, C, H, W)
        logits = self.head(feat)
        return {"semantic_logits": logits}

Register the Head in the Model

Subclass GCTStream (or GCTBase for non-streaming variants) to register your head. Implement _build_<head_name>_head to instantiate the module and store it as an instance attribute.

# lingbot_map/models/gct_stream_extended.py

from lingbot_map.models.gct_stream import GCTStream
from lingbot_map.heads.semantic_head import SemanticHead

class GCTStreamWithSemantic(GCTStream):
    def __init__(self, enable_semantic: bool = True, **kwargs):
        self.enable_semantic = enable_semantic
        super().__init__(**kwargs)

    def _build_semantic_head(self):
        return SemanticHead(
            dim_in=2 * self.embed_dim,
            patch_size=self.patch_size,
            num_classes=21,
        )

Expose the Prediction Output

Override the corresponding _predict_* method to invoke the head and attach its output to the result dictionary. This ensures the output flows through the inference pipeline defined in lingbot_map/models/gct_base.py.

    def _predict_semantic(self, aggregated_tokens_list, images, **kwargs):
        if not self.enable_semantic:
            return {}
        return self.semantic_head(aggregated_tokens_list, images)

Implement a Custom Aggregator Strategy

For research requiring alternative temporal aggregation—such as memory-augmented attention or hierarchical sliding windows—extend AggregatorBase rather than modifying AggregatorStream directly.

Override Key Aggregator Methods

Subclass AggregatorBase and override _build_blocks, forward, and optionally _init_kv_cache. The _build_blocks method constructs per-frame and global attention blocks, while forward implements your custom causal logic. You may reuse KV-cache helpers from the base class.

# lingbot_map/aggregator/custom.py

from lingbot_map.aggregator.base import AggregatorBase

class MyAggregator(AggregatorBase):
    def _build_blocks(self, block_fn, **kwargs):
        self.frame_blocks = nn.ModuleList([block_fn(**kwargs) for _ in range(2)])
        self.global_blocks = nn.ModuleList([block_fn(**kwargs) for _ in range(2)])

    def forward(self, tokens, **kw):
        # Custom causal aggregation logic

        # Return list of token tensors as expected by heads

        pass

Plug the custom aggregator into your model by overriding _build_aggregator in a GCTStream subclass:

class GCTStreamCustomAgg(GCTStream):
    def _build_aggregator(self):
        return MyAggregator(
            sliding_window_size=self.sliding_window_size,
            num_frame_for_scale=self.num_frame_for_scale,
            # Forward kwargs used by AggregatorStream

        )

Extend the Data Pipeline for New Sensors

To extend lingbot-map functionality for new sensor modalities like LiDAR or event cameras, implement a dataset class following the pattern in benchmark/datasets/. The class must yield dictionaries with keys such as image, intrinsics, and pose, plus your new modality.

# benchmark/datasets/my_sensor.py

import torch
from .base import BaseDataset

class MySensorDataset(BaseDataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform)

    def __getitem__(self, idx):
        return {
            "image": self._load_image(self.image_paths[idx]),
            "lidar": self._load_lidar(self.lidar_paths[idx]),
            "intrinsics": self.intrinsics[idx],
            "pose": self.poses[idx],
        }

Register the dataset in benchmark/datasets/__init__.py to enable discovery via the CLI flag --dataset MySensorDataset.

Integrate New Features into the CLI and Visualization

Expose new functionality through command-line arguments in demo.py or benchmark/run.py using argparse. For example, add --enable_semantic to toggle custom heads at runtime.

To visualize new outputs (e.g., semantic masks), modify lingbot_map/vis/viser_wrapper.py to check for the new tensor keys in the model output dictionary and add them to the viewer:

if "semantic_logits" in outputs:
    semantic = outputs["semantic_logits"].argmax(dim=1)
    viewer.add_point_cloud(semantic_colors, name="semantic")

Summary

• Extend lingbot-map functionality by subclassing GCTStream to add prediction heads via _build_<head>_head methods and corresponding _predict_* routines.
• Customize temporal aggregation by inheriting from AggregatorBase and overriding _build_blocks and forward to implement alternative attention mechanisms while preserving KV-cache compatibility.
• Add new sensor support by creating dataset classes in benchmark/datasets/ that return standardized dictionaries, enabling seamless integration with the benchmark CLI.
• Wire features to the interface by updating argument parsers in demo.py and rendering logic in lingbot_map/vis/viser_wrapper.py.

Frequently Asked Questions

How do I add a new prediction head to LingBot-Map?

Create a head class in lingbot_map/heads/ that inherits from nn.Module and returns a dictionary from its forward method. Then subclass GCTStream, implement _build_<head>_head to instantiate it, and override _predict_<head> to expose the output. According to the source code in lingbot_map/models/gct_base.py, the base model automatically detects and calls these builder methods during initialization.

Can I replace the default streaming aggregator with custom temporal logic?

Yes. Subclass AggregatorBase from lingbot_map/aggregator/base.py and override _build_blocks to define your attention mechanism and forward to implement your causal logic. You can reuse the KV-cache management utilities from AggregatorStream in lingbot_map/aggregator/stream.py to maintain FlashInfer or SDPA compatibility. Inject your aggregator by overriding _build_aggregator in a GCTStream subclass.

What is the recommended way to add support for new sensors like LiDAR?

Implement a new dataset class in benchmark/datasets/ that inherits from BaseDataset and returns dictionaries containing your sensor data alongside standard keys like image and pose. Register the class in benchmark/datasets/__init__.py to enable CLI usage via --dataset YourDataset. The benchmark runner in benchmark/run.py expects this interface for all data loaders.

How do I visualize custom outputs in the LingBot-Map demo?

Update lingbot_map/vis/viser_wrapper.py to check for your custom output keys in the model's prediction dictionary. Extract the tensors, convert them to the appropriate format (e.g., colorized labels for segmentation), and pass them to the viewer's add_point_cloud or equivalent method. The wrapper already handles camera poses and depth; follow the existing pattern for new modalities.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}