pixelSplat

When providing camera data to pixelSplat, follow these conventions:

• Extrinsics: OpenCV-style camera-to-world matrices.
  • +Z: Camera look vector
  • +X: Camera right vector
  • -Y: Camera up vector
• Intrinsics: Normalized intrinsics. The first row is divided by image width, and the second row is divided by image height.

Pre-trained checkpoints can be downloaded from this Google Drive folder.

Note that original checkpoints (from before the camera-ready architectural improvements) are available separately here.

To install pixelSplat, create a Python 3.10+ virtual environment and install the core dependencies. If you are on Ubuntu, ensure python3.11-dev is installed.

python3.10 -m venv venv
source venv/bin/activate
# Install these first!
pip install wheel torch torchvision torchaudio
pip install -r requirements.txt

python3.10 -m venv venv
source venv/bin/activate
# Install these first! Also, make sure you have python3.11-dev installed if using Ubuntu.
pip install wheel torch torchvision torchaudio
pip install -r requirements.txt

pixelSplat uses versions of the RealEstate10k and ACID datasets split into ~100 MB chunks.

1. Download the chunks from the provided Google Drive link.
2. Unzip the files into a datasets folder in the project root directory.

Alternatively, you can find preprocessed versions at http://schadenfreude.csail.mit.edu:8000/.

A DatasetShim is a callable used to wrap or transform a dataset before it is passed to a DataLoader. This is useful for applying stage-specific logic (e.g., applying different augmentations for 'train' vs 'val') or wrapping the dataset in a decorator.

Signature: DatasetShim = Callable[[Dataset, Stage], Dataset]

Example usage when initializing DataModule:

def my_custom_shim(dataset, stage):
    if stage == "train":
        return SomeAugmentedDataset(dataset)
    return dataset

module = DataModule(..., dataset_shim=my_custom_shim)

Training is performed using src/main.py. The default re10k experiment requires a single GPU with 80 GB of VRAM (e.g., A100 or H100).

To train with the re10k configuration:

python3 -m src.main +experiment=re10k

To reduce memory usage, you can decrease the batch size (the value provided is per-GPU):

python3 -m src.main +experiment=re10k data_loader.train.batch_size=1

Multi-GPU training is supported.

Ablations can be run by specifying the corresponding experiment configuration. For example, to run the ablation without the epipolar transformer:

python3 -m src.main +experiment=re10k_ablation_no_epipolar_transformer

Pre-trained checkpoints for these ablations are included in the main checkpoint repository.

The DataLoaderCfg and DataLoaderStageCfg classes define the configuration for training, testing, and validation data loaders. You can specify batch sizes, worker counts, and seeds for each stage to ensure reproducible and efficient data loading.

DataLoaderStageCfg fields:

• batch_size: Number of samples per batch.
• num_workers: Number of subprocesses to use for data loading.
• persistent_workers: Whether to keep workers in memory between epochs.
• seed: Integer seed for the random number generator (set to None for no seed).

@dataclass
class DataLoaderStageCfg:
    batch_size: int
    num_workers: int
    persistent_workers: bool
    seed: int | None

@dataclass
class DataLoaderCfg:
    train: DataLoaderStageCfg
    test: DataLoaderStageCfg
    val: DataLoaderStageCfg

Use sample_image_grid() to generate normalized (0 to 1) coordinates and integer pixel indices for a given image shape.

Parameters:

• shape: tuple[int, ...] (The dimensions of the image, e.g., (H, W)).
• device: torch.device (The device to create tensors on).

Returns:

• coordinates: Float[Tensor, "*shape dim"] (Normalized (x, y) coordinates using xy indexing).
• indices: Int64[Tensor, "*shape dim"] (Integer (row, col) indices using ij indexing).

The DataModule class (inheriting from LightningDataModule) manages the lifecycle of datasets and dataloaders for training, validation, and testing. It integrates with get_dataset to instantiate datasets based on a DatasetCfg and applies an optional dataset_shim to modify the dataset behavior.

To instantiate a DataModule, you need:

• dataset_cfg: A DatasetCfg object defining the dataset.
• data_loader_cfg: A DataLoaderCfg object defining loader parameters.
• step_tracker (optional): A StepTracker instance.
• dataset_shim (optional): A callable with signature (Dataset, Stage) -> Dataset used to wrap or modify the dataset.
• global_rank (optional): Integer used for seed offsets in distributed training.