Temp save

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -80,6 +80,8 @@ class FullImageDatamanagerConfig(DataManagerConfig):
     new images. If -1, never pick new images."""
     eval_image_indices: Optional[Tuple[int, ...]] = (0,)
     """Specifies the image indices to use during eval; if None, uses all."""
+    cache_images: Literal["no-cache", "cpu", "gpu"] = "cpu"
+    """Whether to cache images in memory. If "numpy", caches as numpy arrays, if "torch", caches as torch tensors."""
 
 TDataset = TypeVar("TDataset", bound=InputDataset, default=InputDataset)
 
@@ -126,6 +128,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train")
         self.train_dataset = self.create_train_dataset()
         self.eval_dataset = self.create_eval_dataset()
+        self.cached_train, self.cached_eval = self.cache_images(self.config.cache_images)
         self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device
         if self.config.masks_on_gpu is True:
             self.exclude_batch_keys_from_device.remove("mask")
@@ -140,6 +143,29 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         super().__init__()
 
 
+    def cache_images(self, cache_images_option):
+        if cache_images_option == "no-cache":
+            cached_train = {}
+            cached_eval = {}
+        elif cache_images_option == "cpu":
+            cached_train = [self.train_dataset.get_data(i) for i in range(len(self.train_dataset))]
+            cached_eval = [self.eval_dataset.get_data(i) for i in range(len(self.eval_dataset))]
+        elif cache_images_option == "gpu":
+            cached_train = [self.train_dataset.get_data(i) for i in range(len(self.train_dataset))]
+            cached_eval = [self.eval_dataset.get_data(i) for i in range(len(self.eval_dataset))]
+            for cache in cached_train:
+                cache["image"] = cache["image"].to(self.device)
+                if "mask" in cache:
+                    cache["mask"] = cache["mask"].to(self.device)
+            for cache in cached_eval:
+                cache["image"] = cache["image"].to(self.device)
+                if "mask" in cache:
+                    cache["mask"] = cache["mask"].to(self.device)
+        else:
+            raise ValueError(f"Unknown cache_images option {cache_images_option}")
+
+        return cached_train, cached_eval
+
     def create_train_dataset(self) -> TDataset:
         """Sets up the data loaders for training"""
         return self.dataset_type(
@@ -207,8 +233,19 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         # Make sure to re-populate the unseen cameras list if we have exhausted it
         if len(self.train_unseen_cameras) == 0:
             self.train_unseen_cameras = [i for i in range(len(self.train_dataset))]
-        data = self.train_dataset.get_data(image_idx)
-        data["image"] = data["image"].to(self.device)
+        
+        if self.config.cache_images == "no-cache":
+            data = self.train_dataset.get_data(image_idx)
+            data["image"] = data["image"].to(self.device)
+        elif self.config.cache_images == "cpu":
+            data = self.cached_train[image_idx]
+            data["image"] = data["image"].to(self.device)
+        elif self.config.cache_images == "gpu":
+            data = self.cached_train[image_idx]
+        else:
+            raise ValueError(f"Unknown cache_images option {self.config.cache_images}")
+        
+        
         assert len(self.train_dataset.cameras.shape) == 1, "Assumes single batch dimension"
         camera = self.train_dataset.cameras[image_idx].to(self.device)
         return camera, data

nerfstudio/models/gaussian_splatting.py

@@ -70,10 +70,13 @@ class GaussianSplattingModelConfig(ModelConfig):
     prune_density_until = 15000
     one_up_sh_every = 1000
     lambda_ssim = 0.2
+    use_diff_rast = True # TODO (jake-austin): remove
 
 class GaussianSplattingModel(Model):
     """Gaussian Splatting model
 
+    TODO (jake-austin): Figure out how to print out on the training log in terminal the number of splats
+
     Args:
         config: Gaussian Splatting configuration to instantiate model
     """
@@ -189,7 +192,8 @@ class GaussianSplattingModel(Model):
     def get_features(self):
         features_dc = self._features_dc
         features_rest = self._features_rest
-        return torch.cat((features_dc, features_rest), dim=1)
+
+        return features_dc.squeeze() # torch.cat((features_dc, features_rest), dim=1)
     
     @property
     def get_opacity(self):
@@ -296,44 +300,116 @@ class GaussianSplattingModel(Model):
         Returns:
             Outputs of model. (ie. rendered colors)
         """
-        ray_bundle = ray_bundle.reshape(())
-        assert ray_bundle.shape == ()
-        R = ray_bundle.camera_to_worlds[..., :3, :3].squeeze()
+
+
+        camera = ray_bundle.reshape(())
+        assert camera.shape == ()
+        R = camera.camera_to_worlds[..., :3, :3].squeeze()
         R[:,0] = -R[:,0]
-        T = (R.T @ ray_bundle.camera_to_worlds[..., :3, 3:4]).squeeze()
-        fovx = gaussian_utils.focal2fov(ray_bundle.fx, ray_bundle.width)
-        fovy = gaussian_utils.focal2fov(ray_bundle.fy, ray_bundle.height)
+        T = (R.T @ camera.camera_to_worlds[..., :3, 3:4]).squeeze()
+        fovx = gaussian_utils.focal2fov(camera.fx, camera.width)
+        fovy = gaussian_utils.focal2fov(camera.fy, camera.height)
 
         world_view_transform = gaussian_utils.getWorld2View2(R, T, torch.tensor([0.0, 0.0, 0.0]).to(self.device), 1.0).transpose(0, 1).to(self.device)
         projection_matrix = gaussian_utils.getProjectionMatrix(znear=.01, zfar=100, fovX=fovx, fovY=fovy).transpose(0,1).to(self.device)
         full_projection_matrix = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
 
