From 5cab9224e4b64473c8fa96cc6ae936831470dd10 Mon Sep 17 00:00:00 2001 From: AdamRashid96 Date: Tue, 19 Sep 2023 07:12:24 +0000 Subject: [PATCH] License + test fixes --- .../data/dataparsers/sdfstudio_dataparser.py | 8 +-- nerfstudio/field_components/encodings.py | 60 ++++++++++++++++++- nerfstudio/fields/sdf_field.py | 1 + nerfstudio/models/neuralangelo.py | 2 +- 4 files changed, 65 insertions(+), 6 deletions(-) diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index 708676ed8..ac0b6a95e 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -1,4 +1,4 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. @@ -41,7 +41,7 @@ CONSOLE = Console() def get_src_from_pairs( - ref_idx, all_imgs, pairs_srcs, neighbors_num=None, neighbors_shuffle=False + ref_idx: int, all_imgs: TensorType, pairs_srcs: TensorType, neighbors_num=None, neighbors_shuffle=False ) -> Dict[str, TensorType]: # src_idx[0] is ref img src_idx = pairs_srcs[ref_idx] @@ -130,7 +130,7 @@ def get_sparse_sfm_points(image_idx: int, sfm_points): # sfm points sparse_sfm_points = sfm_points[image_idx] - sparse_sfm_points = BasicImages([sparse_sfm_points]) + # sparse_sfm_points = BasicImages([sparse_sfm_points]) return {"sparse_sfm_points": sparse_sfm_points} @@ -361,7 +361,7 @@ class SDFStudio(DataParser): 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 + [get_image(image_filename) for image_filename in sorted(image_filenames)], dim=0 ).cuda() additional_inputs_dict["pairs"] = { diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index c8f536f79..546a64d90 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -491,11 +491,13 @@ class TensorVMEncoding(Encoding): resolution: int = 128, num_components: int = 24, init_scale: float = 0.1, + smoothstep: bool = False, ) -> None: super().__init__(in_dim=3) self.resolution = resolution self.num_components = num_components + self.smoothstep = smoothstep self.plane_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, resolution))) self.line_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, 1))) @@ -503,6 +505,61 @@ class TensorVMEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components * 3 + def index_fn(self, x: torch.Tensor, y: torch.Tensor, width: int, height: int): + y.clamp_max_(height - 1) + x.clamp_max_(width - 1) + + if y.max() >= height or x.max() >= width: + breakpoint() + + index = y * width + x + feature_offset = width * height * torch.arange(3) + index += feature_offset.to(x.device)[:, None, None] + + return index.long() + + def grid_sample_2d(self, feature, coord, type="plane"): + if type == "plane": + height, width = self.resolution, self.resolution + else: + height, width = self.resolution, 1 + + scaled = coord * torch.tensor([width, height]).to(coord.device)[None, None] + scaled_c = torch.ceil(scaled).type(torch.int32) + scaled_f = torch.floor(scaled).type(torch.int32) + + offset = scaled - scaled_f + + # smooth version of offset + if self.smoothstep: + offset = offset * offset * (3.0 - 2.0 * offset) + + offset = offset[..., None, :] + + index_0 = self.index_fn(scaled_c[..., 0:1], scaled_c[..., 1:2], height, width) # [..., num_levels] + index_2 = self.index_fn(scaled_f[..., 0:1], scaled_f[..., 1:2], height, width) + index_1 = self.index_fn(scaled_c[..., 0:1], scaled_f[..., 1:2], height, width) + index_3 = self.index_fn(scaled_f[..., 0:1], scaled_c[..., 1:2], height, width) + + # breakpoint() + if type == "plane": + f_0 = feature[index_0] # [..., num_levels, features_per_level] + f_1 = feature[index_1] + f_2 = feature[index_2] + f_3 = feature[index_3] + + f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1]) + f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1]) + + f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2]) + + return f0312 + else: + f_0 = feature[index_0] # [..., num_levels, features_per_level] + f_2 = feature[index_2] + f_02 = f_0 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1]) + return f_02 + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: """Compute encoding for each position in in_positions @@ -519,7 +576,8 @@ class TensorVMEncoding(Encoding): plane_coord = plane_coord.view(3, -1, 1, 2).detach() line_coord = line_coord.view(3, -1, 1, 2).detach() - plane_features = F.grid_sample(self.plane_coef, plane_coord, align_corners=True) # [3, Components, -1, 1] + # plane_features = F.grid_sample(self.plane_coef, plane_coord, align_corners=True) # [3, Components, -1, 1] + plane_features = self.grid_sample_2d(self.plane_coef, plane_coord, type="plane") # [3, -1, 1, Components] line_features = F.grid_sample(self.line_coef, line_coord, align_corners=True) # [3, Components, -1, 1] features = plane_features * line_features # [3, Components, -1, 1] diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index a8491ea12..973d78802 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/nerfstudio/fields/sdf_field.py @@ -457,6 +457,7 @@ class SDFField(Field): if self.spatial_distortion is not None and not skip_spatial_distortion: x = self.spatial_distortion(x) + points_sdf = torch.zeros_like(x[..., :1]) # compute gradient in contracted space if self.config.use_numerical_gradients: # https://github.com/bennyguo/instant-nsr-pl/blob/main/models/geometry.py#L173 diff --git a/nerfstudio/models/neuralangelo.py b/nerfstudio/models/neuralangelo.py index 88f796b6b..ce043615f 100644 --- a/nerfstudio/models/neuralangelo.py +++ b/nerfstudio/models/neuralangelo.py @@ -1,4 +1,4 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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.