From ba08a8f530960486b9f21e8fcd7e591d6209ae86 Mon Sep 17 00:00:00 2001 From: Mingwu Zheng Date: Sun, 28 May 2023 05:00:40 +0800 Subject: [PATCH] Add Metadata Handling to `Cameras` Collation Function (#2001) * Add metadata collate for cameras * checks passed --------- Co-authored-by: Matthew Tancik --- nerfstudio/data/utils/nerfstudio_collate.py | 8 ++++++++ 1 file changed, 8 insertions(+) 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: