tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-16 04:43:52 +08:00
Code Activity

Fix typing errors (#229)

Browse Source
This commit is contained in:
Matthew Tancik authored and GitHub committed 2022-07-29 14:17:38 -07:00
1 parent be001e8ebe
commit 58212543f5
30 files changed
+236 -147
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

configs/graph_default.yaml
+2
View File
@@ -109,4 +109,6 @@ optimizers:
viewer:
enable: False
zmq_url: tcp://127.0.0.1:6000
min_render_image_height: 64
max_render_image_height: 1024
num_rays_per_chunk: 4096
environment/jupyter.py
-3
View File
@@ -1,3 +0,0 @@
## https://github.com/jupyterlab/jupyterlab/issues/2049
## Allow access to hidden files
c.ContentsManager.allow_hidden = True
nerfactory/cameras/camera_paths.py
+6 -5
View File
@@ -18,13 +18,14 @@ Code for camera paths.
import copy
from dataclasses import dataclass
from typing import List, Tuple
from typing import List, Optional, Tuple
import torch
import nerfactory.cameras.utils as camera_utils
from nerfactory.cameras.cameras import (
Camera,
PinholeCamera,
get_camera,
get_intrinsics_from_intrinsics_matrix,
)
@@ -38,7 +39,7 @@ class CameraPath:
cameras: List[Camera]
def get_interpolated_camera_path(camera_a: Camera, camera_b: Camera, steps: int) -> CameraPath:
def get_interpolated_camera_path(camera_a: PinholeCamera, camera_b: PinholeCamera, steps: int) -> CameraPath:
"""Generate a camera path between two cameras.
Args:
@@ -72,10 +73,10 @@ def get_interpolated_camera_path(camera_a: Camera, camera_b: Camera, steps: int)
def get_spiral_path(
camera: Camera,
camera: PinholeCamera,
steps: int = 30,
radius: float = None,
radiuses: Tuple[float] = None,
radius: Optional[float] = None,
radiuses: Optional[Tuple[float]] = None,
rots: int = 2,
zrate: float = 0.5,
) -> CameraPath:
nerfactory/cameras/cameras.py
+1 -2
View File
@@ -25,7 +25,6 @@ from torch.nn.functional import normalize
from torchtyping import TensorType
from nerfactory.cameras.rays import RayBundle
from nerfactory.utils.misc import is_not_none
class Camera:
@@ -157,7 +156,7 @@ class Camera:
camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1).to(device) # (num_rays, 3, 4)
coords = self.get_image_coords().to(device)
ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords)
if is_not_none(self.camera_index):
if self.camera_index is not None:
ray_bundle.set_camera_indices(camera_index=self.camera_index)
return ray_bundle
nerfactory/cameras/rays.py
+3 -4
View File
@@ -23,7 +23,6 @@ import torch
from torchtyping import TensorType
from nerfactory.utils.math import Gaussians, conical_frustum_to_gaussian
from nerfactory.utils.misc import is_not_none
from nerfactory.utils.tensor_dataclass import TensorDataclass
@@ -143,7 +142,7 @@ class RaySamples(TensorDataclass):
Returns:
RaySamples: New set of masked samples.
"""
if is_not_none(self.valid_mask):
if self.valid_mask is not None:
return self[self.valid_mask[..., 0]]
return self
@@ -169,7 +168,7 @@ class RayBundle(TensorDataclass):
nears: Optional[TensorType["num_rays", 1]] = None
fars: Optional[TensorType["num_rays", 1]] = None
valid_mask: Optional[TensorType["num_rays", 1, bool]] = None
num_rays_per_chunk: int = None
num_rays_per_chunk: Optional[int] = None
def set_camera_indices(self, camera_index: int) -> None:
"""Sets all of the the camera indices to a specific camera index.
@@ -229,7 +228,7 @@ class RayBundle(TensorDataclass):
dists = bin_ends - bin_starts # [..., N_samples, 1]
deltas = dists * torch.norm(self.directions[:, :], dim=-1)[..., None, None]
if is_not_none(self.camera_indices):
if self.camera_indices is not None:
camera_indices = self.camera_indices[..., None]
else:
camera_indices = None
nerfactory/cameras/utils.py
+3 -1
View File
@@ -143,7 +143,7 @@ def quaternion_from_matrix(matrix, isprecise=False):
K /= 3.0
# quaternion is eigenvector of K that corresponds to largest eigenvalue
w, V = np.linalg.eigh(K)
q = V[[3, 0, 1, 2], np.argmax(w)]
q = V[np.array([3, 0, 1, 2]), np.argmax(w)]
if q[0] < 0.0:
np.negative(q, q)
return q
@@ -167,6 +167,8 @@ def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True):
"""
q0 = unit_vector(quat0[:4])
q1 = unit_vector(quat1[:4])
if q0 is None or q1 is None:
raise ValueError("Input quaternions invalid.")
if fraction == 0.0:
return q0
elif fraction == 1.0:
nerfactory/cuda/backend.py
+1
View File
@@ -11,6 +11,7 @@
# 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.
"""Setup cuda backend."""
import glob
import os
from subprocess import DEVNULL, call
nerfactory/data/colmap_utils.py
+2 -3
View File
@@ -427,8 +427,7 @@ def read_model(path, ext=""):
elif detect_model_format(path, ".txt"):
ext = ".txt"
else:
print("Provide model format: '.bin' or '.txt'")
return
raise ValueError("Provide model format: '.bin' or '.txt'")
if ext == ".txt":
cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
@@ -489,7 +488,7 @@ def rotmat2qvec(R):
/ 3.0
)
eigvals, eigvecs = np.linalg.eigh(K)
qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
qvec = eigvecs[np.array([3, 0, 1, 2]), np.argmax(eigvals)]
if qvec[0] < 0:
qvec *= -1
return qvec
nerfactory/data/dataloader.py
+11 -2
View File
@@ -59,10 +59,10 @@ def setup_dataset_train(config: DataConfig, device: str) -> Tuple[DatasetInputs,
image_dataset_train = instantiate_from_dict_config(config.image_dataset_train, **dataset_inputs_train.as_dict())
# ImageSampler
image_sampler_train = instantiate_from_dict_config(
config.dataloader_train.image_sampler, dataset=image_dataset_train, device=device
config.dataloader_train.image_sampler, dataset=image_dataset_train, device=device # type: ignore
)
# PixelSampler
pixel_sampler_train = instantiate_from_dict_config(config.dataloader_train.pixel_sampler)
pixel_sampler_train = instantiate_from_dict_config(config.dataloader_train.pixel_sampler) # type: ignore
# Dataloader
dataloader_train = TrainDataloader(image_sampler_train, pixel_sampler_train)
return dataset_inputs_train, dataloader_train
@@ -255,6 +255,15 @@ class EvalDataloader: # pylint: disable=too-few-public-methods
self.device = device
self.kwargs = kwargs
@abstractmethod
def __iter__(self):
"""Iterates over the dataset"""
return self
@abstractmethod
def __next__(self) -> Tuple[RayBundle, Dict]:
"""Returns the next batch of data"""
def get_camera(self, image_idx) -> Camera:
"""Get camera for the given image index"""
intrinsics = self.intrinsics[image_idx].to(self.device)
nerfactory/data/datasets.py
+1 -1
View File
@@ -462,7 +462,7 @@ class Friends(Dataset):
camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale
# --- semantics ---
semantics = Semantics()
semantics = None
if self.include_semantics:
thing_filenames = [
image_filename.replace("/images/", "/segmentations/thing/").replace(".jpg", ".png")
nerfactory/data/image_dataset.py
+5 -5
View File
@@ -21,16 +21,16 @@ from typing import Dict, List, Optional, Union
import numpy as np
import numpy.typing as npt
import PIL
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchtyping import TensorType
from nerfactory.data.structs import Semantics
from nerfactory.utils.misc import is_not_none
class ImageDataset(torch.utils.data.Dataset):
class ImageDataset(Dataset):
"""Dataset that returns images."""
def __init__(
@@ -79,7 +79,7 @@ class ImageDataset(torch.utils.data.Dataset):
f"Image height {image_height} is not divisible by downscale_factor {self.downscale_factor}"
)
pil_image = pil_image.resize(
(image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.BILINEAR
(image_width // self.downscale_factor, image_height // self.downscale_factor), Image.BILINEAR
)
image = np.array(pil_image, dtype="uint8") # shape is (h, w, 3 or 4)
assert len(image.shape) == 3
@@ -115,10 +115,10 @@ class ImageDataset(torch.utils.data.Dataset):
data = {"image_idx": image_idx}
assert is_not_none(image)
data["image"] = image
if is_not_none(mask):
if mask is not None:
assert mask.shape[:2] == image.shape[:2]
data["mask"] = mask
if is_not_none(semantics):
if semantics is not None:
assert semantics.shape[:2] == image.shape[:2]
data["semantics"] = semantics
return data
nerfactory/data/image_sampler.py
+17 -6
View File
@@ -21,13 +21,23 @@ from abc import abstractmethod
from typing import Union
import torch
from torch.utils.data import default_collate
from torch.utils.data import IterableDataset, default_collate
from nerfactory.data.image_dataset import ImageDataset
from nerfactory.utils.misc import get_dict_to_torch
class ImageSampler(torch.utils.data.IterableDataset):
"""Samples image_batch's."""
class ImageSampler(IterableDataset):
"""Samples image_batch's.
Attributes:
dataset (torch.utils.data.Dataset): Dataset to sample from.
"""
# dataset: ImageDataset
def __init__(self, dataset: ImageDataset):
super().__init__()
self.dataset = dataset
@abstractmethod
def __getitem__(self, idx):
@@ -51,14 +61,14 @@ class CacheImageSampler(ImageSampler):
def __init__(
self,
dataset: torch.utils.data.Dataset,
dataset: ImageDataset,
num_images_to_sample_from: int = -1,
num_times_to_repeat_images: int = 0,
device: Union[torch.device, str] = "cpu",
):
super().__init__()
super().__init__(dataset=dataset)
self.dataset = dataset
self.cache_all_images = num_images_to_sample_from is -1
self.cache_all_images = num_images_to_sample_from == -1
self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
self.num_times_to_repeat_images = num_times_to_repeat_images
self.device = device
@@ -76,6 +86,7 @@ class CacheImageSampler(ImageSampler):
def _get_batch_list(self):
"""Returns a list of batches from the dataset attribute."""
indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from)
batch_list = [self.dataset.__getitem__(idx) for idx in indices]
return batch_list
nerfactory/data/pixel_sampler.py
+3 -2
View File
@@ -43,6 +43,7 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool
c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)]
image = batch["image"][c, y, x]
mask, semantics = None, None
if "mask" in batch:
mask = batch["mask"][c, y, x]
if "semantics" in batch:
@@ -57,9 +58,9 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool
"indices": indices, # with the abs camera indices
"image": image,
}
if "mask" in batch:
if mask is not None:
collated_batch["mask"] = mask
if "semantics" in batch:
if semantics is not None:
collated_batch["semantics"] = semantics
if keep_full_image:
nerfactory/data/structs.py
+14 -14
View File
@@ -17,7 +17,7 @@ Dataset input structures.
"""
from dataclasses import dataclass
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
import torch
from torchtyping import TensorType
@@ -37,12 +37,12 @@ class PointCloud:
class Semantics:
"""Dataclass for semantic labels."""
stuff_classes: Optional[List[str]] = None
stuff_colors: Optional[List[str]] = None
stuff_filenames: Optional[List[str]] = None
thing_classes: Optional[List[str]] = None
thing_colors: Optional[List[str]] = None
thing_filenames: Optional[List[str]] = None
stuff_filenames: List[str]
stuff_classes: List[str]
thing_filenames: List[str]
thing_classes: List[str]
stuff_colors: Optional[torch.Tensor] = None
thing_colors: Optional[torch.Tensor] = None
@dataclass
@@ -68,7 +68,7 @@ class SceneBounds:
diff = self.aabb[1] - self.aabb[0]
return self.aabb[0] + diff / 2.0
def get_centered_and_scaled_scene_bounds(self, scale_factor=1.0):
def get_centered_and_scaled_scene_bounds(self, scale_factor: Union[float, torch.Tensor] = 1.0):
"""Returns a new box that has been shifted and rescaled to be centered
about the origin."""
return SceneBounds(aabb=(self.aabb - self.get_center()) * scale_factor)
@@ -110,13 +110,13 @@ class DatasetInputs:
"""
image_filenames: List[str]
intrinsics: torch.Tensor
camera_to_world: torch.Tensor
downscale_factor: int = 1
intrinsics: torch.tensor = None
camera_to_world: torch.tensor = None
mask_filenames: List[str] = None
depth_filenames: List[str] = None
mask_filenames: Optional[List[str]] = None
depth_filenames: Optional[List[str]] = None
scene_bounds: SceneBounds = SceneBounds()
semantics: Semantics = Semantics()
semantics: Optional[Semantics] = None
point_cloud: PointCloud = PointCloud()
alpha_color: Optional[TensorType[3]] = None
@@ -151,5 +151,5 @@ class BaseDataContainer:
"""
rays: RayBundle # Raybundle and the cameras will be merged into one thing in a later PR
cameras: torch.tensor
cameras: torch.Tensor
ground_truth_pixels: Optional[TensorType["num_pixels", 3]] = None
nerfactory/engine/trainer.py
+43 -32
View File
@@ -18,23 +18,26 @@ Code to train model.
import functools
import logging
import os
from typing import Dict
from typing import Dict, List, Optional
import torch
import torch.distributed as dist
from torch.cuda.amp import GradScaler
from torch.cuda.amp.grad_scaler import GradScaler
from torch.nn.parallel import DistributedDataParallel as DDP
from torchtyping import TensorType
from nerfactory.cameras.rays import RayBundle
from nerfactory.data.dataloader import (
EvalDataloader,
TrainDataloader,
setup_dataset_eval,
setup_dataset_train,
)
from nerfactory.graphs.base import setup_graph
from nerfactory.optimizers.optimizers import setup_optimizers
from nerfactory.data.structs import DatasetInputs
from nerfactory.graphs.base import Graph, setup_graph
from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers
from nerfactory.utils import profiler, writer
from nerfactory.utils.callbacks import Callback
from nerfactory.utils.config import Config
from nerfactory.utils.decorators import check_main_thread
from nerfactory.utils.writer import EventName, TimeWriter
@@ -62,12 +65,12 @@ class Trainer:
self.mixed_precision = False
logging.warning("Mixed precision is disabled for CPU training.")
# dataset variables
self.dataset_inputs_train = None
self.dataloader_train = None
self.dataloader_eval = None
self.dataset_inputs_train: DatasetInputs
self.dataloader_train: TrainDataloader
self.dataloader_eval: EvalDataloader
# model variables
self.graph = None
self.optimizers = None
self.graph = Graph
self.optimizers: Optimizers
self.start_step = 0
# logging variables
writer.setup_event_writers(config.logging, max_iter=config.trainer.max_num_iterations)
@@ -76,7 +79,7 @@ class Trainer:
self.visualizer_state = viewer_utils.VisualizerState(config.viewer)
self.grad_scaler = GradScaler(enabled=self.mixed_precision)
# training callbacks
self.callbacks = None
self.callbacks: List[Callback]
def setup(self, test_mode=False):
"""Setup the Trainer by calling other setup functions.
@@ -89,14 +92,13 @@ class Trainer:
self.graph = setup_graph(self.config.graph, self.dataset_inputs_train, device=self.device)
self.optimizers = setup_optimizers(self.config.optimizers, self.graph.get_param_groups())
if self.config.trainer.resume_train.load_dir:
self._load_checkpoint()
self._load_checkpoint()
if self.world_size > 1:
self.graph = DDP(self.graph, device_ids=[self.local_rank])
dist.barrier(device_ids=[self.local_rank])
self.callbacks = self.graph.get_training_callbacks()
self.callbacks = self.graph.get_training_callbacks() # type: ignore
@classmethod
def get_aggregated_loss(cls, loss_dict: Dict[str, torch.Tensor]) -> torch.Tensor:
@@ -131,7 +133,7 @@ class Trainer:
if step % self.config.trainer.steps_per_test == 0:
self.eval_with_dataloader(self.dataloader_eval, step=step)
self._write_out_storage(step)
self.visualizer_state.update_scene(step, self.graph)
self.visualizer_state.update_scene(step, self.graph) # type: ignore
self._write_out_storage(num_iterations)
@@ -152,15 +154,18 @@ class Trainer:
def _load_checkpoint(self) -> None:
"""Helper function to load graph and optimizer from prespecified checkpoint"""
load_config = self.config.trainer.resume_train
load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt")
assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist"
loaded_state = torch.load(load_path, map_location="cpu")
self.start_step = loaded_state["step"] + 1
# load the checkpoints for graph and optimizer
self.graph.load_graph(loaded_state)
self.optimizers.load_optimizers(loaded_state)
self.grad_scaler.load_state_dict(loaded_state["scaler"])
logging.info("done loading checkpoint from %s", load_path)
if load_config.load_dir is not None and load_config.load_step is not None:
load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt")
assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist"
loaded_state = torch.load(load_path, map_location="cpu")
self.start_step = loaded_state["step"] + 1
# load the checkpoints for graph and optimizer
self.graph.load_graph(loaded_state) # type: ignore
self.optimizers.load_optimizers(loaded_state)
self.grad_scaler.load_state_dict(loaded_state["scaler"])
logging.info("done loading checkpoint from %s", load_path)
else:
logging.info("No checkpoints to load, training from scratch")
@check_main_thread
def _save_checkpoint(self, output_dir: str, step: int) -> None:
@@ -173,10 +178,14 @@ class Trainer:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt")
if hasattr(self.graph, "module"):
model = self.graph.module.state_dict()
else:
model = self.graph.state_dict()
torch.save(
{
"step": step,
"model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(),
"model": model,
"optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()},
"scalers": self.grad_scaler.state_dict(),
},
@@ -197,9 +206,9 @@ class Trainer:
"""
self.optimizers.zero_grad_all()
with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision):
_, loss_dict, metrics_dict = self.graph.forward(ray_indices, batch=batch)
_, loss_dict, metrics_dict = self.graph.forward(ray_indices=ray_indices, batch=batch) # type: ignore
loss = sum(loss_dict.values())
self.grad_scaler.scale(loss).backward()
self.grad_scaler.scale(loss).backward() # type: ignore
self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
self.grad_scaler.update()
@@ -213,7 +222,7 @@ class Trainer:
return loss_dict
@profiler.time_function
def test_image(self, camera_ray_bundle: RayBundle, batch: dict, step: int = None) -> float:
def test_image(self, camera_ray_bundle: RayBundle, batch: dict, step: Optional[int] = None) -> float:
"""Test a specific image.
Args:
@@ -224,14 +233,16 @@ class Trainer:
Returns:
float: PSNR
"""
self.graph.eval()
self.graph.eval() # type: ignore
if camera_ray_bundle.camera_indices is None:
raise ValueError("camera_ray_bundle.camera_indices is None during testing")
image_idx = int(camera_ray_bundle.camera_indices[0, 0])
outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs)
self.graph.train()
outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) # type: ignore
psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs) # type: ignore
self.graph.train() # type: ignore
return psnr
def eval_with_dataloader(self, dataloader: EvalDataloader, step: int = None) -> None:
def eval_with_dataloader(self, dataloader: EvalDataloader, step: Optional[int] = None) -> None:
"""Run evaluation with a given dataloader.
Args:
nerfactory/fields/base.py
+2 -3
View File
@@ -25,7 +25,6 @@ from torchtyping import TensorType
from nerfactory.cameras.rays import Frustums, RaySamples
from nerfactory.fields.modules.field_heads import FieldHeadNames
from nerfactory.utils.misc import is_not_none
class Field(nn.Module):
@@ -79,7 +78,7 @@ class Field(nn.Module):
"""
valid_mask = ray_samples.valid_mask
if is_not_none(valid_mask):
if valid_mask is not None:
# Hacky handling of empty masks. Tests on a single ray but doesn't use results
if not valid_mask.any():
ray_samples = RaySamples(frustums=Frustums.get_mock_frustum().to(valid_mask.device))
@@ -105,5 +104,5 @@ class Field(nn.Module):
density, density_embedding = self.get_density(ray_samples)
field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding)
field_outputs[FieldHeadNames.DENSITY] = density
field_outputs[FieldHeadNames.DENSITY] = density # type: ignore
return field_outputs
nerfactory/fields/density_fields/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2022 The Plenoptix 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.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
nerfactory/fields/density_fields/density_grid.py
+5 -3
View File
@@ -51,8 +51,10 @@ class DensityGrid(nn.Module):
on the scale but all shares the same center.
For example:
- `center=0`, `base_scale=3`, `num_cascades=1` would create one grid (lvl=0) covers the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]].
- `center=0`, `base_scale=3`, `num_cascades=2` would create one grid (lvl=0) covers the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]],
- `center=0`, `base_scale=3`, `num_cascades=1` would create one grid (lvl=0) covers
the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]].
- `center=0`, `base_scale=3`, `num_cascades=2` would create one grid (lvl=0) covers
the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]],
and one more grid (lvl=1) covers the aabb [[-3, -3, -3], [3, 3, 3]].
Normally you can simply set the `center` to the scene center, and `base_scale` to be the scene scale, and use
@@ -225,6 +227,6 @@ class DensityGrid(nn.Module):
densities = densities.view(*positions.shape[:-1], 1)
return densities
def forward(self, x):
def forward(self, x): # pylint: disable=no-self-use
"""Not implemented."""
raise RuntimeError("Shouldn't be called!")
nerfactory/fields/instant_ngp_field.py
+8 -3
View File
@@ -17,10 +17,11 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, ....
"""
from typing import Tuple
from typing import Optional, Tuple
import torch
from torch.nn.parameter import Parameter
from torchtyping import TensorType
from nerfactory.cameras.rays import RaySamples
from nerfactory.data.structs import SceneBounds
@@ -110,13 +111,17 @@ class TCNNInstantNGPField(Field):
density = trunc_exp(density_before_activation.to(positions))
return density, base_mlp_out
def get_outputs(self, ray_samples: RaySamples, density_embedding=None):
def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
# TODO: add valid_mask masking!
# tcnn requires directions in the range [0,1]
directions = get_normalized_directions(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
d = self.direction_encoding(directions_flat)
h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
if density_embedding is None:
positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
h = torch.cat([d, positions.view(-1, 3)], dim=-1)
else:
h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
return {FieldHeadNames.RGB: rgb}
nerfactory/fields/modules/base.py
+2 -1
View File
@@ -16,6 +16,7 @@
The field module baseclass.
"""
from abc import abstractmethod
from typing import Optional
from torch import nn
from torchtyping import TensorType
@@ -29,7 +30,7 @@ class FieldModule(nn.Module):
out_dim (int, optional): Ouput dimension to module.
"""
def __init__(self, in_dim: int = None, out_dim: int = None) -> None:
def __init__(self, in_dim: Optional[int] = None, out_dim: Optional[int] = None) -> None:
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
nerfactory/fields/modules/encoding.py
+8
View File
@@ -50,6 +50,8 @@ class Identity(Encoding):
"""Identity encoding (Does not modify input)"""
def get_out_dim(self) -> int:
if self.in_dim is None:
raise ValueError("Input dimension has not been set")
return self.in_dim
def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
@@ -72,6 +74,8 @@ class ScalingAndOffset(Encoding):
self.offset = offset
def get_out_dim(self) -> int:
if self.in_dim is None:
raise ValueError("Input dimension has not been set")
return self.in_dim
def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
@@ -101,6 +105,8 @@ class NeRFEncoding(Encoding):
self.include_input = include_input
def get_out_dim(self) -> int:
if self.in_dim is None:
raise ValueError("Input dimension has not been set")
out_dim = self.in_dim * self.num_frequencies * 2
if self.include_input:
out_dim += self.in_dim
@@ -156,6 +162,8 @@ class RFFEncoding(Encoding):
if not scale > 0:
raise ValueError("RFF encoding scale should be greater than zero")
self.scale = scale
if self.in_dim is None:
raise ValueError("Input dimension has not been set")
b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies))
self.register_buffer(name="b_matrix", tensor=b_matrix)
self.include_input = include_input
nerfactory/fields/modules/mlp.py
+5 -4
View File
@@ -42,7 +42,7 @@ class MLP(FieldModule):
num_layers: int,
layer_width: int,
out_dim: Optional[int] = None,
skip_connections: Tuple[int] = (),
skip_connections: Optional[Tuple[int]] = None,
activation: Optional[nn.Module] = nn.ReLU(),
out_activation: Optional[nn.Module] = None,
) -> None:
@@ -67,9 +67,10 @@ class MLP(FieldModule):
else:
for i in range(self.num_layers - 1):
if i == 0:
assert i not in list(self.skip_connections), "Skip connection at layer 0 doesn't make sense."
if self.skip_connections is not None:
assert i not in list(self.skip_connections), "Skip connection at layer 0 doesn't make sense."
layers.append(nn.Linear(self.in_dim, self.layer_width))
elif i in self.skip_connections:
elif self.skip_connections is not None and i in self.skip_connections:
layers.append(nn.Linear(self.layer_width + self.in_dim, self.layer_width))
else:
layers.append(nn.Linear(self.layer_width, self.layer_width))
@@ -87,7 +88,7 @@ class MLP(FieldModule):
"""
x = in_tensor
for i, layer in enumerate(self.layers):
if i in self.skip_connections:
if self.skip_connections is not None and i in self.skip_connections:
x = torch.cat([in_tensor, x], -1)
x = layer(x)
if self.activation:
nerfactory/graphs/base.py
+14 -12
View File
@@ -39,18 +39,6 @@ from nerfactory.utils.misc import (
)
@profiler.time_function
def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) -> "Graph":
"""Setup the graph. The dataset inputs should be set with the training data.
Args:
dataset_inputs: The inputs which will be used to define the camera parameters.
"""
graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict())
graph.to(device)
return graph
class AbstractGraph(nn.Module):
"""Highest level graph class. Somewhat useful to lift code up and out of the way."""
@@ -243,3 +231,17 @@ class Graph(AbstractGraph):
def load_graph(self, loaded_state: Dict[str, Any]) -> None:
"""Load the checkpoint from the given path"""
self.load_state_dict({key.replace("module.", ""): value for key, value in loaded_state["model"].items()})
@profiler.time_function
def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) -> Graph:
"""Setup the graph. The dataset inputs should be set with the training data.
Args:
dataset_inputs: The inputs which will be used to define the camera parameters.
"""
graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict())
graph.to(device)
if not isinstance(graph, Graph):
raise TypeError(f"Unable to instantiate Graph, got {type(graph)}")
return graph
nerfactory/optimizers/optimizers.py
+8 -6
View File
@@ -16,19 +16,21 @@
Optimizers class.
"""
from typing import Any, Dict, List
from typing import Any, Dict, List, Union
import numpy as np
from hydra.utils import instantiate
from omegaconf import DictConfig
from torch.cuda.amp import GradScaler
from torch.nn import Parameter
from torch.cuda.amp.grad_scaler import GradScaler
from torch.nn.parameter import Parameter
from torch.optim.lr_scheduler import LambdaLR
from nerfactory.utils import writer
def setup_optimizers(config: DictConfig, param_groups: Dict[str, List[Parameter]]) -> "Optimizers":
def setup_optimizers(
config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]]
) -> "Optimizers":
"""Helper to set up the optimizers
Args:
@@ -82,7 +84,7 @@ class Optimizers:
param_groups: A dictionary of parameter groups to optimize.
"""
def __init__(self, config: DictConfig, param_groups: Dict[str, List[Parameter]]):
def __init__(self, config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]]):
self.config = config
self.optimizers = {}
self.schedulers = {}
@@ -108,7 +110,7 @@ class Optimizers:
Args:
param_group_name: name of scheduler to step forward
"""
if self.config.param_group_name.scheduler:
if self.config.param_group_name.scheduler: # type: ignore
self.schedulers[param_group_name].step()
def zero_grad_all(self) -> None:
nerfactory/renderers/renderers.py
+14 -12
View File
@@ -51,10 +51,10 @@ class RGBRenderer(nn.Module):
@classmethod
def combine_rgb(
cls,
rgb: TensorType[..., "num_samples", 3],
weights: TensorType[..., "num_samples", 1],
rgb: TensorType["bs":..., "num_samples", 3],
weights: TensorType["bs":..., "num_samples", 1],
background_color: Optional[TensorType[3]] = None,
) -> TensorType[..., 3]:
) -> TensorType["bs":..., 3]:
"""Composite samples along ray and render color image
Args:
@@ -76,9 +76,9 @@ class RGBRenderer(nn.Module):
def forward(
self,
rgb: TensorType[..., "num_samples", 3],
weights: TensorType[..., "num_samples", 1],
) -> TensorType[..., 3]:
rgb: TensorType["bs":..., "num_samples", 3],
weights: TensorType["bs":..., "num_samples", 1],
) -> TensorType["bs":..., 3]:
"""Composite samples along ray and render color image
Args:
@@ -149,8 +149,8 @@ class AccumulationRenderer(nn.Module):
@classmethod
def forward(
cls,
weights: TensorType[..., "num_samples", 1],
) -> TensorType:
weights: TensorType["bs":..., "num_samples", 1],
) -> TensorType["bs":..., 1]:
"""Composite samples along ray and calculate accumulation.
Args:
@@ -205,8 +205,8 @@ class UncertaintyRenderer(nn.Module):
@classmethod
def forward(
cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples", 1]
) -> TensorType:
cls, betas: TensorType["bs":..., "num_samples", 1], weights: TensorType["bs":..., "num_samples", 1]
) -> TensorType["bs":..., 1]:
"""Calculate uncertainty along the ray.
Args:
@@ -225,8 +225,10 @@ class SemanticRenderer(nn.Module):
@classmethod
def forward(
cls, semantics: TensorType[..., "num_samples", "num_classes"], weights: TensorType[..., "num_samples", 1]
) -> TensorType[..., "num_classes"]:
cls,
semantics: TensorType["bs":..., "num_samples", "num_classes"],
weights: TensorType["bs":..., "num_samples", 1],
) -> TensorType["bs":..., "num_classes"]:
"""_summary_"""
sem = torch.sum(weights * semantics, dim=-2)
return sem
nerfactory/utils/config.py
+1 -1
View File
@@ -122,7 +122,7 @@ class Config:
data: DataConfig = MISSING
graph: GraphConfig = MISSING
optimizers: Dict[str, Any] = MISSING
viewer: Dict[str, Any] = MISSING
viewer: ViewerConfig = MISSING
# additional optional parameters here
hydra: Optional[Dict[str, Any]] = None
nerfactory/utils/misc.py
+11 -7
View File
@@ -20,7 +20,7 @@ import hashlib
import json
from math import floor, log
from pydoc import locate
from typing import Any, Callable, Dict
from typing import Any, Callable, Dict, Union
import torch
from omegaconf import DictConfig
@@ -34,11 +34,11 @@ class DotDict(dict):
def __getattr__(self, attr):
return self[attr]
__setattr__ = dict.__setitem__
__delattr__ = dict.__delitem__
__setattr__ = dict.__setitem__ # type: ignore
__delattr__ = dict.__delitem__ # type: ignore
def get_dict_to_torch(stuff, device="cpu"):
def get_dict_to_torch(stuff, device: Union[torch.device, str] = "cpu"):
"""Set everything in the dict to the specified torch device."""
if isinstance(stuff, dict):
for k, v in stuff.items():
@@ -74,13 +74,17 @@ def get_masked_dict(d, mask):
return masked_dict
def instantiate_from_dict_config(dict_config: DictConfig, **kwargs):
"""Our version of hydra's instantiate function."""
def instantiate_from_dict_config(dict_config: Any, **kwargs):
"""Our version of hydra's instantiate function.
Args:
dict_config: DictConfig object to instantiate. It can be a dataclass or a dict but must have a `_target_` field.
"""
dict_config_kwargs = {k: v for k, v in dict_config.items() if k != "_target_"}
uninstantiated_class = locate(dict_config._target_) # pylint: disable=protected-access
all_kwargs = dict_config_kwargs
all_kwargs.update(kwargs)
instantiated_class = uninstantiated_class(**all_kwargs)
instantiated_class = uninstantiated_class(**all_kwargs) # type: ignore
return instantiated_class
nerfactory/utils/writer.py
+2 -2
View File
@@ -202,8 +202,8 @@ class TimeWriter:
self.step = step
self.write = write
self.start = None
self.duration = None
self.start: float = 0.0
self.duration: float = 0.0
def __enter__(self):
self.start = time()
pyproject.toml
+17 -1
View File
@@ -43,4 +43,20 @@ disable = [
addopts = "--typeguard-packages=nerfactory --torchtyping-patch-typeguard --disable-warnings --ignore=tests/cuda --ignore=tests/viewer"
testpaths = [
"tests",
]
]
# pyright
[tool.pyright]
include = ["nerfactory"]
exclude = ["**/node_modules",
"**/__pycache__",
]
ignore = ["nerfactory/viewer", "nerfactory/cuda"]
defineConstant = { DEBUG = true }
reportMissingImports = true
reportMissingTypeStubs = false
reportPrivateImportUsage = false
pythonVersion = "3.8"
pythonPlatform = "Linux"
tests/viewer/test_viewer.py
+14 -12
View File
@@ -12,9 +12,9 @@ import umsgpack
from tqdm import tqdm
import nerfactory.viewer.server.geometry as g
from nerfactory.cameras.camera_paths import InterpolatedCameraPath
from nerfactory.cameras import camera_paths
from nerfactory.cameras.cameras import get_camera
from nerfactory.data.utils import get_dataset_inputs
from nerfactory.data.datasets import Blender
from nerfactory.utils.io import get_absolute_path
from nerfactory.viewer.server.transformations import (
get_rotation_matrix,
@@ -24,6 +24,7 @@ from nerfactory.viewer.server.viewer_utils import get_default_vis, set_camera
def test_drawing():
"""Test drawing objects in the viewer."""
vis = get_default_vis()
vis.delete()
@@ -110,20 +111,21 @@ def test_drawing():
def test_streaming():
pass
"""Test streaming of data to the viewer."""
def test_rendering():
pass
"""Test rendering."""
def test_camera_trajectory():
print("test_camera_trajectory")
"""Test generating camera trajectory."""
vis = get_default_vis()
# vis.delete()
# sample poses from a dataset
dataset_inputs = get_dataset_inputs(data_directory="data/blender/lego", dataset_format="blender", split="train")
dataset = Blender(data_directory="data/blender/lego")
dataset_inputs = dataset.get_dataset_inputs(split="train")
num_cameras = len(dataset_inputs.intrinsics)
intrinsics = dataset_inputs.intrinsics
camera_to_world = dataset_inputs.camera_to_world
@@ -136,11 +138,10 @@ def test_camera_trajectory():
estimated_time = num_steps / fps
print("estimated_time:", estimated_time)
camera_path = InterpolatedCameraPath(camera_a, camera_b)
cameras = camera_path.get_path(steps=num_steps)
camera_path = camera_paths.get_interpolated_camera_path(camera_a, camera_b, steps=num_steps)
start = time.time()
for camera in cameras:
for camera in camera_path.cameras:
set_camera(vis, camera)
time.sleep(1 / fps)
time_elapsed = time.time() - start
@@ -148,8 +149,9 @@ def test_camera_trajectory():
def test_send_image_stream():
"""Test sending image stream."""
vis = get_default_vis()
vis_Background = vis["/Background"]
vis_background = vis["/Background"]
video_filename = get_absolute_path("data/instant_ngp/bear/bear.MOV")
images = []
@@ -162,16 +164,16 @@ def test_send_image_stream():
ret, image = cap.read()
# break
num_images = len(images)
fps = 30
for image in tqdm(images):
rgba = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
encoded = umsgpack.Ext(0x12, rgba.tobytes())
vis_Background.set_image(encoded)
vis_background.set_image(encoded)
time.sleep(1 / fps)
def test_perspective_camera():
"""Test perspective camera in viewer."""
vis = get_default_vis()
vis.set_object(g.Box([0.5, 0.5, 0.5]))
camera = g.PerspectiveCamera(fov=90)