The splat trainer, run by cg.splats.train inside the GPU interpreter.
Only torch, nerfstudio and numpy are imported here, never capturegraph: the
GPU environment is a separate installation, so the stored Array[PosedImage]
is read by its on-disk layout. The run writes progress.json twice a second
and the trained gaussians as a PLY the SplatViewer reads.
python trainer.py --frames F --output O --iterations N \\
--seed S.ply --splat OUT.ply --progress P.json
reporting
Writes {"step", "total", "loss", "gaussians"} to progress while the block trains.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| class reporting:
"""Writes ``{"step", "total", "loss", "gaussians"}`` to ``progress`` while the block trains."""
def __init__(self, trainer: Trainer, total: int, progress: Path) -> None:
"""Report ``trainer``'s run of ``total`` steps into ``progress``."""
self._trainer = trainer
self._total = total
self._progress = progress
self._loss: float | None = None
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
def __enter__(self) -> reporting:
"""Start reporting; every training iteration records its loss."""
iterate = self._trainer.train_iteration
def recorded(step: int) -> tuple[object, dict[str, object], dict[str, object]]:
loss, loss_dict, metrics = iterate(step)
self._loss = float(loss)
return loss, loss_dict, metrics
self._trainer.train_iteration = recorded
self._thread.start()
return self
def __exit__(self, *_: object) -> None:
"""Stop reporting and write the final count."""
self._stop.set()
self._thread.join()
self._write(self._total)
def _run(self) -> None:
while not self._stop.wait(REPORT_SECONDS):
self._write(int(getattr(self._trainer, "step", 0)))
def _write(self, step: int) -> None:
payload = {
"step": min(step, self._total),
"total": self._total,
"loss": self._loss,
"gaussians": int(self._trainer.pipeline.model.num_points),
"written_at": time.time(),
}
staging = self._progress.with_suffix(".tmp")
staging.write_text(json.dumps(payload))
os.replace(staging, self._progress)
|
__enter__()
Start reporting; every training iteration records its loss.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def __enter__(self) -> reporting:
"""Start reporting; every training iteration records its loss."""
iterate = self._trainer.train_iteration
def recorded(step: int) -> tuple[object, dict[str, object], dict[str, object]]:
loss, loss_dict, metrics = iterate(step)
self._loss = float(loss)
return loss, loss_dict, metrics
self._trainer.train_iteration = recorded
self._thread.start()
return self
|
__exit__(*_)
Stop reporting and write the final count.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def __exit__(self, *_: object) -> None:
"""Stop reporting and write the final count."""
self._stop.set()
self._thread.join()
self._write(self._total)
|
__init__(trainer, total, progress)
Report trainer's run of total steps into progress.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def __init__(self, trainer: Trainer, total: int, progress: Path) -> None:
"""Report ``trainer``'s run of ``total`` steps into ``progress``."""
self._trainer = trainer
self._total = total
self._progress = progress
self._loss: float | None = None
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
|
build_trainer(data, output, iterations, *, optimize_cameras)
A splatfacto trainer over data in the capture's own metric frame, degree-0 colour.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def build_trainer(data: Path, output: Path, iterations: int, *, optimize_cameras: bool) -> Trainer:
"""A splatfacto trainer over ``data`` in the capture's own metric frame, degree-0 colour."""
from nerfstudio.configs.method_configs import ( # pyright: ignore[reportMissingImports]
all_methods,
)
config = copy.deepcopy(all_methods["splatfacto"])
model = config.pipeline.model
model.sh_degree = 0
model.stop_split_at = iterations // 2
model.rasterize_mode = "classic"
model.camera_optimizer.mode = "SO3xR3" if optimize_cameras else "off"
datamanager = config.pipeline.datamanager
datamanager.cache_images = "cpu"
datamanager.cache_images_type = "uint8"
dataparser = datamanager.dataparser
dataparser.data = data
dataparser.load_3D_points = True
dataparser.orientation_method = "none"
dataparser.center_method = "none"
dataparser.auto_scale_poses = False
dataparser.downscale_factor = 1
dataparser.eval_mode = "all"
never = iterations + 1
config.output_dir = output
config.max_num_iterations = iterations
config.steps_per_eval_image = never
config.steps_per_eval_all_images = never
config.steps_per_eval_batch = never
config.steps_per_save = iterations
config.logging.steps_per_log = STEPS_PER_LOG
config.vis = "tensorboard"
config.set_timestamp()
return config.setup(local_rank=0, world_size=1)
|
camera_to_world(pose)
The 4×4 camera-to-world matrix of a stored Pose.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def camera_to_world(pose: dict) -> np.ndarray:
"""The 4×4 camera-to-world matrix of a stored ``Pose``."""
w, x, y, z = (pose[f"quaternion_{axis}"] for axis in "wxyz")
norm = math.sqrt(w * w + x * x + y * y + z * z) or 1.0
w, x, y, z = w / norm, x / norm, y / norm, z / norm
matrix = np.eye(4)
matrix[:3, :3] = [
[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
[2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
[2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
]
matrix[:3, 3] = [pose["position_x"], pose["position_y"], pose["position_z"]]
return matrix
|
main()
Train on the dataset the arguments name and write the splat.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def main() -> None:
"""Train on the dataset the arguments name and write the splat."""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--frames", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--iterations", type=int, required=True)
parser.add_argument("--seed", type=Path, required=True)
parser.add_argument("--splat", type=Path, required=True)
parser.add_argument("--progress", type=Path, required=True)
parser.add_argument("--optimize-cameras", action="store_true")
arguments = parser.parse_args()
data = stage_dataset(arguments.frames, arguments.output / "data", arguments.seed)
trainer = build_trainer(
data,
arguments.output / "runs",
arguments.iterations,
optimize_cameras=arguments.optimize_cameras,
)
trainer.setup()
with reporting(trainer, arguments.iterations, arguments.progress):
trainer.train()
write_splat(trainer.pipeline.model, arguments.splat)
|
stage_dataset(frames, data, seed)
Nerfstudio's transforms.json under data for the Array[PosedImage] at frames.
The stored array is one directory per frame holding image.<ext> beside
pose.json, the pose being a capturegraph Pose (position, quaternion
and intrinsics at the image's size, in the OpenGL camera-to-world
convention); seed is the PLY the gaussians start from.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def stage_dataset(frames: Path, data: Path, seed: Path) -> Path:
"""Nerfstudio's ``transforms.json`` under ``data`` for the ``Array[PosedImage]`` at ``frames``.
The stored array is one directory per frame holding ``image.<ext>`` beside
``pose.json``, the pose being a capturegraph ``Pose`` (position, quaternion
and intrinsics at the image's size, in the OpenGL camera-to-world
convention); ``seed`` is the PLY the gaussians start from.
"""
records = []
for frame in sorted(path for path in frames.iterdir() if path.is_dir()):
pose = json.loads((frame / "pose.json").read_text())
(image,) = [path for path in frame.iterdir() if path.stem == "image"]
records.append(
{
"file_path": str(image.resolve()),
"w": int(pose["reference_width"]),
"h": int(pose["reference_height"]),
"fl_x": pose["focal_x"],
"fl_y": pose["focal_y"],
"cx": pose["center_x"],
"cy": pose["center_y"],
"transform_matrix": camera_to_world(pose).tolist(),
}
)
data.mkdir(parents=True, exist_ok=True)
(data / "transforms.json").write_text(
json.dumps(
{
"camera_model": "OPENCV",
"ply_file_path": str(seed.resolve()),
"frames": records,
},
indent=1,
)
)
return data
|
write_splat(model, splat)
The model's gaussians as a binary PLY of the standard 3DGS properties.
Source code in capturegraph-lib/capturegraph/recipes/splats/trainer.py
| def write_splat(model: SplatfactoModel, splat: Path) -> None:
"""The model's gaussians as a binary PLY of the standard 3DGS properties."""
positions = model.means.detach().cpu().numpy()
harmonics_degree_zero = model.features_dc.detach().cpu().numpy()
scales = model.scales.detach().cpu().numpy()
quaternions = model.quats.detach().cpu().numpy()
opacities = model.opacities.detach().cpu().numpy().reshape(-1, 1)
columns = np.concatenate(
[
positions,
np.zeros_like(positions),
harmonics_degree_zero,
opacities,
scales,
quaternions,
],
axis=1,
).astype(np.float32)
columns = columns[np.isfinite(columns).all(axis=1)]
record = np.zeros(len(columns), dtype=[(name, "<f4") for name in PLY_PROPERTIES])
for index, name in enumerate(PLY_PROPERTIES):
record[name] = columns[:, index]
header = (
"ply\nformat binary_little_endian 1.0\n"
f"element vertex {len(columns)}\n"
+ "".join(f"property float {name}\n" for name in PLY_PROPERTIES)
+ "end_header\n"
)
splat.write_bytes(header.encode("ascii") + record.tobytes())
|