add model, dataparser, sdf_field

nerfstudio/data/dataparsers/sdfstudio_dataparser.py

@@ -1,4 +1,4 @@
-# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,15 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Datapaser for sdfstudio formatted data"""
 
+"""Data parser for friends dataset"""
 from __future__ import annotations
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Type
+from typing import Dict, Optional, Type
 
+import numpy as np
 import torch
+from PIL import Image
+from rich.console import Console
+from torchtyping import TensorType
 
 from nerfstudio.cameras import camera_utils
 from nerfstudio.cameras.cameras import Cameras, CameraType
@@ -29,8 +33,105 @@ 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
 
+CONSOLE = Console()
+
+
+def get_src_from_pairs(
+    ref_idx, all_imgs, pairs_srcs, neighbors_num=None, neighbors_shuffle=False
+) -> Dict[str, TensorType]:
+    # src_idx[0] is ref img
+    src_idx = pairs_srcs[ref_idx]
+    # randomly sample neighbors
+    if neighbors_num and neighbors_num > -1 and neighbors_num < len(src_idx) - 1:
+        if neighbors_shuffle:
+            perm_idx = torch.randperm(len(src_idx) - 1) + 1
+            src_idx = torch.cat([src_idx[[0]], src_idx[perm_idx[:neighbors_num]]])
+        else:
+            src_idx = src_idx[: neighbors_num + 1]
+    src_idx = src_idx.to(all_imgs.device)
+    return {"src_imgs": all_imgs[src_idx], "src_idxs": src_idx}
+
+
+def get_image(image_filename, alpha_color=None) -> TensorType["image_height", "image_width", "num_channels"]:
+    """Returns a 3 channel image.
+
+    Args:
+        image_idx: The image index in the dataset.
+    """
+    pil_image = Image.open(image_filename)
+    np_image = np.array(pil_image, dtype="uint8")  # shape is (h, w, 3 or 4)
+    assert len(np_image.shape) == 3
+    assert np_image.dtype == np.uint8
+    assert np_image.shape[2] in [3, 4], f"Image shape of {np_image.shape} is in correct."
+    image = torch.from_numpy(np_image.astype("float32") / 255.0)
+    if alpha_color is not None and image.shape[-1] == 4:
+        assert image.shape[-1] == 4
+        image = image[:, :, :3] * image[:, :, -1:] + alpha_color * (1.0 - image[:, :, -1:])
+    else:
+        image = image[:, :, :3]
+    return image
+
+
+def get_depths_and_normals(image_idx: int, depths, normals):
+    """function to process additional depths and normal information
+
+    Args:
+        image_idx: specific image index to work with
+        semantics: semantics data
+    """
+
+    # depth
+    depth = depths[image_idx]
+    # normal
+    normal = normals[image_idx]
+
+    return {"depth": depth, "normal": normal}
+
+
+def get_sensor_depths(image_idx: int, sensor_depths):
+    """function to process additional sensor depths
+
+    Args:
+        image_idx: specific image index to work with
+        sensor_depths: semantics data
+    """
+
+    # sensor depth
+    sensor_depth = sensor_depths[image_idx]
+
+    return {"sensor_depth": sensor_depth}
+
+
+def get_foreground_masks(image_idx: int, fg_masks):
+    """function to process additional foreground_masks
+
+    Args:
+        image_idx: specific image index to work with
+        fg_masks: foreground_masks
+    """
+
+    # sensor depth
+    fg_mask = fg_masks[image_idx]
+
+    return {"fg_mask": fg_mask}
+
+
+def get_sparse_sfm_points(image_idx: int, sfm_points):
+    """function to process additional sparse sfm points
+
+    Args:
+        image_idx: specific image index to work with
+        sfm_points: sparse sfm points
+    """
+
+    # sfm points
+    sparse_sfm_points = sfm_points[image_idx]
+    sparse_sfm_points = BasicImages([sparse_sfm_points])
+    return {"sparse_sfm_points": sparse_sfm_points}
+
 
 @dataclass
 class SDFStudioDataParserConfig(DataParserConfig):
