tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Fix torchfield #1749 #1539 (#1859)

Browse Source 
* fix bugs in TorchNerfactoField

* remove unused package
This commit is contained in:
Sophus 2023-05-05 00:25:30 +08:00 committed by GitHub
parent 61bad4dab3
commit 68f2edcb05
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
nerfstudio/fields/nerfacto_field.py
View File

@ -381,8 +381,8 @@ class TorchNerfactoField(Field):
positions = self.spatial_distortion(positions)
else:
positions = ray_samples.frustums.get_positions()
encoded_xyz = self.position_encoding(positions)
base_mlp_out = self.mlp_base(encoded_xyz)
encoded_xyz = self.position_encoding(positions.view(-1, 3)).to(torch.float32)
base_mlp_out = self.mlp_base(encoded_xyz).view(*ray_samples.frustums.shape, -1)
density = self.field_output_density(base_mlp_out)
return density, base_mlp_out
@ -404,18 +404,20 @@ class TorchNerfactoField(Field):
outputs = {}
for field_head in self.field_heads:
encoded_dir = self.direction_encoding(ray_samples.frustums.directions)
directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
encoded_dir = self.direction_encoding(directions_flat)
mlp_out = self.mlp_head(
torch.cat(
[
encoded_dir,
density_embedding, # type:ignore
encoded_dir.view(-1, self.direction_encoding.get_out_dim()),
density_embedding.view(-1, self.mlp_base.get_out_dim()), # type:ignore
embedded_appearance.view(-1, self.appearance_embedding_dim),
],
dim=-1, # type:ignore
)
)
outputs[field_head.field_head_name] = field_head(mlp_out)
outputs[field_head.field_head_name] = field_head(mlp_out).view(*outputs_shape, -1).to(directions)
return outputs