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

> Extend LingBot-Map functionality with custom heads, aggregators, and pipelines. Learn how to subclass components and integrate new sensors without altering the core streaming pipeline.

- Repository: [Robbyant/lingbot-map](https://github.com/Robbyant/lingbot-map)
- Tags: how-to-guide
- Published: 2026-07-19

---

**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`](https://github.com/Robbyant/lingbot-map/blob/main/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.

```python

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

```python

# 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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/models/gct_base.py).

```python
    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.

```python

# 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:

```python
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.

```python

# 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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) or [`benchmark/run.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/lingbot_map/vis/viser_wrapper.py) to check for the new tensor keys in the model output dictionary and add them to the viewer:

```python
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`](https://github.com/Robbyant/lingbot-map/blob/main/demo.py) and rendering logic in [`lingbot_map/vis/viser_wrapper.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/benchmark/datasets/__init__.py) to enable CLI usage via `--dataset YourDataset`. The benchmark runner in [`benchmark/run.py`](https://github.com/Robbyant/lingbot-map/blob/main/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`](https://github.com/Robbyant/lingbot-map/blob/main/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.