Merge branch 'dev/231106' into main

configs/neuralangelo-colmap_dense.yaml

@@ -5,7 +5,7 @@ seed: 42
 dataset:
   name: colmap
   root_dir: ???
-  img_downscale: 1 # specify training image size by either img_wh or img_downscale
+  img_downscale: 2 # specify training image size by either img_wh or img_downscale
   up_est_method: ground # if true, use estimated ground plane normal direction as up direction
   center_est_method: lookat
   n_test_traj_steps: 30
@@ -15,7 +15,7 @@ dataset:
 
 model:
   name: neus
-  radius: 2.5
+  radius: 1.5
   num_samples_per_ray: 1024
   train_num_rays: 128
   max_train_num_rays: 8192

configs/neuralangelo-colmap_sparse-wmask.yaml

@@ -134,8 +134,8 @@ trainer:
   max_steps: 20000
   log_every_n_steps: 100
   num_sanity_val_steps: 0
-  val_check_interval: 5000
+  val_check_interval: 10000
   limit_train_batches: 1.0
-  limit_val_batches: 2
+  limit_val_batches: 1
   enable_progress_bar: true
   precision: 16

configs/neuralangelo-colmap_sparse-wreflection.yaml

@@ -0,0 +1,183 @@
+name: neuralangelo-colmap_sparse-wreflection-${basename:${dataset.root_dir}}
+tag: ""
+seed: 42
+
+dataset:
+  name: colmap
+  root_dir: ???
+  img_downscale: 4 # specify training image size by either img_wh or img_downscale
+  up_est_method: ground # if true, use estimated ground plane normal direction as up direction
+  center_est_method: lookat
+  n_test_traj_steps: 30
+  apply_mask: false
+  load_data_on_gpu: false
+  dense_pcd_path: null
+
+model:
+  name: neus
+  radius: 1.5
+  num_samples_per_ray: 1024
+  train_num_rays: 128
+  max_train_num_rays: 8192
+  grid_prune: true
+  grid_prune_occ_thre: 0.001
+  dynamic_ray_sampling: true
+  batch_image_sampling: true
+  randomized: true
+  ray_chunk: 2048
+  cos_anneal_end: 20000
+  learned_background: true
+  background_color: random
+  variance:
+    init_val: 0.3
+    modulate: false
+  geometry:
+    name: volume-sdf
+    radius: ${model.radius}
+    feature_dim: 65
+    grad_type: analytic
+    finite_difference_eps: progressive
+    isosurface:
+      method: mc
+      resolution: 512
+      chunk: 2097152
+      threshold: 0.001
+    xyz_encoding_config:
+      otype: ProgressiveBandHashGrid
+      n_levels: 16
+      n_features_per_level: 2
+      log2_hashmap_size: 19
+      base_resolution: 32
+      per_level_scale: 1.3195079107728942
+      include_xyz: true
+      start_level: 4
+      start_step: 5000
+      update_steps: 1000
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 64
+      n_hidden_layers: 2
+      sphere_init: true
+      sphere_init_radius: 0.5
+      weight_norm: true
+  texture:
+    name: volume-dual-colorV2
+    input_feature_dim: ${add:${model.geometry.feature_dim},3} # surface normal as additional input
+    diffuse_warmup_steps: 5000 
+    dir_encoding_config:
+      otype: SphericalHarmonics
+      degree: 4
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 256
+      n_hidden_layers: 4
+    color_activation: sigmoid
+  # background model configurations
+  num_samples_per_ray_bg: 256
+  geometry_bg:
+    name: volume-density
+    radius: ${model.radius}
+    feature_dim: 8
+    density_activation: trunc_exp
+    density_bias: -1
+    isosurface: null
+    xyz_encoding_config:
+      otype: ProgressiveBandHashGrid
+      n_levels: 16
+      n_features_per_level: 2
+      log2_hashmap_size: 19
+      base_resolution: 32
+      per_level_scale: 1.3195079107728942
+      include_xyz: true
+      start_level: 4
+      start_step: 5000
+      update_steps: 1000
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 64
+      n_hidden_layers: 1
+  texture_bg:
+    name: volume-radiance
+    input_feature_dim: ${model.geometry_bg.feature_dim}
+    dir_encoding_config:
+      otype: SphericalHarmonics
+      degree: 4
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 64
+      n_hidden_layers: 2
+    color_activation: sigmoid
+
+system:
+  name: neus-system
+  loss:
+    lambda_sdf_l1: 0
+    lambda_normal: 0.
+    lambda_rgb_mse: 5.
+    lambda_rgb_l1: 0.
+    lambda_mask: 0.0
+    lambda_eikonal: 0.1
+    lambda_curvature: [0, 0, 5.e-2, 5000]
+    lambda_sparsity: 0.0
+    lambda_distortion: 0.0
+    lambda_distortion_bg: 0.0
+    lambda_opaque: 0.0
+    sparsity_scale: 1.
+  optimizer:
+    name: AdamW
+    args:
+      lr: 0.01
+      betas: [0.9, 0.99]
+      eps: 1.e-15
+    params:
+      geometry:
+        lr: 0.01
+      texture:
+        lr: 0.01
+      geometry_bg:
+        lr: 0.01
+      texture_bg:
+        lr: 0.01
+      variance:
+        lr: 0.001
+  warmup_steps: 500
+  scheduler:
+    name: SequentialLR
+    interval: step
+    milestones:
+      - ${system.warmup_steps}
+    schedulers:
+      - name: LinearLR # linear warm-up in the first system.warmup_steps steps
+        args:
+          start_factor: 0.01
+          end_factor: 1.0
+          total_iters: ${system.warmup_steps}
+      - name: ExponentialLR
+        args:
+          gamma: ${calc_exp_lr_decay_rate:0.1,${sub:${trainer.max_steps},${system.warmup_steps}}}
+
+checkpoint:
+  save_top_k: -1
+  every_n_train_steps: ${trainer.max_steps}
+
+export:
+  chunk_size: 2097152
+  export_vertex_color: True
+
+trainer:
+  max_steps: 20000
+  log_every_n_steps: 100
+  num_sanity_val_steps: 0
+  val_check_interval: 5000
+  limit_train_batches: 1.0
+  limit_val_batches: 1
+  enable_progress_bar: true
+  precision: 16

