fixing eval dataloader indices issue (#196)

* added some logic for the indices when none

* forgot to comment out code

* fix linting

* fixed typing

pyrad/data/dataloader.py

@@ -17,7 +17,7 @@ Data loader.
 """
 
 import random
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 from omegaconf import ListConfig
 
 from torchtyping import TensorType
@@ -142,12 +142,20 @@ class FixedIndicesEvalDataloader(EvalDataloader):
         intrinsics,
         camera_to_world,
         num_rays_per_chunk: int,
-        image_indices: Union[List[int], ListConfig],
+        image_indices: Optional[Union[List[int], ListConfig]] = None,
         device="cpu",
         **kwargs,
     ):
+        """
+        Args:
+            image_dataset: ImageDataset to load data from
+            image_indices: List of image indices to load data from. If None, then use all images.
+        """
         super().__init__(image_dataset, intrinsics, camera_to_world, num_rays_per_chunk, device, **kwargs)
-        self.image_indices = image_indices
+        if image_indices is None:
+            self.image_indices = list(range(len(image_dataset)))
+        else:
+            self.image_indices = image_indices
         self.num_rays_per_chunk = num_rays_per_chunk
         self.count = 0
 

scripts/benchmarking/run_benchmark.py

@@ -41,6 +41,7 @@ def _load_best_ckpt(hydra_dir: str, config: DictConfig) -> str:
     step = os.path.splitext(latest_ckpt)[0].split("-")[-1]
     config.resume_train.load_dir = model_dir
     config.resume_train.load_step = int(step)
+    config.data.dataloader_eval.image_indices = None
     return os.path.join(model_dir, latest_ckpt)
 
 

scripts/run_eval.py

@@ -184,6 +184,7 @@ def main(
     config = compose(config_name, overrides=config_overrides)
 
     config.trainer.resume_train.load_dir = checkpoint_dir
+    config.data.dataloader_eval.image_indices = None
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # setup graph and dataset