How hengyiwang/co-slam works
Purpose and Entry Points
Co-SLAM is a Python/PyTorch neural RGB-D SLAM system for real-time camera tracking and dense scene reconstruction using joint coordinate and sparse parametric encodings. Run the single-process pipeline through coslam.py; use coslam_mp.py for separated multiprocessing tracking/mapping. Both accept a scene YAML configuration.
Configuration and Dataset Boundary
Experiment YAMLs define dataset-specific paths, camera/scene parameters, and SLAM settings, with shared and scene-level presets such as configs/Replica/replica.yaml, configs/Replica/office0.yaml, and configs/ScanNet/scannet.yaml. Dataset loading is centralized behind datasets/dataset.py, with dataset-format utilities in datasets/utils.py; this isolates Replica, ScanNet, TUM, Azure/iPhone, and synthetic RGB-D inputs from the core SLAM logic.
Neural Scene Representation
The learned map is implemented in model/scene_rep.py, combining neural decoders (model/decoder.py) with coordinate/parametric encodings (model/encodings.py). This representation consumes sampled RGB-D observations and queried 3D positions, producing scene quantities used for rendering, alignment, optimization, and reconstruction.
Tracking, Mapping, and Keyframes
The single-process coordinator in coslam.py owns the sequential SLAM lifecycle. The multiprocessing variant divides responsibilities into a tracker (mp_slam/tracker.py) and mapper (mp_slam/mapper.py), creating a concurrency boundary between per-frame pose estimation and map/keyframe updates. Keyframe selection/storage is represented by model/keyframe.py, enabling mapping to optimize selected historical observations rather than every frame.
Pose and Alignment Optimization
Camera poses are managed through optimization/pose_array.py. Frame-to-map alignment and tracking losses are concentrated in optimization/frame_align.py, while shared optimization operations reside in optimization/utils.py. Control flow is RGB-D frame ingestion → pose initialization/alignment against the current neural scene → keyframe/map update → refined pose and representation state.
Reconstruction and Outputs
The trained implicit scene is converted to a dense surface via the bundled marching-cubes extension under external/NumpyMarchingCubes/marching_cubes/, built from external/NumpyMarchingCubes/setup.py. General pipeline functions are shared in utils.py; trajectory-related inspection/evaluation utilities include tools/eval_ate.py and tools/vis_cameras.py.
Runtime and External Dependencies
The system targets Python 3.7, PyTorch 1.10.1 with CUDA 11.3, and GPU execution. Its material acceleration dependency is NVIDIA tiny-cuda-nn for neural encoding computation; PyTorch3D is also required. Surface extraction relies on a compiled C++/Cython NumPy marching-cubes module. Input data is local RGB-D dataset storage, populated by scripts such as scripts/download_replica.sh and scripts/download_tum.sh.
Components 18
Runtime
- Scene YAML configs · Dataset paths, camera settings, and SLAM parameters.
configs/Replica/replica.yaml - Single-process pipeline · Owns the sequential SLAM lifecycle.
coslam.py - Multiprocess pipeline · Starts separated tracking and mapping.
coslam_mp.py - Dataset loader · Normalizes supported RGB-D dataset formats for SLAM.
datasets/dataset.py - Dataset format utilities
datasets/utils.py
Core SLAM
- Neural scene representation · Produces scene quantities from observations and queried 3D positions.
model/scene_rep.py - Neural decoders
model/decoder.py - Coordinate encodings · Uses tiny-cuda-nn accelerated encodings.
model/encodings.py - Pose state
optimization/pose_array.py - Frame alignment · Aligns incoming RGB-D frames to the current neural map.
optimization/frame_align.py - Optimization utilities
optimization/utils.py - Keyframe store · Selects observations retained for map optimization.
model/keyframe.py
Multiprocessing
Reconstruction & Evaluation
- Marching-cubes extension · Converts the implicit scene to a dense surface.
- Extension build
external/NumpyMarchingCubes/setup.py - Trajectory evaluation
tools/eval_ate.py - Camera visualization
tools/vis_cameras.py
Connections 25
- Scene YAML configs → Single-process pipeline: configures
- Scene YAML configs → Multiprocess pipeline: configures
- Single-process pipeline → Dataset loader: loads RGB-D frames
- Dataset format utilities → Dataset loader: supports formats
- Dataset loader → Frame alignment: RGB-D observations
- Single-process pipeline → Frame alignment: coordinates tracking
- Frame alignment → Pose state: refines poses
- Pose state → Frame alignment: pose initialization
- Frame alignment → Neural scene representation: queries map
- Neural scene representation → Frame alignment: rendered scene quantities
- Neural decoders → Neural scene representation: compose
- Coordinate encodings → Neural scene representation: encode coordinates
- Optimization utilities → Frame alignment: shared operations
- Single-process pipeline → Keyframe store: updates selected frames
- Keyframe store → Neural scene representation: mapping observations
- Multiprocess pipeline → Tracker: starts
- Multiprocess pipeline → Mapper: starts
- Dataset loader → Tracker: incoming frames
- Tracker → Mapper: poses and frame updates · Concurrency boundary between tracking and mapping.
- Mapper → Keyframe store: manages
- Mapper → Neural scene representation: optimizes map
- Neural scene representation → Marching-cubes extension: implicit field
- Extension build → Marching-cubes extension: builds
- Pose state → Trajectory evaluation: trajectory
- Pose state → Camera visualization: camera poses