fix bilagrid

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -408,6 +408,9 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         )  # Note that this must happen before indexing cameras, as it can modify the cameras in the dataset during undistortion
         cameras = nerfstudio_collate([self.train_dataset.cameras[i : i + 1].to(self.device) for i in camera_indices])
 
+        if cameras.metadata is None:
+            cameras.metadata = {}
+        cameras.metadata["cam_idx"] = torch.tensor(camera_indices, device=self.device, dtype=torch.long)
         return cameras, data
 
     def next_eval(self, step: int) -> Tuple[Cameras, Dict]:

nerfstudio/models/splatfacto.py

@@ -473,7 +473,11 @@ class SplatfactoModel(Model):
             raise ValueError(f"Unknown background color {self.config.background_color}")
         return background
 
-    def _apply_bilateral_grid(self, rgb: torch.Tensor, cam_idx: int, H: int, W: int) -> torch.Tensor:
+    def _apply_bilateral_grid(self, rgb: torch.Tensor, cam_idxs: torch.Tensor, H: int, W: int) -> torch.Tensor:
+        """
+        rgb: [B, H, W, 3]
+        cam_idxs: [B]
+        """
         # make xy grid
         grid_y, grid_x = torch.meshgrid(
             torch.linspace(0, 1.0, H, device=self.device),
@@ -481,14 +485,13 @@ class SplatfactoModel(Model):
             indexing="ij",
         )
         grid_xy = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0)
-
         out = slice(
             bil_grids=self.bil_grids,
-            rgb=rgb,
-            xy=grid_xy,
-            grid_idx=torch.tensor(cam_idx, device=self.device, dtype=torch.long),
+            rgb=rgb,  # Process the entire batch in parallel
+            xy=grid_xy.expand(rgb.shape[0], -1, -1, -1),  # Expand grid_xy to match batch size
+            grid_idx=cam_idxs.unsqueeze(-1),
         )
-        return out["rgb"]
+        return out["rgb"]  # Return the processed RGB directly
 
     def get_outputs(self, cameras: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
         """Takes in cameras and returns a dictionary of outputs.
@@ -598,6 +601,8 @@ class SplatfactoModel(Model):
         if self.config.use_bilateral_grid and self.training:
             if cameras.metadata is not None and "cam_idx" in cameras.metadata:
                 rgb = self._apply_bilateral_grid(rgb, cameras.metadata["cam_idx"], H, W)
+            else:
+                raise ValueError("Camera index not found in metadata, bilateral grid cannot be applied.")
 
         if render_mode == "RGB+ED":
             depth_im = render[:, ..., 3:4]