@@ -42,19 +143,35 @@ class SDFStudioDataParserConfig(DataParserConfig):
     """Directory specifying location of data."""
     include_mono_prior: bool = False
     """whether or not to load monocular depth and normal """
-    depth_unit_scale_factor: float = 1e-3
-    """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
+    include_sensor_depth: bool = False
+    """whether or not to load sensor depth"""
     include_foreground_mask: bool = False
     """whether or not to load foreground mask"""
+    include_sfm_points: bool = False
+    """whether or not to load sfm points"""
+    depth_unit_scale_factor: float = 1e-3
+    """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
     downscale_factor: int = 1
     scene_scale: float = 2.0
     """
     Sets the bounding cube to have edge length of this size.
-    The longest dimension of the axis-aligned bbox will be scaled to this value.
+    The longest dimension of the Friends axis-aligned bbox will be scaled to this value.
     """
+    load_pairs: bool = False
+    """whether to load pairs for multi-view consistency"""
+    neighbors_num: Optional[int] = None
+    neighbors_shuffle: Optional[bool] = False
+    pairs_sorted_ascending: Optional[bool] = True
+    """if src image pairs are sorted in ascending order by similarity i.e. 
+    the last element is the most similar to the first (ref)"""
     skip_every_for_val_split: int = 1
     """sub sampling validation images"""
+    train_val_no_overlap: bool = False
+    """remove selected / sampled validation images from training set"""
     auto_orient: bool = True
+    """automatically orient the scene such that the up direction is the same as the viewer's up direction"""
+    load_dtu_highres: bool = False
+    """load high resolution images from DTU dataset, should only be used for the preprocessed DTU dataset"""
 
 
 @dataclass
@@ -63,7 +180,7 @@ class SDFStudio(DataParser):
 
     config: SDFStudioDataParserConfig
 
-    def _generate_dataparser_outputs(self, split="train"):
+    def _generate_dataparser_outputs(self, split="train"):  # pylint: disable=unused-argument,too-many-statements
         # load meta data
         meta = load_from_json(self.config.data / "meta_data.json")
 
@@ -71,10 +188,17 @@ class SDFStudio(DataParser):
         # subsample to avoid out-of-memory for validation set
         if split != "train" and self.config.skip_every_for_val_split >= 1:
             indices = indices[:: self.config.skip_every_for_val_split]
+        else:
+            # if you use this option, training set should not contain any image in validation set
+            if self.config.train_val_no_overlap:
+                indices = [i for i in indices if i % self.config.skip_every_for_val_split != 0]
 
         image_filenames = []
-        depth_filenames = []
-        normal_filenames = []
+        depth_images = []
+        normal_images = []
+        sensor_depth_images = []
+        foreground_mask_images = []
+        sfm_points = []
         transform = None
         fx = []
         fy = []
@@ -86,28 +210,71 @@ class SDFStudio(DataParser):
                 continue
 
             image_filename = self.config.data / frame["rgb_path"]
-            depth_filename = frame.get("mono_depth_path")
-            normal_filename = frame.get("mono_normal_path")
 
             intrinsics = torch.tensor(frame["intrinsics"])
             camtoworld = torch.tensor(frame["camtoworld"])
 
             # append data
             image_filenames.append(image_filename)
-            if depth_filename is not None and normal_filename is not None:
-                depth_filenames.append(self.config.data / depth_filename)
-                normal_filenames.append(self.config.data / normal_filename)
             fx.append(intrinsics[0, 0])
             fy.append(intrinsics[1, 1])
             cx.append(intrinsics[0, 2])
             cy.append(intrinsics[1, 2])
             camera_to_worlds.append(camtoworld)
 
