From 2dc1fb6bc9c48df1d6cd0ac822395d0af6926c19 Mon Sep 17 00:00:00 2001
From: Nicholas Pfaff <53228351+nepfaff@users.noreply.github.com>
Date: Tue, 25 Apr 2023 13:04:07 -0400
Subject: [PATCH] Allow masks with instant ngp dataparser (#1741)

* Add masks to instant-ngp-dataparser

* Load entire dataset to GPU if using masks

---------

Co-authored-by: machenmusik <machenmusik@users.noreply.github.com>
---
 nerfstudio/data/dataparsers/instant_ngp_dataparser.py | 5 +++++
 nerfstudio/data/utils/dataloaders.py                  | 3 ++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
index 2fdc2df85..8b4bf724c 100644
--- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
+++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py
@@ -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,
         )
 
diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py
index 23d712d7e..656202c1d 100644
--- a/nerfstudio/data/utils/dataloaders.py
+++ b/nerfstudio/data/utils/dataloaders.py
@@ -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):