configs/neuralangelo-colmap_sparse.yaml

@@ -8,7 +8,7 @@ dataset:
   img_downscale: 2 # specify training image size by either img_wh or img_downscale
   up_est_method: ground # if true, use estimated ground plane normal direction as up direction
   center_est_method: lookat
-  n_test_traj_steps: 30
+  n_test_traj_steps: 6
   apply_mask: false
   load_data_on_gpu: false
   dense_pcd_path: null
@@ -176,8 +176,8 @@ trainer:
   max_steps: 20000
   log_every_n_steps: 100
   num_sanity_val_steps: 0
-  val_check_interval: 5000
+  val_check_interval: 10000
   limit_train_batches: 1.0
-  limit_val_batches: 2
+  limit_val_batches: 1
   enable_progress_bar: true
   precision: 16

datasets/colmap.py

@@ -314,7 +314,6 @@ def setup(self, config, split):
         self.all_points = self.all_points.float()
         self.pts3d_normal = self.pts3d_normal.float()
         self.all_points_ = contract_to_unisphere(self.all_points, 1.0, ContractionType.AABB) # points normalized to (0, 1)
-        self.all_fg_indexs, self.all_bg_indexs = self.all_fg_indexs.to(self.rank), self.all_bg_indexs.to(self.rank)
 
     def query_radius_occ(self, query_points, radius=0.01):
 

export.py

@@ -5,56 +5,47 @@
 import logging
 from datetime import datetime
 import trimesh
+import numpy as np
 
 logging.basicConfig(level=logging.INFO)
 
-def decimate_mesh(mesh: str, decimation_factor: float):
-    logging.info(f"Original mesh with {len(mesh.faces)} faces.")
-
-    # Decimate the mesh
-    if decimation_factor < 1:
-        decimation_factor = int(len(mesh.faces) * decimation_factor)
-    else:
-        decimation_factor = int(decimation_factor)
-
-    mesh = mesh.simplify_quadratic_decimation(decimation_factor)
-    logging.info(f"Decimated mesh to {len(mesh.faces)} faces.")
-
-    return mesh
 
 def main():
     logging.info("Start exporting.")
     parser = argparse.ArgumentParser()
     parser.add_argument('--gpu', default='0', help='GPU(s) to be used')
     parser.add_argument('--exp_dir', required=True)