+            # here is hard coded for DTU high-res images
+            if self.config.load_dtu_highres:
+                image_filename = self.config.data / "image" / frame["rgb_path"].replace("_rgb", "")
+                intrinsics[:2, :] *= 1200 / 384.0
+                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
+                depth = np.load(self.config.data / frame["mono_depth_path"])
+                depth_images.append(torch.from_numpy(depth).float())
+
+                # load mono normal
+                normal = np.load(self.config.data / frame["mono_normal_path"])
+
+                # transform normal to world coordinate system
+                normal = normal * 2.0 - 1.0  # omnidata output is normalized so we convert it back to normal here
+                normal = torch.from_numpy(normal).float()
+
+                rot = camtoworld[:3, :3]
+
+                normal_map = normal.reshape(3, -1)
+                normal_map = torch.nn.functional.normalize(normal_map, p=2, dim=0)
+
+                normal_map = rot @ normal_map
+                normal_map = normal_map.permute(1, 0).reshape(*normal.shape[1:], 3)
+                normal_images.append(normal_map)
+
+            if self.config.include_sensor_depth:
+                assert meta["has_sensor_depth"]
+                # load sensor depth
+                sensor_depth = np.load(self.config.data / frame["sensor_depth_path"])
+                sensor_depth_images.append(torch.from_numpy(sensor_depth).float())
+
+            if self.config.include_foreground_mask:
+                assert meta["has_foreground_mask"]
+                # load foreground mask
+                foreground_mask = np.array(Image.open(self.config.data / frame["foreground_mask"]), dtype="uint8")
+                foreground_mask = foreground_mask[..., :1]
+                foreground_mask_images.append(torch.from_numpy(foreground_mask).float() / 255.0)
+
+            if self.config.include_sfm_points:
+                assert meta["has_sparse_sfm_points"]
+                # load sparse sfm points
+                sfm_points_view = np.loadtxt(self.config.data / frame["sfm_sparse_points_view"])
+                sfm_points.append(torch.from_numpy(sfm_points_view).float())
+
         fx = torch.stack(fx)
         fy = torch.stack(fy)
         cx = torch.stack(cx)
         cy = torch.stack(cy)
-        c2w_colmap = torch.stack(camera_to_worlds)
         camera_to_worlds = torch.stack(camera_to_worlds)
 
         # Convert from COLMAP's/OPENCV's camera coordinate system to nerfstudio
@@ -120,6 +287,15 @@ class SDFStudio(DataParser):
                 center_method="none",
             )
 
+            # we should also transform normal accordingly
+            normal_images_aligned = []
+            for normal_image in normal_images:
+                h, w, _ = normal_image.shape
+                normal_image = transform[:3, :3] @ normal_image.reshape(-1, 3).permute(1, 0)
+                normal_image = normal_image.permute(1, 0).reshape(h, w, 3)
+                normal_images_aligned.append(normal_image)
+            normal_images = normal_images_aligned
+
         # scene box from meta data
         meta_scene_box = meta["scene_box"]
         aabb = torch.tensor(meta_scene_box["aabb"], dtype=torch.float32)
@@ -141,19 +317,69 @@ class SDFStudio(DataParser):
 
         # TODO supports downsample
         # cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor)
+
         if self.config.include_mono_prior:
