formatting

nerfstudio/configs/method_configs.py

@@ -41,7 +41,7 @@ from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParser
 from nerfstudio.data.datasets.depth_dataset import DepthDataset
 from nerfstudio.data.datasets.sdf_dataset import SDFDataset
 from nerfstudio.data.datasets.semantic_dataset import SemanticDataset
-from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig, AdamWOptimizerConfig
+from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
 from nerfstudio.engine.schedulers import (
     CosineDecaySchedulerConfig,
     ExponentialDecaySchedulerConfig,
@@ -629,19 +629,18 @@ method_configs["neuralangelo"] = TrainerConfig(
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
             # "scheduler": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=500000),
             "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
-
         },
         "field_background": {
             # "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15),
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
             "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
-
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
     vis="viewer",
 )
 
+
 def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True):
     """Merge new methods and descriptions into existing methods and descriptions.
     Args:

nerfstudio/data/dataparsers/sdfstudio_dataparser.py

@@ -33,6 +33,7 @@ from nerfstudio.data.dataparsers.base_dataparser import (
     DataparserOutputs,
 )
 from nerfstudio.data.scene_box import SceneBox
+
 # from nerfstudio.utils.images import BasicImages
 from nerfstudio.utils.io import load_from_json
 
@@ -229,7 +230,7 @@ class SDFStudio(DataParser):
                 intrinsics[0, 2] += 200
                 height, width = 1200, 1600
                 meta["height"], meta["width"] = height, width
-            
+
             if self.config.include_mono_prior:
                 assert meta["has_mono_prior"]
                 # load mono depth
@@ -347,7 +348,10 @@ class SDFStudio(DataParser):
         if pairs_path.exists() and split == "train" and self.config.load_pairs:
             with open(pairs_path, "r") as f:
                 pairs = f.readlines()
-            split_ext = lambda x: x.split(".")[0]
+
+            def split_ext(x):
+                return x.split(".")[0]
+
             pairs_srcs = []
             for sources_line in pairs:
                 sources_array = [int(split_ext(img_name)) for img_name in sources_line.split(" ")]

nerfstudio/engine/optimizers.py

@@ -60,6 +60,7 @@ class AdamOptimizerConfig(OptimizerConfig):
     weight_decay: float = 0
     """The weight decay to use."""
 
+
 @dataclass
 class AdamWOptimizerConfig(OptimizerConfig):
     """Basic optimizer config with AdamW"""
@@ -68,6 +69,7 @@ class AdamWOptimizerConfig(OptimizerConfig):
     weight_decay: float = 0
     """The weight decay to use."""
 
+
 @dataclass
 class RAdamOptimizerConfig(OptimizerConfig):
     """Basic optimizer config with RAdam"""

nerfstudio/engine/schedulers.py

@@ -17,7 +17,7 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Literal, Optional, Tuple, Type, Any, List
+from typing import Literal, Optional, Tuple, Type, List
 
 import numpy as np
 from torch.optim import Optimizer, lr_scheduler
@@ -173,6 +173,7 @@ class CosineDecayScheduler(Scheduler):
         scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func)
         return scheduler
 
+
 @dataclass
 class MultiStepWarmupSchedulerConfig(SchedulerConfig):
     """Basic scheduler config with self-defined exponential decay schedule"""
@@ -186,6 +187,7 @@ class MultiStepWarmupSchedulerConfig(SchedulerConfig):
     gamma: float = 0.33
     """The learning rate decay factor."""
 
+
 class MultiStepWarmupScheduler(Scheduler):
     """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler"""
 
@@ -196,8 +198,8 @@ class MultiStepWarmupScheduler(Scheduler):
             if step < self.config.warm_up_end:
                 learning_factor = step / self.config.warm_up_end
             else:
-                index = np.searchsorted(self.config.milestones, step, side='left')
-                learning_factor = self.config.gamma ** index
+                index = np.searchsorted(self.config.milestones, step, side="left")
+                learning_factor = self.config.gamma**index
             return learning_factor
 
         scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func)

nerfstudio/field_components/encodings.py

@@ -169,10 +169,10 @@ class NeRFEncoding(Encoding):
         if self.in_dim is None:
             raise ValueError("Input dimension has not been set")
         out_dim = self.in_dim * self.num_frequencies * 2
-        
+
         if self.off_axis:
             out_dim = self.P.shape[1] * self.num_frequencies * 2
-        
+
         if self.include_input:
             out_dim += self.in_dim
         return out_dim
@@ -776,7 +776,6 @@ class PeriodicVolumeEncoding(Encoding):
         hash_init_scale: float = 0.001,
         smoothstep: bool = False,
     ) -> None:
-
         super().__init__(in_dim=3)
         self.num_levels = num_levels
         self.features_per_level = features_per_level

nerfstudio/fields/sdf_field.py

@@ -69,6 +69,7 @@ except ModuleNotFoundError:
 #         beta = self.beta.abs() + self.beta_min
 #         return beta
 
+
 class LearnedVariance(nn.Module):
     """Variance network in NeuS
 
