diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index 430a8aca8..d000515aa 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/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, }, diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index 2f730d2c5..807363789 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/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, diff --git a/nerfstudio/models/neus_facto.py b/nerfstudio/models/neus_facto.py index 5c1a3ec9d..33e73fd4b 100644 --- a/nerfstudio/models/neus_facto.py +++ b/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]