-    parser.add_argument('--output-dir', required=True)
-
-    parser.add_argument('--decimate', type=float, help='Specifies the desired final size of the mesh. \
-                        If the number is less than 1, it represents the final size as a percentage of the initial size. \
-                        If the number is greater than 1, it represents the desired number of faces.')
+    parser.add_argument('--flip', action='store_true')
+    parser.add_argument('--res', default=1024)
+    parser.add_argument('--output-dir', default='results')
     args, extras = parser.parse_known_args()
 
     # set CUDA_VISIBLE_DEVICES then import pytorch-lightning
     os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
     os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
     n_gpus = len(args.gpu.split(','))
 
-    code_dir = os.path.join(args.exp_dir, 'code')
+    # code_dir = os.path.join(args.exp_dir, 'code')
     ckpt_dir = os.path.join(args.exp_dir, 'ckpt')
     latest_ckpt = sorted(os.listdir(ckpt_dir), key=lambda s: int(s.split('-')[0].split('=')[1]), reverse=True)[0]
     latest_ckpt = os.path.join(ckpt_dir, latest_ckpt)
     config_path = os.path.join(args.exp_dir, 'config', 'parsed.yaml')
 
-    logging.info(f"Importing modules from cached code: {code_dir}")
-    sys.path.append(code_dir)
-    import datasets
+    # logging.info(f"Importing modules from cached code: {code_dir}")
+    # sys.path.append(code_dir)
     import systems
     import pytorch_lightning as pl
     from utils.misc import load_config    
 
     # parse YAML config to OmegaConf
     logging.info(f"Loading configuration: {config_path}")
     config = load_config(config_path, cli_args=extras)
+
+    # Update level of ProgressiveBandHashGrid
+    if  config.model.geometry.xyz_encoding_config.otype == 'ProgressiveBandHashGrid':
+        config.model.geometry.xyz_encoding_config.start_level = config.model.geometry.xyz_encoding_config.n_levels
+    config.model.geometry.isosurface.resolution = args.res
+    config.export.export_vertex_color = True
     config.cmd_args = vars(args)
 
     if 'seed' not in config:
@@ -63,19 +54,29 @@ def main():
     logging.info(f"Creating system: {config.system.name}")
     system = systems.make(config.system.name, config, load_from_checkpoint=latest_ckpt)
     system.model.cuda()
-    mesh = system.model.isosurface()
+    mesh = system.model.export(config.export)
+
+    mesh['v_pos'] = mesh['v_pos'][:, [0, 2, 1]].numpy()
+    if args.flip:
+        mesh['t_pos_idx'] = mesh['t_pos_idx'].numpy()[:, [0, 2, 1]]
+    else:
+        mesh['t_pos_idx'] = np.fliplr(mesh['t_pos_idx'].numpy())[:, [0, 2, 1]]
+
     mesh = trimesh.Trimesh(
-        vertices=mesh['v_pos'].numpy(),
-        faces=mesh['t_pos_idx'].numpy()
+            vertices=mesh['v_pos'],
+            faces=mesh['t_pos_idx'],
+            vertex_colors=mesh['v_rgb'].numpy(),
+            vertex_normals=mesh['v_norm'].numpy()
+        )
+    mesh.visual.material = trimesh.visual.material.PBRMaterial(
+        metallicFactor=0.25,
+        roughnessFactor=0.25
     )
 
-    if args.decimate > 0:
-        logging.info("Decimating mesh.")
-        mesh = decimate_mesh(mesh, args.decimate)
-
     os.makedirs(args.output_dir, exist_ok=True)
     logging.info("Exporting mesh.")
-    mesh.export(os.path.join(args.output_dir, 'iso_mesh.ply'))
+    mesh.export(os.path.join(args.output_dir, f'{config.name}.glb'))
+    mesh.export(os.path.join(args.output_dir, f'{config.name}.obj'))
     logging.info("Export finished successfully.")
 
 if __name__ == '__main__':