@@ -197,10 +198,10 @@ class SDFField(Field):
         self.divide_factor = self.config.divide_factor
 
         self.num_levels = self.config.num_levels
-        self.max_res = self.config.max_res 
-        self.base_res = self.config.base_res 
-        self.log2_hashmap_size = self.config.log2_hashmap_size 
-        self.features_per_level = self.config.features_per_level 
+        self.max_res = self.config.max_res
+        self.base_res = self.config.base_res
+        self.log2_hashmap_size = self.config.log2_hashmap_size
+        self.features_per_level = self.config.features_per_level
         self.smoothstep = self.config.smoothstep
         self.growth_factor = np.exp((np.log(config.max_res) - np.log(config.base_res)) / (config.num_levels - 1))
 
@@ -236,7 +237,6 @@ class SDFField(Field):
             print("using tensor vm")
             self.encoding = TensorVMEncoding(128, 24, smoothstep=config.smoothstep)
 
-
         # we concat inputs position ourselves
         self.position_encoding = NeRFEncoding(
             in_dim=3,
@@ -260,7 +260,6 @@ class SDFField(Field):
         # deviation_network to compute alpha from sdf from NeuS
         self.deviation_network = LearnedVariance(init_val=self.config.beta_init)
 
-
         # diffuse and specular tint layer
         if self.config.use_diffuse_color:
             self.diffuse_color_pred = nn.Linear(self.config.geo_feat_dim, 3)
@@ -290,15 +289,15 @@ class SDFField(Field):
         dims = [in_dim] + dims + [3]
         self.num_layers_color = len(dims)
 
-        for l in range(0, self.num_layers_color - 1):
-            out_dim = dims[l + 1]
-            lin = nn.Linear(dims[l], out_dim)
+        for i in range(0, self.num_layers_color - 1):
+            out_dim = dims[i + 1]
+            lin = nn.Linear(dims[i], out_dim)
             torch.nn.init.kaiming_uniform_(lin.weight.data)
             torch.nn.init.zeros_(lin.bias.data)
 
             if self.config.weight_norm:
                 lin = nn.utils.weight_norm(lin)
-            setattr(self, "clin" + str(l), lin)
+            setattr(self, "clin" + str(i), lin)
 
         self.softplus = nn.Softplus(beta=100)
         self.relu = nn.ReLU()
@@ -356,7 +355,7 @@ class SDFField(Field):
 
     def update_mask(self, level: int):
         self.hash_encoding_mask[:] = 1.0
-        self.hash_encoding_mask[level * self.features_per_level:] = 0
+        self.hash_encoding_mask[level * self.features_per_level :] = 0
 
     def forward_geonetwork(self, inputs: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch geo_features+1"]:
         """forward the geonetwork"""
@@ -398,7 +397,7 @@ class SDFField(Field):
         hidden_output = self.forward_geonetwork(positions_flat).view(*ray_samples.frustums.shape, -1)
         sdf, _ = torch.split(hidden_output, [1, self.config.geo_feat_dim], dim=-1)
         return sdf
-    
+
     def set_numerical_gradients_delta(self, delta: float) -> None:
         """Set the delta value for numerical gradient."""
         self.numerical_gradients_delta = delta
@@ -573,7 +572,7 @@ class SDFField(Field):
             inputs = self.spatial_distortion(inputs)
         points_norm = inputs.norm(dim=-1)
         # compute gradient in constracted space
-        
+
         inputs.requires_grad_(True)
         with torch.enable_grad():
             hidden_output = self.forward_geonetwork(inputs)
@@ -584,11 +583,18 @@ class SDFField(Field):
                 skip_spatial_distortion=True,
                 return_sdf=True,
             )
-            sampled_sdf = sampled_sdf.view(-1, *ray_samples.frustums.directions.shape[:-1]).permute(1, 2, 0).contiguous()
+            sampled_sdf = (
+                sampled_sdf.view(-1, *ray_samples.frustums.directions.shape[:-1]).permute(1, 2, 0).contiguous()
+            )
         else:
             d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)
             gradients = torch.autograd.grad(
-                outputs=sdf, inputs=inputs, grad_outputs=d_output, create_graph=True, retain_graph=True, only_inputs=True
+                outputs=sdf,
+                inputs=inputs,
+                grad_outputs=d_output,
+                create_graph=True,
+                retain_graph=True,
+                only_inputs=True,
             )[0]
             sampled_sdf = None
 