-            assert meta["has_mono_prior"], f"no mono prior in {self.config.data}"
+            additional_inputs_dict = {
+                "cues": {"func": get_depths_and_normals, "kwargs": {"depths": depth_images, "normals": normal_images}}
+            }
+        else:
+            additional_inputs_dict = {}
+
+        if self.config.include_sensor_depth:
+            additional_inputs_dict["sensor_depth"] = {
+                "func": get_sensor_depths,
+                "kwargs": {"sensor_depths": sensor_depth_images},
+            }
+
+        if self.config.include_foreground_mask:
+            additional_inputs_dict["foreground_masks"] = {
+                "func": get_foreground_masks,
+                "kwargs": {"fg_masks": foreground_mask_images},
+            }
+
+        # if self.config.include_sfm_points:
+        #     additional_inputs_dict["sfm_points"] = {
+        #         "func": get_sparse_sfm_points,
+        #         "kwargs": {"sfm_points": sfm_points},
+        #     }
+        # load pair information
+        pairs_path = self.config.data / "pairs.txt"
+        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]
+            pairs_srcs = []
+            for sources_line in pairs:
+                sources_array = [int(split_ext(img_name)) for img_name in sources_line.split(" ")]
+                if self.config.pairs_sorted_ascending:
+                    # invert (flip) the source elements s.t. the most similar source is in index 1 (index 0 is reference)
+                    sources_array = [sources_array[0]] + sources_array[:1:-1]
+                pairs_srcs.append(sources_array)
+            pairs_srcs = torch.tensor(pairs_srcs)
+            all_imgs = torch.stack(
+                [get_image(image_filename) for image_filename in sorted(image_filenames)], axis=0
+            ).cuda()
+
+            additional_inputs_dict["pairs"] = {
+                "func": get_src_from_pairs,
+                "kwargs": {
+                    "all_imgs": all_imgs,
+                    "pairs_srcs": pairs_srcs,
+                    "neighbors_num": self.config.neighbors_num,
+                    "neighbors_shuffle": self.config.neighbors_shuffle,
+                },
+            }
 
         dataparser_outputs = DataparserOutputs(
             image_filenames=image_filenames,
             cameras=cameras,
             scene_box=scene_box,
             metadata={
-                "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
-                "normal_filenames": normal_filenames if len(normal_filenames) > 0 else None,
+                "depth_filenames": depth_images if len(depth_images) > 0 else None,
+                "normal_filenames": normal_images if len(normal_images) > 0 else None,
                 "transform": transform,
                 # required for normal maps, these are in colmap format so they require c2w before conversion
-                "camera_to_worlds": c2w_colmap if len(c2w_colmap) > 0 else None,
+                "camera_to_worlds": camera_to_worlds if len(camera_to_worlds) > 0 else None,
                 "include_mono_prior": self.config.include_mono_prior,
                 "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
             },

nerfstudio/fields/sdf_field.py

@@ -148,7 +148,14 @@ class SDFField(Field):
         self.use_grid_feature = self.config.use_grid_feature
         self.divide_factor = self.config.divide_factor
 
-        growth_factor = np.exp((np.log(config.max_res) - np.log(config.base_res)) / (config.num_levels - 1))
+        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.use_hash = self.config.use_hash
+        self.smoothstep = self.config.smoothstep
+        self.growth_factor = np.exp((np.log(config.max_res) - np.log(config.base_res)) / (config.num_levels - 1))
 
         if self.config.encoding_type == "hash":
             # feature encoding
@@ -160,10 +167,14 @@ class SDFField(Field):
                     "n_features_per_level": config.features_per_level,
                     "log2_hashmap_size": config.log2_hashmap_size,
                     "base_resolution": config.base_res,
-                    "per_level_scale": growth_factor,
+                    "per_level_scale": self.growth_factor,
                     "interpolation": "Smoothstep" if config.smoothstep else "Linear",
                 },
             )
