updates to dataloader

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -105,7 +105,16 @@ class Nerfstudio(DataParser):
         width = []
         distort = []
 
+        # sort the frames by fname
+        fnames = []
         for frame in meta["frames"]:
+            filepath = PurePath(frame["file_path"])
+            fname = self._get_fname(filepath, data_dir)
+            fnames.append(fname)
+        inds = np.argsort(fnames)
+        frames = [meta["frames"][ind] for ind in inds]
+
+        for frame in frames:
             filepath = PurePath(frame["file_path"])
             fname = self._get_fname(filepath, data_dir)
             if not fname.exists():
@@ -167,16 +176,39 @@ class Nerfstudio(DataParser):
         You should check that mask_path is specified for every frame (or zero frames) in transforms.json.
         """
 
+        # sort the filenames according to image_filenames
+
+
         # filter image_filenames and poses based on train/eval split percentage
         num_images = len(image_filenames)
         num_train_images = math.ceil(num_images * self.config.train_split_percentage)
         num_eval_images = num_images - num_train_images
         i_all = np.arange(num_images)
-        i_train = np.linspace(
-            0, num_images - 1, num_train_images, dtype=int
+        i_eval = np.linspace(
+            0, num_images - 1, num_eval_images, dtype=int
         )  # equally spaced training images starting and ending at 0 and num_images-1
-        i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
+        i_train = np.setdiff1d(i_all, i_eval)  # eval images are the remaining images
+        assert len(i_train) == num_train_images
         assert len(i_eval) == num_eval_images
+
+        # # print(image_filenames)
+        # inds = np.argsort(image_filenames)
+        # image_filenames = [image_filenames[i] for i in inds]
+        # poses = [poses[i] for i in inds]
+        # if len(mask_filenames) != 0:
+        #     mask_filenames = [mask_filenames[i] for i in inds]
+        
+        # mipnerf360 code...
+        all_indices = np.arange(num_images)
+        train_indices = all_indices[all_indices % 8 != 0]
+        eval_indices = all_indices[all_indices % 8 == 0]
+        i_train = train_indices
+        i_eval = eval_indices
+
+        print(image_filenames)
+        print(i_train)
+        print(i_eval)
+        
         if split == "train":
             indices = i_train
         elif split in ["val", "test"]:
@@ -184,6 +216,24 @@ class Nerfstudio(DataParser):
         else:
             raise ValueError(f"Unknown dataparser split {split}")
 
+        # # print("\nhere\n")
+        # # print(i_train)
+        # # print(i_eval)
+
+        # # mipnerf360 code...
+        # eval_interval = 1.0 - self.config.train_split_percentage
+        # all_indices = np.arange(num_images)
+        # train_indices = all_indices[all_indices % 8 != 0]
+        # eval_indices = all_indices[all_indices % 8 == 0]
+        # # print("\nhere\n")
+        # # print(train_indices)
+        # # print(eval_indices)
+
+        # # import sys
+        # # sys.exit()
+
+        # indices
+
         if "orientation_override" in meta:
             orientation_method = meta["orientation_override"]
             CONSOLE.log(f"[yellow] Dataset is overriding orientation method to {orientation_method}")