Fixed frustum device and updated model checkpoint defaults (#78)

* fix tensor device and update default model save dir

* removed viewer reference

configs/graph_default.yaml

@@ -26,8 +26,8 @@ logging:
   enable_profiler: True
 
 trainer:
-  model_dir: /tmp/pyrad_models/
-  steps_per_save: 100
+  model_dir: "./checkpoints"
+  steps_per_save: 1000
   steps_per_test: 1000
   max_num_iterations: 1000000
   resume_train:

pyrad/cameras/cameras.py

@@ -366,7 +366,7 @@ def get_camera_model(num_intrinsics_params: int) -> Type[Camera]:
     raise NotImplementedError
 
 
-def get_camera(intrinsics, camera_to_world, camera_index):
+def get_camera(intrinsics, camera_to_world, camera_index=None):
     assert len(intrinsics.shape) == 1, "The intrinsics object should be a flat tensor."
     num_intrinsics_params = len(intrinsics)
     camera_class = get_camera_model(num_intrinsics_params)

pyrad/cameras/rays.py

@@ -69,18 +69,18 @@ class Frustums:
         )
 
     @classmethod
-    def get_mock_frustum(cls) -> "Frustums":
+    def get_mock_frustum(cls, device="cpu") -> "Frustums":
         """Helper function to generate a placeholder frustum.
 
         Returns:
             Frustums: A size 1 frustum with meaningless values.
         """
         return Frustums(
-            origins=torch.ones((1, 3)),
-            directions=torch.ones((1, 3)),
-            frustum_starts=torch.ones((1, 1)),
-            frustum_ends=torch.ones((1, 1)),
-            pixel_area=torch.ones((1, 1)),
+            origins=torch.ones((1, 3)).to(device),
+            directions=torch.ones((1, 3)).to(device),
+            frustum_starts=torch.ones((1, 1)).to(device),
+            frustum_ends=torch.ones((1, 1)).to(device),
+            pixel_area=torch.ones((1, 1)).to(device),
         )
 
     def apply_masks(self, mask: TensorType) -> "Frustums":

pyrad/fields/base.py

@@ -82,7 +82,7 @@ class Field(nn.Module):
         if is_not_none(valid_mask):
             # Hacky handling of empty masks. Tests on a single ray but doesn't use results
             if not valid_mask.any():
-                ray_samples = RaySamples(frustums=Frustums.get_mock_frustum())
+                ray_samples = RaySamples(frustums=Frustums.get_mock_frustum(device=valid_mask.device))
             else:
                 ray_samples = ray_samples.apply_masks()
             density_masked, density_embedding_masked = self.get_density(ray_samples)