-        camera_center = world_view_transform.inverse()[3, :3]
-        data = {
-            "xyz": self.get_xyz,
-            "features": self.get_features,
-            "opacity": self.get_opacity,
-            "scaling": self.get_scaling,
-            "rotation": self.get_rotation,
-            "active_sh_degree": self.active_sh_degree,
-            "max_sh_degree": self.max_sh_degree,
-            "FoVx": fovx,
-            "FoVy": fovy,
-            "image_height": ray_bundle.width,
-            "image_width": ray_bundle.width,
-            "world_view_transform": world_view_transform,
-            "full_proj_transform": full_projection_matrix,
-            "camera_center": camera_center,
-        }
+        if not self.config.use_diff_rast:
 
-        image = gaussian_utils.render_from_dict(
-            data,
-            bg_color=torch.ones(3).to(self.device),
-        )["render"].permute(1, 2, 0) #HWC
+            camera_center = world_view_transform.inverse()[3, :3]
+            data = {
+                "xyz": self.get_xyz,
+                "features": self.get_features,
+                "opacity": self.get_opacity,
+                "scaling": self.get_scaling,
+                "rotation": self.get_rotation,
+                "active_sh_degree": self.active_sh_degree,
+                "max_sh_degree": self.max_sh_degree,
+                "FoVx": fovx,
+                "FoVy": fovy,
+                "image_height": camera.width,
+                "image_width": camera.width,
+                "world_view_transform": world_view_transform,
+                "full_proj_transform": full_projection_matrix,
+                "camera_center": camera_center,
+            }
 
-        # TODO (jake-austin): need to clip the images since there are color overflow issues
+            image = gaussian_utils.render_from_dict(
+                data,
+                bg_color=torch.tensor([1,1,1]).to(torch.float32).to(self.device),
+            )["render"].permute(1, 2, 0) #HWC
 
-        return {"rgb": image}
+            return {"rgb": image}
+        
+        else:
+
+            c2w = camera.camera_to_worlds
+            c2w = pose_utils.to4x4(c2w)
+            w2c = torch.linalg.inv(c2w)
+            proj = self._get_proj_matrix(
+                w2c=w2c,
+                fx=camera.fx.item(),
+                fy=camera.fy.item(),
+                width=camera.width.item(),
+                height=camera.height.item(),
+            )
+
+            BLOCK_X = 16
+            BLOCK_Y = 16
+            W = camera.width.item()
+            H = camera.height.item()
+            tile_bounds = (W + BLOCK_X - 1) // BLOCK_X, (H + BLOCK_Y - 1) // BLOCK_Y, 1
+
+            glob_scale = 1
+            xys, depths, radii, conics, num_tiles_hit = diff_rast.project_gaussians(
+                means3d=self.get_xyz,
+                scales=self.get_scaling,
+                glob_scale=glob_scale,
+                quats=self.get_rotation,
+                viewmat=w2c,
+                projmat=proj,
+                img_height=camera.width.item(),
+                img_width=camera.height.item(),
+                fx=camera.fx.item(),
+                fy=camera.fy.item(),
+                tile_bounds=tile_bounds,
+            )
+
+            rgb = diff_rast.rasterize(
+                xys,
+                depths,
+                radii,
+                conics,
+                num_tiles_hit,
+                self.get_features,
+                self.get_opacity.squeeze(),
+                camera.height.item(),
+                camera.width.item(),
+            )
+
+            return {"rgb": rgb}
+
+    def _get_proj_matrix(
+        self, w2c, fx: float, fy: float, width: int, height: int, znear=0.01, zfar=100
+    ):
+        top = 0.5 * height / fy * znear
+        bottom = -top
+        right = 0.5 * width / fx * znear
+        left = -right
+
+        P = torch.zeros(4, 4).to(self.device)
+
+        z_sign = 1.0
+
+        P[0, 0] = 2.0 * znear / (right - left)
+        P[1, 1] = 2.0 * znear / (top - bottom)
+        P[0, 2] = (right + left) / (right - left)
+        P[1, 2] = (top + bottom) / (top - bottom)
+        P[3, 2] = z_sign
+        P[2, 2] = z_sign * zfar / (zfar - znear)
+        P[2, 3] = -(zfar * znear) / (zfar - znear)
+
+        return P @ w2c
 
     def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:
         """Compute and returns metrics.
@@ -353,6 +429,8 @@ class GaussianSplattingModel(Model):
             batch: ground truth batch corresponding to outputs
             metrics_dict: dictionary of metrics, some of which we can use for loss
         """
+        plt.imsave("test.png", np.concatenate([outputs["rgb"].detach().cpu().numpy(), batch["image"].detach().cpu().numpy()], axis=1))
+
         Ll1 = gaussian_utils.l1_loss(batch['image'], outputs['rgb'])
         loss = (1.0 - self.config.lambda_ssim) * Ll1 + self.config.lambda_ssim * (1.0 - gaussian_utils.ssim(batch['image'], outputs['rgb']))
         return {"main_loss": loss}
@@ -364,7 +442,9 @@ class GaussianSplattingModel(Model):
         Args:
             camera_ray_bundle: ray bundle to calculate outputs over
         """
-        return self.get_outputs(camera_ray_bundle)
+        outs = self.get_outputs(camera_ray_bundle.to(self.device))
+        outs["rgb"] = torch.clamp(outs["rgb"], 0.0, 1.0)
+        return outs
 
     def get_image_metrics_and_images(
         self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]