correct vis mask for colmap data

configs/neuralangelo-colmap.yaml

@@ -193,7 +193,7 @@ trainer:
  max_steps: 1e5
  log_every_n_steps: 500
  num_sanity_val_steps: 0
- val_check_interval: 1e5
+ val_check_interval: 5e4
  limit_train_batches: 1.0
  limit_val_batches: 4
  enable_progress_bar: true

configs/neus-colmap.yaml

@@ -191,7 +191,7 @@ trainer:
  max_steps: 1e5
  log_every_n_steps: 500
  num_sanity_val_steps: 0
- val_check_interval: 1e5
+ val_check_interval: 5e4
  limit_train_batches: 1.0
  limit_val_batches: 4
  enable_progress_bar: true

datasets/colmap.py

@@ -246,10 +246,10 @@ def setup(self, config, split):
  mask = TF.to_tensor(mask)[0]
  else:
  mask = torch.ones_like(img[...,0], device=img.device)
-  vis_mask = torch.ones_like(img[...,0], device=img.device)
+ vis_mask = torch.ones_like(img[...,0], device=img.device)
  all_fg_masks.append(mask) # (h, w)
+ all_vis_masks.append(vis_mask) # (h, w)
  all_images.append(img)
- all_vis_masks.append(vis_mask)
  if self.apply_depth:
  # load estimated or recorded depth map
  depth_path = os.path.join(self.config.root_dir, f"{frame['depth_path']}")
@@ -271,7 +271,7 @@ def setup(self, config, split):
  all_depths.append(torch.zeros_like(img[...,0], device=img.device))
  all_depth_masks.append(torch.zeros_like(img[...,0], device=img.device))
 
- all_c2w, all_images, all_fg_masks, all_depths, all_depth_masks, self.all_vis_masks = \
+ all_c2w, all_images, all_fg_masks, all_depths, all_depth_masks, all_vis_masks = \
  torch.stack(all_c2w, dim=0).float(), \
  torch.stack(all_images, dim=0).float(), \
  torch.stack(all_fg_masks, dim=0).float(), \
@@ -298,6 +298,7 @@ def setup(self, config, split):
  'all_fg_masks': all_fg_masks,
  'all_depths': all_depths,
  'all_depth_masks': all_depth_masks,
+ 'all_vis_masks': all_vis_masks,
  }
 
  ColmapDatasetBase.initialized = True
@@ -309,6 +310,7 @@ def setup(self, config, split):
  self.all_c2w = create_spheric_poses(self.all_c2w[:,:,3], n_steps=self.config.n_test_traj_steps)
  self.all_images = torch.zeros((self.config.n_test_traj_steps, self.h, self.w, 3), dtype=torch.float32)
  self.all_fg_masks = torch.zeros((self.config.n_test_traj_steps, self.h, self.w), dtype=torch.float32)
+ self.all_vis_masks = torch.ones((self.config.n_test_traj_steps, self.h, self.w), dtype=torch.float32)
 
  """
  # for debug use
@@ -342,6 +344,7 @@ def setup(self, config, split):
  if self.config.load_data_on_gpu:
  self.all_images = self.all_images.to(self.rank) 
  self.all_fg_masks = self.all_fg_masks.to(self.rank)
+ self.all_vis_masks = self.all_vis_masks.to(self.rank)
 
 
 class ColmapDataset(Dataset, ColmapDatasetBase):

requirements.txt

@@ -1,4 +1,6 @@
 pytorch-lightning<2
+torchvision
+torchtyping
 omegaconf==2.2.3
 nerfacc==0.3.3
 matplotlib