nerfstudio/models/neuralangelo.py

@@ -19,24 +19,17 @@ Implementation of Neuralangelo model.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Type
 
-import numpy as np
 import torch
-from torch.nn import Parameter
 
-from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.engine.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
     TrainingCallbackLocation,
 )
 from nerfstudio.field_components.field_heads import FieldHeadNames
-from nerfstudio.fields.density_fields import HashMLPDensityField
-from nerfstudio.model_components.losses import interlevel_loss
-from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
 from nerfstudio.models.neus import NeuSModel, NeuSModelConfig
-from nerfstudio.utils import colormaps
 
 
 @dataclass
@@ -44,7 +37,7 @@ class NeuralangeloModelConfig(NeuSModelConfig):
     """Neuralangelo Model Config"""
 
     _target: Type = field(default_factory=lambda: NeuralangeloModel)
-    #TODO move to base model config since it can be used in all models
+    # TODO move to base model config since it can be used in all models
     enable_progressive_hash_encoding: bool = True
     """whether to use progressive hash encoding"""
     enable_numerical_gradients_schedule: bool = True
@@ -60,6 +53,7 @@ class NeuralangeloModelConfig(NeuSModelConfig):
     steps_per_level: int = 5000
     """steps per level of multi-resolution hash encoding"""
 