+            self.hash_encoding_mask = torch.ones(
+                self.num_levels * self.features_per_level,
+                dtype=torch.float32,
+            )
 
         # we concat inputs position ourselves
         self.position_encoding = NeRFEncoding(
@@ -206,6 +217,7 @@ class SDFField(Field):
         self.sigmoid = torch.nn.Sigmoid()
 
         self._cos_anneal_ratio = 1.0
+        self.numerical_gradients_delta = 0.0001
 
         if self.use_grid_feature:
             assert self.spatial_distortion is not None, "spatial distortion must be provided when using grid feature"
@@ -257,6 +269,10 @@ class SDFField(Field):
         """Set the anneal value for the proposal network."""
         self._cos_anneal_ratio = anneal
 
+    def update_mask(self, level: int):
+        self.hash_encoding_mask[:] = 1.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"""
         if self.use_grid_feature:
@@ -265,6 +281,8 @@ class SDFField(Field):
             # map range [-2, 2] to [0, 1]
             positions = (positions + 2.0) / 4.0
             feature = self.encoding(positions)
+            # mask feature
+            feature = feature * self.hash_encoding_mask.to(feature.device)
         else:
             feature = torch.zeros_like(inputs[:, :1].repeat(1, self.encoding.n_output_dims))
 
@@ -295,6 +313,10 @@ 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
 
     def get_alpha(
         self,

nerfstudio/models/neus_facto.py

@@ -77,6 +77,29 @@ class NeuSFactoModelConfig(NeuSModelConfig):
     """Max num iterations for the annealing function."""
     use_single_jitter: bool = True
     """Whether use single jitter or not for the proposal networks."""
+    use_anneal_beta: bool = False
+    """whether to anneal beta of neus or not similar to bakedsdf"""
+    beta_anneal_max_num_iters: int = 1000_000
+    """max num iterations for the annealing function of beta"""
+    beta_anneal_init: float = 0.05
+    """initial beta for annealing function"""
+    beta_anneal_end: float = 0.0002
+    """final beta for annealing function"""
+    # TODO move to base model config since it can be used in all models
+    enable_progressive_hash_encoding: bool = False
+    """whether to use progressive hash encoding"""
+    enable_numerical_gradients_schedule: bool = False
+    """whether to use numerical gradients delta schedule"""
+    enable_curvature_loss_schedule: bool = False
+    """whether to use curvature loss weight schedule"""
+    curvature_loss_multi: float = 0.0
+    """curvature loss weight"""
+    curvature_loss_warmup_steps: int = 20_000
+    """curvature loss warmup steps"""
+    level_init: int = 4
+    """initial level of multi-resolution hash encoding"""
+    steps_per_level: int = 10_000
+    """steps per level of multi-resolution hash encoding"""
 
 
 class NeuSFactoModel(NeuSModel):
@@ -172,6 +195,99 @@ class NeuSFactoModel(NeuSModel):
                 )
             )
 
+        if self.config.use_anneal_beta:
+            # anneal the beta of volsdf before each training iterations
+            M = self.config.beta_anneal_max_num_iters
+            beta_init = self.config.beta_anneal_init
+            beta_end = self.config.beta_anneal_end
+
+            def set_beta(step):
+                # bakedsdf's beta schedule adapted to neus
+                train_frac = np.clip(step / M, 0, 1)
+                beta = beta_init / (1 + (beta_init - beta_end) / beta_end * (train_frac**0.8))
+                beta = np.log(1.0 / beta) / 10.0
+                self.field.deviation_network.variance.data[...] = beta
+
+            callbacks.append(
+                TrainingCallback(
+                    where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
+                    update_every_num_iters=1,
+                    func=set_beta,
+                )
+            )
+
+        # 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.0 / base_res
+        end_delta = 1.0 / max_res
+
+        if self.config.enable_numerical_gradients_schedule:
+
+            def set_delta(step):
+                delta = 1.0 / (base_res * growth_factor ** (step / steps_per_level))
+                delta = max(1.0 / (4.0 * max_res), delta)
+                self.field.set_numerical_gradients_delta(
+                    delta * 4.0
+                )  # TODO because we divide 4 to normalize points to [0, 1]
+
+            callbacks.append(
+                TrainingCallback(
+                    where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
+                    update_every_num_iters=1,
+                    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
+                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],
+                    update_every_num_iters=1,
+                    func=set_mask,
+                )
+            )
+
+        # schedule the curvature loss weight
+        # linear warmup for 5000 steps to 5e-4 and then decay as delta
+        self.curvature_loss_multi_factor = 1.0
+        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.0 / (
+                        base_res * growth_factor ** ((step - self.config.curvature_loss_warmup_steps) / steps_per_level)
+                    )
+                    delta = max(1.0 / (max_res * 10.0), delta)
+                    factor = delta / init_delta
+
+                self.curvature_loss_multi_factor = factor
+
+            callbacks.append(
+                TrainingCallback(
+                    where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
+                    update_every_num_iters=1,
+                    func=set_curvature_loss_mult_factor,
+                )
+            )
+
         return callbacks
 
     def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict[str, Any]:
@@ -208,7 +324,33 @@ class NeuSFactoModel(NeuSModel):
                 outputs["weights_list"], outputs["ray_samples_list"]
             )
 
+        # 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"]
+
+            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
+    
+    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_image_metrics_and_images(
         self, outputs: Dict[str, Any], batch: Dict[str, Any]