From 1a214ba129bfad8eedcff4cafa87443e5d54081d Mon Sep 17 00:00:00 2001 From: "Jake Austin jake-austin@berkeley.edu" Date: Tue, 12 Sep 2023 23:19:47 +0000 Subject: [PATCH] Temp save --- .../datamanagers/full_images_datamanager.py | 41 ++++- nerfstudio/models/gaussian_splatting.py | 142 ++++++++++++++---- 2 files changed, 150 insertions(+), 33 deletions(-) diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 27a49f8c0..9576b8604 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/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 diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index c022aaf3d..a41e97933 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/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]