diff --git a/nerfstudio/data/utils/nerfstudio_collate.py b/nerfstudio/data/utils/nerfstudio_collate.py index 922b025c0..8d61c695e 100644 --- a/nerfstudio/data/utils/nerfstudio_collate.py +++ b/nerfstudio/data/utils/nerfstudio_collate.py @@ -162,6 +162,13 @@ def nerfstudio_collate(batch: Any, extra_mappings: Union[Dict[type, Callable], N else: op = torch.cat + # Create metadata dictionary + metadata_keys = batch[0].metadata.keys() + assert all( + (cam.metadata.keys() == metadata_keys for cam in batch) + ), "All cameras must have the same metadata keys." + metadata = {key: op([cam.metadata[key] for cam in batch], dim=0) for key in metadata_keys} + return Cameras( op([cameras.camera_to_worlds for cameras in batch], dim=0), op([cameras.fx for cameras in batch], dim=0), @@ -184,6 +191,7 @@ def nerfstudio_collate(batch: Any, extra_mappings: Union[Dict[type, Callable], N [cameras.times if cameras.times is not None else -torch.ones_like(cameras.times) for cameras in batch], dim=0, ), + metadata=metadata, ) for type_key in extra_mappings: