tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Allow masks with instant ngp dataparser (#1741)

Browse Source 
* Add masks to instant-ngp-dataparser

* Load entire dataset to GPU if using masks

---------

Co-authored-by: machenmusik <machenmusik@users.noreply.github.com>
This commit is contained in:
Nicholas Pfaff 2023-04-25 13:04:07 -04:00 committed by GitHub
parent 231573bd8e
commit 2dc1fb6bc9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
nerfstudio/data/dataparsers/instant_ngp_dataparser.py
View File

@ -65,6 +65,7 @@ class InstantNGP(DataParser):
data_dir = self.config.data
image_filenames = []
mask_filenames = []
poses = []
num_skipped_image_filenames = 0
for frame in meta["frames"]:
@ -87,6 +88,9 @@ class InstantNGP(DataParser):
meta["h"] = h
image_filenames.append(fname)
poses.append(np.array(frame["transform_matrix"]))
if "mask_path" in frame:
mask_fname = data_dir / Path(frame["mask_path"])
mask_filenames.append(mask_fname)
if num_skipped_image_filenames >= 0:
CONSOLE.print(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.")
assert (
@ -144,6 +148,7 @@ class InstantNGP(DataParser):
image_filenames=image_filenames,
cameras=cameras,
scene_box=scene_box,
mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
dataparser_scale=self.config.scene_scale,
)

3
nerfstudio/data/utils/dataloaders.py
View File

@ -116,7 +116,8 @@ class CacheDataloader(DataLoader):
"""Returns a collated batch."""
batch_list = self._get_batch_list()
collated_batch = self.collate_fn(batch_list)
collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
exclude = [] if "mask" in collated_batch.keys() else ["image"] # type: ignore
collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=exclude)
return collated_batch
def __iter__(self):