From f2971b504e6506defa17bbcf6ac1e8d282fea6dc Mon Sep 17 00:00:00 2001 From: maturk Date: Sat, 22 Jul 2023 11:30:19 +0300 Subject: [PATCH] camera opt in tensorf --- nerfstudio/models/tensorf.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 742818b64..345050a4c 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -53,6 +53,7 @@ from nerfstudio.model_components.renderers import ( from nerfstudio.model_components.scene_colliders import AABBBoxCollider from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps, colors, misc +from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig @dataclass @@ -89,6 +90,8 @@ class TensoRFModelConfig(ModelConfig): tensorf_encoding: Literal["triplane", "vm", "cp"] = "vm" regularization: Literal["none", "l1", "tv"] = "l1" """Regularization method used in tensorf paper""" + camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3") + """Config of the camera optimizer to use""" class TensoRFModel(Model): @@ -254,6 +257,11 @@ class TensoRFModel(Model): if self.config.tensorf_encoding == "cp" and self.config.regularization == "tv": raise RuntimeError("TV reg not supported for CP decomposition") + # (optional) camera optimizer + self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup( + num_cameras=self.num_train_data, device="cpu" + ) + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} @@ -265,11 +273,14 @@ class TensoRFModel(Model): param_groups["encodings"] = list(self.field.color_encoding.parameters()) + list( self.field.density_encoding.parameters() ) + self.camera_optimizer.get_param_groups(param_groups=param_groups) return param_groups def get_outputs(self, ray_bundle: RayBundle): # uniform sampling + if self.training: + self.camera_optimizer.apply_to_raybundle(ray_bundle) ray_samples_uniform = self.sampler_uniform(ray_bundle) dens = self.field.get_density(ray_samples_uniform) weights = ray_samples_uniform.get_weights(dens) @@ -327,6 +338,8 @@ class TensoRFModel(Model): else: raise ValueError(f"Regularization {self.config.regularization} not supported") + self.camera_optimizer.get_loss_dict(loss_dict) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict @@ -359,5 +372,7 @@ class TensoRFModel(Model): "ssim": float(ssim.item()), "lpips": float(lpips.item()), } + self.camera_optimizer.get_metrics_dict(metrics_dict) + images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth} return metrics_dict, images_dict