+
 class NeuralangeloModel(NeuSModel):
     """Neuralangelo model
 
@@ -82,24 +76,25 @@ class NeuralangeloModel(NeuSModel):
         # read the hash encoding parameters from field
         level_init = self.config.level_init
         # schedule the delta in numerical gradients computation
-        num_levels = self.field.num_levels
         max_res = self.field.max_res
         base_res = self.field.base_res
         growth_factor = self.field.growth_factor
-        
+
         steps_per_level = self.config.steps_per_level
-        
-        
-        init_delta = 1. / base_res
-        end_delta = 1. / max_res
-        
+
+        init_delta = 1.0 / base_res
+        1.0 / max_res
+
         # compute the delta based on level
         if self.config.enable_numerical_gradients_schedule:
+
             def set_delta(step):
-                delta = 1. / (base_res * growth_factor ** ( step / steps_per_level))
-                delta = max(1. / max_res, delta)
-                self.field.set_numerical_gradients_delta(delta * 2.) # TODO because we divide 4 to normalize points to [0, 1]
-                
+                delta = 1.0 / (base_res * growth_factor ** (step / steps_per_level))
+                delta = max(1.0 / max_res, delta)
+                self.field.set_numerical_gradients_delta(
+                    delta * 2.0
+                )  # TODO because we divide 4 to normalize points to [0, 1]
+
             callbacks.append(
                 TrainingCallback(
                     where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
@@ -107,15 +102,16 @@ class NeuralangeloModel(NeuSModel):
                     func=set_delta,
                 )
             )
-        
+
         # schedule the current level of multi-resolution hash encoding
         if self.config.enable_progressive_hash_encoding:
+
             def set_mask(step):
-                #TODO make this consistent with delta schedule
+                # TODO make this consistent with delta schedule
                 level = int(step / steps_per_level) + 1
                 level = max(level, level_init)
                 self.field.update_mask(level)
-    
+
             callbacks.append(
                 TrainingCallback(
                     where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
@@ -126,16 +122,19 @@ class NeuralangeloModel(NeuSModel):
         # schedule the curvature loss weight
         # linear warmup for 5000 steps to 5e-4 and then decay as delta
         if self.config.enable_curvature_loss_schedule:
+
             def set_curvature_loss_mult_factor(step):
                 if step < self.config.curvature_loss_warmup_steps:
                     factor = step / self.config.curvature_loss_warmup_steps
                 else:
-                    delta = 1. / (base_res * growth_factor ** ( (step - self.config.curvature_loss_warmup_steps) / steps_per_level))
-                    delta = max(1. / max_res, delta)
+                    delta = 1.0 / (
+                        base_res * growth_factor ** ((step - self.config.curvature_loss_warmup_steps) / steps_per_level)
+                    )
+                    delta = max(1.0 / max_res, delta)
                     factor = delta / init_delta
 
                 self.curvature_loss_multi_factor = factor
-            
+
             callbacks.append(
                 TrainingCallback(
                     where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
@@ -143,38 +142,38 @@ class NeuralangeloModel(NeuSModel):
                     func=set_curvature_loss_mult_factor,
                 )
             )
-        
-        
-        #TODO switch to analytic gradients after delta is small enough?
-        
+
+        # TODO switch to analytic gradients after delta is small enough?
+
         return callbacks
-    
+
     def get_metrics_dict(self, outputs, batch) -> Dict:
         metrics_dict = super().get_metrics_dict(outputs, batch)
-        
+
         if self.training:
             # training statics
             metrics_dict["activated_encoding"] = self.field.hash_encoding_mask.mean().item()
             metrics_dict["numerical_gradients_delta"] = self.field.numerical_gradients_delta
             metrics_dict["curvature_loss_multi"] = self.curvature_loss_multi_factor * self.config.curvature_loss_multi
-            
+
         return metrics_dict
-    
+
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict:
         loss_dict = super().get_loss_dict(outputs, batch, metrics_dict)
 
         # curvature loss
         if self.training and self.config.curvature_loss_multi > 0.0:
             delta = self.field.numerical_gradients_delta
-            centered_sdf = outputs['field_outputs'][FieldHeadNames.SDF]
-            sourounding_sdf = outputs['field_outputs']["sampled_sdf"]
-            
+            centered_sdf = outputs["field_outputs"][FieldHeadNames.SDF]
+            sourounding_sdf = outputs["field_outputs"]["sampled_sdf"]
+
             sourounding_sdf = sourounding_sdf.reshape(centered_sdf.shape[:2] + (3, 2))
-            
+
             # (a - b)/d - (b -c)/d = (a + c - 2b)/d
             # ((a - b)/d - (b -c)/d)/d = (a + c - 2b)/(d*d)
             curvature = (sourounding_sdf.sum(dim=-1) - 2 * centered_sdf) / (delta * delta)
-            loss_dict["curvature_loss"] = torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor
-            
-        return loss_dict
+            loss_dict["curvature_loss"] = (
+                torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor
+            )
 
+        return loss_dict

nerfstudio/models/neus_facto.py

@@ -219,7 +219,6 @@ class NeuSFactoModel(NeuSModel):
         # read the hash encoding parameters from field
         level_init = self.config.level_init
         # schedule the delta in numerical gradients computation
-        num_levels = self.field.num_levels
         max_res = self.field.max_res
         base_res = self.field.base_res
         growth_factor = self.field.growth_factor
@@ -227,7 +226,7 @@ class NeuSFactoModel(NeuSModel):
         steps_per_level = self.config.steps_per_level
 
         init_delta = 1.0 / base_res
-        end_delta = 1.0 / max_res
+        1.0 / max_res
 
         if self.config.enable_numerical_gradients_schedule:
 
@@ -245,8 +244,8 @@ class NeuSFactoModel(NeuSModel):
                     func=set_delta,
                 )
             )
-        
-         # schedule the current level of multi-resolution hash encoding
+
+        # schedule the current level of multi-resolution hash encoding
         if self.config.enable_progressive_hash_encoding:
 
             def set_mask(step):
@@ -339,7 +338,7 @@ class NeuSFactoModel(NeuSModel):
                 torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor
             )
         return loss_dict
-    
+
     def get_metrics_dict(self, outputs, batch) -> Dict:
         metrics_dict = super().get_metrics_dict(outputs, batch)
 
@@ -351,7 +350,6 @@ class NeuSFactoModel(NeuSModel):
 
         return metrics_dict
 
-
     def get_image_metrics_and_images(
         self, outputs: Dict[str, Any], batch: Dict[str, Any]
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: