From d50c71ac43c687b73dbf5d139d3af8e8073ff43a Mon Sep 17 00:00:00 2001 From: evonneng Date: Sat, 10 Sep 2022 23:09:44 -0700 Subject: [PATCH] update docstrings and random cleanups (#387) * fixing docstrings and small nits; up to fileds * finished adding and fixing docstrings --- docs/reference/api/config.rst | 8 + docs/reference/api/index.rst | 1 + nerfactory/cameras/camera_paths.py | 3 +- nerfactory/cameras/cameras.py | 12 +- nerfactory/cameras/rays.py | 96 +++++------ nerfactory/cameras/utils.py | 152 ++++++------------ nerfactory/configs/base.py | 129 ++++++++++++++- nerfactory/datamanagers/base.py | 6 +- nerfactory/datamanagers/dataloaders.py | 35 ++-- nerfactory/datamanagers/dataparsers/base.py | 16 +- .../dataparsers/blender_parser.py | 5 - .../dataparsers/friends_parser.py | 13 +- .../dataparsers/instant_ngp_parser.py | 8 +- .../dataparsers/mipnerf_parser.py | 10 +- .../dataparsers/nerfactory_parser.py | 10 +- .../dataparsers/record3d_parser.py | 10 +- nerfactory/datamanagers/datasets.py | 31 ++-- nerfactory/datamanagers/pixel_sampler.py | 25 ++- nerfactory/datamanagers/structs.py | 65 +++++--- nerfactory/engine/trainer.py | 37 ++--- nerfactory/fields/base.py | 26 ++- .../fields/density_fields/density_grid.py | 42 ++--- nerfactory/fields/instant_ngp_field.py | 28 +++- nerfactory/fields/modules/base.py | 17 +- nerfactory/fields/modules/embedding.py | 10 +- nerfactory/fields/modules/encoding.py | 10 +- nerfactory/fields/nerf_field.py | 18 +-- nerfactory/fields/nerfw_field.py | 20 +-- nerfactory/fields/semantic_nerf_field.py | 13 +- nerfactory/models/base.py | 51 +++--- nerfactory/models/instant_ngp.py | 3 +- nerfactory/models/mipnerf.py | 6 +- nerfactory/models/mipnerf_360.py | 6 +- nerfactory/models/modules/ray_generator.py | 3 +- nerfactory/models/modules/ray_losses.py | 5 +- nerfactory/models/modules/ray_sampler.py | 111 +++++++------ nerfactory/models/modules/scene_colliders.py | 25 ++- nerfactory/models/nerfw.py | 6 +- nerfactory/models/semantic_nerf.py | 7 +- nerfactory/models/tensorf.py | 6 +- nerfactory/models/vanilla_nerf.py | 7 +- nerfactory/pipelines/base.py | 30 +++- nerfactory/utils/callbacks.py | 16 +- nerfactory/utils/colmap_utils.py | 62 +++---- nerfactory/utils/decorators.py | 6 +- nerfactory/utils/math.py | 45 +++--- nerfactory/utils/misc.py | 51 ++++-- nerfactory/utils/poses.py | 18 +-- nerfactory/utils/profiler.py | 6 +- nerfactory/utils/tensor_dataclass.py | 37 +++-- nerfactory/utils/writer.py | 113 +++++++------ nerfactory/viewer/server/viewer_utils.py | 16 +- 52 files changed, 846 insertions(+), 646 deletions(-) create mode 100644 docs/reference/api/config.rst diff --git a/docs/reference/api/config.rst b/docs/reference/api/config.rst new file mode 100644 index 000000000..3a807c20a --- /dev/null +++ b/docs/reference/api/config.rst @@ -0,0 +1,8 @@ +.. _configs: + +Configs +============ + +.. automodule:: nerfactory.configs.base + :members: + :show-inheritance: diff --git a/docs/reference/api/index.rst b/docs/reference/api/index.rst index 431430a5e..e9443b495 100644 --- a/docs/reference/api/index.rst +++ b/docs/reference/api/index.rst @@ -7,6 +7,7 @@ TODO: Explanation of each component .. toctree:: + config data fields fields/modules/index diff --git a/nerfactory/cameras/camera_paths.py b/nerfactory/cameras/camera_paths.py index e6f86837a..96b5050ad 100644 --- a/nerfactory/cameras/camera_paths.py +++ b/nerfactory/cameras/camera_paths.py @@ -30,8 +30,7 @@ def get_interpolated_camera_path(cameras: Cameras, steps: int) -> Cameras: """Generate a camera path between two cameras. Args: - camera_a: The first camera. - camera_b: The second camera. + cameras: Cameras object containing intrinsics of all cameras. steps: The number of steps to interpolate between the two cameras. Returns: diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py index 13eb520be..3e537fc01 100644 --- a/nerfactory/cameras/cameras.py +++ b/nerfactory/cameras/cameras.py @@ -57,6 +57,8 @@ class Cameras: fy: Focal length y. cx: Principal point x. cy: Principal point y. + width: Image width. + height: Image height. distortion_params: OpenCV 6 radial distortion coefficients. camera_type: Type of camera model. """ @@ -114,10 +116,10 @@ class Cameras: def to(self, device: Union[torch.device, str]) -> "Cameras": """ Args: - Device to move the camera to. + device: Device to move the camera onto Returns: - Cameras: Cameras on the specified device. + Cameras on the specified device. """ distortion_params = self.distortion_params.to(device) if self.distortion_params is not None else None return Cameras( @@ -135,10 +137,10 @@ class Cameras: def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["height", "width", 2]: """ Args: - pixel_offset (float): Offset for each pixel. Defaults to center of pixel (0.5) + pixel_offset: Offset for each pixel. Defaults to center of pixel (0.5) Returns: - TensorType["image_height", "image_width", 2]: Grid of image coordinates. + Grid of image coordinates. """ image_height = self.image_height image_width = self.image_width @@ -251,7 +253,7 @@ class Cameras: """Convert a camera to a json dictionary. Args: - camera_idx (int): Index of the camera to convert. + camera_idx: Index of the camera to convert. image: An image in range [0, 1] that is encoded to a base64 string. Defaults to None. max_size: Max size to resize the image to. Defaults to None. diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index e6d2d8c35..9ea9bdcc2 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -28,35 +28,32 @@ from nerfactory.utils.tensor_dataclass import TensorDataclass @dataclass class Frustums(TensorDataclass): - """Describes region of space as a frustum. - - Attributes: - origins: xyz coordinate for ray origin. - directions: Direction of ray. - starts: Where the frustum starts along a ray. - ends: Where the frustum ends along a ray. - pixel_area: Projected area of pixel a distance 1 away from origin. - """ + """Describes region of space as a frustum.""" origins: TensorType["bs":..., 3] + """xyz coordinate for ray origin.""" directions: TensorType["bs":..., 3] + """Direction of ray.""" starts: TensorType["bs":..., 1] + """Where the frustum starts along a ray.""" ends: TensorType["bs":..., 1] + """Where the frustum ends along a ray.""" pixel_area: TensorType["bs":..., 1] + """Projected area of pixel a distance 1 away from origin.""" def get_positions(self) -> TensorType[..., 3]: """Calulates "center" position of frustum. Not weighted by mass. Returns: - TensorType[..., 3]: xyz positions. + xyz positions. """ return self.origins + self.directions * (self.starts + self.ends) / 2 def get_gaussian_blob(self) -> Gaussians: - """Calculates of guassian approximation of conical frustum. + """Calculates guassian approximation of conical frustum. Resturns: - Gaussians: Conical frustums approximated by gaussian distribution. + Conical frustums approximated by gaussian distribution. """ # Cone radius is set such that the square pixel_area matches the cone area. cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159 # r = sqrt(pixel_area / pi) @@ -73,7 +70,7 @@ class Frustums(TensorDataclass): """Helper function to generate a placeholder frustum. Returns: - Frustums: A size 1 frustum with meaningless values. + A size 1 frustum with meaningless values. """ return Frustums( origins=torch.ones((1, 3)).to(device), @@ -86,30 +83,27 @@ class Frustums(TensorDataclass): @dataclass class RaySamples(TensorDataclass): - """Samples along a ray - - Args: - frustums: Frustums along ray. - camera_indices: Camera index. - valid_mask: Rays that are valid. - deltas: "width" of each sample. - ts: sample-camera distance of each sample. - """ + """Samples along a ray""" frustums: Frustums + """Frustums along ray.""" camera_indices: Optional[TensorType["bs":..., 1]] = None + """Camera index.""" valid_mask: Optional[TensorType["bs":..., 1]] = None + """Rays that are valid.""" deltas: Optional[TensorType["bs":..., 1]] = None + """"width" of each sample.""" metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None + """addtional information relevant to generating ray samples""" def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]: """Return weights based on predicted densities Args: - densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray + densities: Predicted densities for samples along ray Returns: - TensorType[..., "num_samples", 1]: Weights for each sample + Weights for each sample """ delta_density = self.deltas * densities @@ -134,14 +128,18 @@ class RaySamples(TensorDataclass): return weights def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: - """Sets valid mask""" + """Sets valid mask + + Args: + valid_mask: the mask to use + """ self.valid_mask = valid_mask def apply_masks(self) -> "RaySamples": """Use valid_mask to mask samples. Returns: - RaySamples: New set of masked samples. + New set of masked samples. """ if self.valid_mask is not None: return self[self.valid_mask[..., 0]] @@ -150,36 +148,33 @@ class RaySamples(TensorDataclass): @dataclass class RayBundle(TensorDataclass): - """A bundle of ray parameters. - - Args: - origins: Ray origins (XYZ) - directions: Unit ray direction vector - pixel_area: Projected area of pixel a distance 1 away from origin - camera_indices: Camera indices - nears: Distance along ray to start sampling - fars: Rays Distance along ray to stop sampling - valid_mask: Rays that are valid - num_rays_per_chunk: Number of rays per chunk - metadata: Additional metadata or data needed for interpolation, will mimic shape of rays - """ + """A bundle of ray parameters.""" # TODO(ethan): make sure the sizes with ... are correct origins: TensorType[..., 3] + """Ray origins (XYZ)""" directions: TensorType[..., 3] + """Unit ray direction vector""" pixel_area: TensorType[..., 1] + """Projected area of pixel a distance 1 away from origin""" camera_indices: Optional[TensorType[..., 1]] = None + """Camera indices""" nears: Optional[TensorType[..., 1]] = None + """Distance along ray to start sampling""" fars: Optional[TensorType[..., 1]] = None + """Rays Distance along ray to stop sampling""" valid_mask: Optional[TensorType[..., 1, bool]] = None + """Rays that are valid""" num_rays_per_chunk: int = 128 + """Number of rays per chunk""" metadata: Optional[Dict[str, TensorType["num_rays", "latent_dims"]]] = None + """Additional metadata or data needed for interpolation, will mimic shape of rays""" def set_camera_indices(self, camera_index: int) -> None: """Sets all of the the camera indices to a specific camera index. Args: - camera_index (int): Camera index. + camera_index: Camera index. """ self.camera_indices = torch.ones_like(self.origins[..., 0:1]).long() * camera_index @@ -191,10 +186,10 @@ class RayBundle(TensorDataclass): """Returns a RayBundle as a subset of rays. Args: - num_rays (int): Number of rays in output RayBundle + num_rays: Number of rays in output RayBundle Returns: - RayBundle: RayBundle with subset of rays. + RayBundle with subset of rays. """ assert num_rays <= len(self) indices = random.sample(range(len(self)), k=num_rays) @@ -204,11 +199,11 @@ class RayBundle(TensorDataclass): """Flattens RayBundle and extracts chunk given start and end indicies. Args: - start_idx (int): Start index of RayBundle chunk. - end_idx (int): End index of RayBundle chunk. + start_idx: Start index of RayBundle chunk. + end_idx: End index of RayBundle chunk. Returns: - RayBundle: Flattened RayBundle with end_idx-start_idx rays. + Flattened RayBundle with end_idx-start_idx rays. """ return self.flatten()[start_idx:end_idx] @@ -216,17 +211,14 @@ class RayBundle(TensorDataclass): def get_ray_samples( self, bin_starts: TensorType["bs":..., "num_samples", 1], bin_ends: TensorType["bs":..., "num_samples", 1] ) -> RaySamples: - """Produces samples for each ray by projection points along the ray direction. - - Currently samples uniformly. - + """Produces samples for each ray by projection points along the ray direction. Currently samples uniformly. Args: - bin_starts (TensorType[..., "num_samples", 1]): Distance from origin to start of bin. - bin_ends (TensorType[..., "num_samples", 1]): Distance from origin to end of bin. + bin_starts: Distance from origin to start of bin. + bin_ends: Distance from origin to end of bin. Returns: - RaySamples: Samples projected along ray. + Samples projected along ray. """ device = self.origins.device diff --git a/nerfactory/cameras/utils.py b/nerfactory/cameras/utils.py index 2cd458b3a..804acb93c 100644 --- a/nerfactory/cameras/utils.py +++ b/nerfactory/cameras/utils.py @@ -18,7 +18,7 @@ Camera transformation helper code. """ import math -from typing import Literal, Tuple +from typing import List, Literal, Optional, Tuple import numpy as np import torch @@ -27,80 +27,31 @@ from torchtyping import TensorType _EPS = np.finfo(float).eps * 4.0 -def unit_vector(data, axis=None, out=None): +def unit_vector(data, axis: Optional[int] = None) -> np.ndarray: """Return ndarray normalized by length, i.e. Euclidean norm, along axis. - >>> v0 = np.random.random(3) - >>> v1 = unit_vector(v0) - >>> np.allclose(v1, v0 / np.linalg.norm(v0)) - True - >>> v0 = np.random.rand(5, 4, 3) - >>> v1 = unit_vector(v0, axis=-1) - >>> v2 = v0 / np.expand_dims(np.sqrt(np.sum(v0*v0, axis=2)), 2) - >>> np.allclose(v1, v2) - True - >>> v1 = unit_vector(v0, axis=1) - >>> v2 = v0 / np.expand_dims(np.sqrt(np.sum(v0*v0, axis=1)), 1) - >>> np.allclose(v1, v2) - True - >>> v1 = np.empty((5, 4, 3)) - >>> unit_vector(v0, axis=1, out=v1) - >>> np.allclose(v1, v2) - True - >>> list(unit_vector([])) - [] - >>> list(unit_vector([1])) - [1.0] + + Args: + axis: the axis along which to normalize into unit vector + out: where to write out the data to. If None, returns a new np ndarray """ - if out is None: - data = np.array(data, dtype=np.float64, copy=True) - if data.ndim == 1: - data /= math.sqrt(np.dot(data, data)) - return data - else: - if out is not data: - out[:] = np.array(data, copy=False) - data = out + data = np.array(data, dtype=np.float64, copy=True) + if data.ndim == 1: + data /= math.sqrt(np.dot(data, data)) + return data length = np.atleast_1d(np.sum(data * data, axis)) np.sqrt(length, length) if axis is not None: length = np.expand_dims(length, axis) data /= length - if out is None: - return data + return data -def quaternion_from_matrix(matrix, isprecise=False): +def quaternion_from_matrix(matrix, isprecise: bool = False) -> np.ndarray: """Return quaternion from rotation matrix. - If isprecise is True, the input matrix is assumed to be a precise rotation - matrix and a faster algorithm is used. - >>> q = quaternion_from_matrix(np.identity(4), True) - >>> np.allclose(q, [1, 0, 0, 0]) - True - >>> q = quaternion_from_matrix(np.diag([1, -1, -1, 1])) - >>> np.allclose(q, [0, 1, 0, 0]) or np.allclose(q, [0, -1, 0, 0]) - True - >>> R = rotation_matrix(0.123, (1, 2, 3)) - >>> q = quaternion_from_matrix(R, True) - >>> np.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786]) - True - >>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0], - ... [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]] - >>> q = quaternion_from_matrix(R) - >>> np.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611]) - True - >>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0], - ... [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]] - >>> q = quaternion_from_matrix(R) - >>> np.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603]) - True - >>> R = random_rotation_matrix() - >>> q = quaternion_from_matrix(R) - >>> is_same_transform(R, quaternion_matrix(q)) - True - >>> R = euler_matrix(0.0, 0.0, np.pi/2.0) - >>> np.allclose(quaternion_from_matrix(R, isprecise=False), - ... quaternion_from_matrix(R, isprecise=True)) - True + + Args: + matrix: rotation matrix to obtain quaternion + isprecise: if True, input matrix is assumed to be precise rotation matrix and a faster algorithm is used. """ M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4] if isprecise: @@ -151,21 +102,14 @@ def quaternion_from_matrix(matrix, isprecise=False): return q -def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True): +def quaternion_slerp(quat0, quat1, fraction: float, spin: int = 0, shortestpath: bool = True) -> np.ndarray: """Return spherical linear interpolation between two quaternions. - >>> q0 = random_quaternion() - >>> q1 = random_quaternion() - >>> q = quaternion_slerp(q0, q1, 0) - >>> np.allclose(q, q0) - True - >>> q = quaternion_slerp(q0, q1, 1, 1) - >>> np.allclose(q, q1) - True - >>> q = quaternion_slerp(q0, q1, 0.5) - >>> angle = math.acos(np.dot(q0, q)) - >>> (np.allclose(2, math.acos(np.dot(q0, q1)) / angle) or - np.allclose(2, math.acos(-np.dot(q0, q1)) / angle)) - True + Args: + quat0: first quaternion + quat1: second quaternion + fraction: how much to interpolate between quat0 vs quat1 (if 0, closer to quat0; if 1, closer to quat1) + spin: how much of an additional spin to place on the interpolation + shortestpath: whether to return the short or long path to rotation """ q0 = unit_vector(quat0[:4]) q1 = unit_vector(quat1[:4]) @@ -192,17 +136,11 @@ def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True): return q0 -def quaternion_matrix(quaternion): +def quaternion_matrix(quaternion) -> np.ndarray: """Return homogeneous rotation matrix from quaternion. - >>> M = quaternion_matrix([0.99810947, 0.06146124, 0, 0]) - >>> np.allclose(M, rotation_matrix(0.123, [1, 0, 0])) - True - >>> M = quaternion_matrix([1, 0, 0, 0]) - >>> np.allclose(M, np.identity(4)) - True - >>> M = quaternion_matrix([0, 1, 0, 0]) - >>> np.allclose(M, np.diag([1, -1, -1, 1])) - True + + Args: + quaternion: value to convert to matrix """ q = np.array(quaternion, dtype=np.float64, copy=True) n = np.dot(q, q) @@ -220,13 +158,12 @@ def quaternion_matrix(quaternion): ) -def get_interpolated_poses(poseA, poseB, steps=10): +def get_interpolated_poses(poseA, poseB, steps: int = 10) -> List[float]: """Return interpolation of poses with specified number of steps. Args: - poseA - poseB - Return: - list of poses + poseA: first pose + poseB: second pose + steps: number of steps the interpolated pose path should contain """ quatA = quaternion_from_matrix(poseA[:3, :3]) @@ -245,11 +182,14 @@ def get_interpolated_poses(poseA, poseB, steps=10): return posesAB -def get_interpolated_K(KA, KB, steps=10): +def get_interpolated_K(KA, KB, steps: int = 10) -> TensorType[3, 4]: """ - Interpolate between two camera poses. + Returns interpolated path between two camera poses with specified number of steps. + Args: - Returns: + KA: camera matrix 1 + KB: camera matrix 2 + steps: number of steps the interpolated pose path should contain """ Ks = [] ts = np.linspace(0, 1, steps) @@ -288,7 +228,7 @@ def get_interpolated_poses_many( def get_swirl_poses(pose): """ Returns a swirl around the camera poses. - # TODO(ethan): add parameters and make this work + # TODO(ethan): add parameters and make this work + add docstrings + return type """ N_frames = 30 * 4 dx = np.linspace(0, 0.03, N_frames) @@ -302,12 +242,12 @@ def get_swirl_poses(pose): return poses -def normalize(x): +def normalize(x) -> TensorType[...]: """Returns a normalized vector.""" return x / torch.linalg.norm(x) -def viewmatrix(lookat, up, pos): +def viewmatrix(lookat, up, pos) -> TensorType[...]: """Returns a camera transformation matrix. Args: @@ -316,7 +256,7 @@ def viewmatrix(lookat, up, pos): pos: The position of the camera. Returns: - torch.Tensor: A camera transformation matrix. + A camera transformation matrix. """ vec2 = normalize(lookat) vec1_avg = normalize(up) @@ -333,7 +273,7 @@ def get_distortion_params( k4: float = 0.0, p1: float = 0.0, p2: float = 0.0, -): +) -> TensorType[...]: """Returns a distortion parameters matrix. Args: @@ -349,8 +289,6 @@ def get_distortion_params( return torch.Tensor([k1, k2, k3, k4, p1, p2]) -# Addapted from MultiNeRF -# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509 def radial_and_tangential_undistort( coords: TensorType["num_coords":..., 2], distortion_params: TensorType["num_coords":..., 6], @@ -358,6 +296,8 @@ def radial_and_tangential_undistort( max_iterations=10, ) -> TensorType["num_coords":..., 2]: """Computes undistorted coords given opencv distortion parameters. + Addapted from MultiNeRF + https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L477-L509 Args: coords: The distorted coordinates. @@ -389,8 +329,6 @@ def radial_and_tangential_undistort( return torch.stack([x, y], dim=-1) -# Adapted from MultiNeRF: -# https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474 def _compute_residual_and_jacobian( x: TensorType["num_coords":...], y: TensorType["num_coords":...], @@ -406,6 +344,8 @@ def _compute_residual_and_jacobian( TensorType["num_coords":...], ]: """Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians. + Adapted from MultiNeRF: + https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474 Args: x: The updated x coordinates. diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 2e3f657eb..a34dec8d8 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -81,10 +81,15 @@ class MachineConfig(PrintableConfig): """Configuration of machine setup""" seed: int = 42 + """random seed initilization""" num_gpus: int = 1 + """total number of gpus available for train/eval""" num_machines: int = 1 + """total number of distributed machines available (for DDP)""" machine_rank: int = 0 + """current machine's rank (for DDP)""" dist_url: str = "auto" + """distributed connection point (for DDP)""" # Logging related configs @@ -93,9 +98,13 @@ class TensorboardWriterConfig(InstantiateConfig): """Tensorboard Writer config""" _target: Type = writer.TensorboardWriter + """target class to instantiate""" enable: bool = False + """if True enables tensorboard logging, else disables""" relative_log_dir: Path = Path("./") + """relative path to save all tensorboard events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set + """auto populated absolute path to saved tensorboard events [Do not set!]""" @dataclass @@ -103,9 +112,13 @@ class WandbWriterConfig(InstantiateConfig): """WandDB Writer config""" _target: Type = writer.WandbWriter + """target class to instantiate""" enable: bool = False + """if True enables wandb logging, else disables""" relative_log_dir: Path = Path("./") + """relative path to save all wandb events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set + """auto populated absolute path to saved wandb events [Do not set!]""" @dataclass @@ -113,7 +126,9 @@ class LocalWriterConfig(InstantiateConfig): """Local Writer config""" _target: Type = writer.LocalWriter + """target class to instantiate""" enable: bool = False + """if True enables local logging, else disables""" stats_to_track: Tuple[writer.EventName, ...] = ( writer.EventName.ITER_TRAIN_TIME, writer.EventName.TRAIN_RAYS_PER_SEC, @@ -121,9 +136,13 @@ class LocalWriterConfig(InstantiateConfig): writer.EventName.VIS_RAYS_PER_SEC, writer.EventName.TEST_RAYS_PER_SEC, ) + """specifies which stats will be logged/printed to terminal""" max_log_size: int = 10 + """maximum number of rows to print before wrapping. if 0, will print everything.""" relative_log_dir: Path = Path("./") + """relative local path to save all events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set + """auto populated absolute local path to saved events [Do not set!]""" def setup(self, banner_messages: Optional[List[str]] = None, **kwargs) -> Any: """Instantiate local writer @@ -139,14 +158,19 @@ class LoggingConfig(PrintableConfig): """Configuration of loggers and profilers""" steps_per_log: int = 10 + """number of steps between logging stats""" max_buffer_size: int = 20 + """maximum history size to keep for computing running averages of stats. + e.g. if 20, averages will be computed over past 20 occurances.""" writer: Tuple[Any, ...] = ( TensorboardWriterConfig(enable=True), WandbWriterConfig(enable=False), LocalWriterConfig(enable=True), ) - # profiler logs run times of functions and prints at end of training + """list of all supported writers. Can turn on/off writers by specifying enable.""" enable_profiler: bool = True + """whether to enable profiling code; prints speed of functions at the end of a program. + profiler logs run times of functions and prints at end of training""" # Trainer related configs @@ -155,15 +179,24 @@ class TrainerConfig(PrintableConfig): """Configuration for training regimen""" steps_per_save: int = 1000 + """number of steps between saves""" steps_per_test: int = 500 + """number of steps between eval""" max_num_iterations: int = 1000000 + """maximum number of iterations to run""" mixed_precision: bool = False + """whether or not to use mixed precision for training""" relative_model_dir: Path = Path("nerfactory_models/") + """relative path to save all checkpoints""" model_dir: Optional[Path] = None # full model dir path to be dynamically set + """auto populated absolute path to saved checkpoints [Do not set!]""" # optional parameters if we want to resume training load_dir: Optional[Path] = None + """optionally specify a pre-trained model directory to load from""" load_step: Optional[int] = None + """optionally specify model step to load from; if none, will find most recent model in load_dir""" load_config: Optional[Path] = None + """optionally specify a pre-defined config to load from""" # Dataset related configs @@ -172,6 +205,7 @@ class DataParserConfig(InstantiateConfig): """Basic dataset config""" _target: Type = DataParser + """_target: target class to instantiate""" @dataclass @@ -179,11 +213,17 @@ class NerfactoryDataParserConfig(DataParserConfig): """Nerfactory dataset config""" _target: Type = Nerfactory + """target class to instantiate""" data_directory: Path = Path("data/ours/posterv2") + """directory specifying location of data""" scale_factor: float = 1.0 + """How much to scale the camera origins by.""" downscale_factor: int = 1 + """How much to downscale images. Defaults to 1.""" scene_scale: float = 4.0 + """How much to scale the scene. Defaults to 0.33""" orientation_method: Literal["pca", "up"] = "up" + """The method to use for orientation. Either "pca" or "up".""" @dataclass @@ -191,9 +231,13 @@ class BlenderDataParserConfig(DataParserConfig): """Blender dataset config""" _target: Type = Blender + """target class to instantiate""" data_directory: Path = Path("data/blender/lego") + """directory specifying location of data""" scale_factor: float = 1.0 + """How much to scale the camera origins by.""" alpha_color: str = "white" + """alpha color of background""" @dataclass @@ -201,8 +245,11 @@ class FriendsDataParserConfig(DataParserConfig): """Friends dataset config""" _target: Type = Friends + """target class to instantiate""" data_directory: Path = Path("data/friends/TBBT-big_living_room") + """directory specifying location of data""" include_semantics: bool = True + """whether or not to include loading of semantics data""" @dataclass @@ -210,11 +257,17 @@ class MipNerf360DataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Mipnerf360 + """target class to instantiate""" data_directory: Path = Path("data/mipnerf_360/garden") + """directory specifying location of data""" downscale_factor: int = 1 + """How much to downscale images. Defaults to 1.""" val_skip: int = 8 + """1/val_skip images to use for validation. Defaults to 8.""" auto_scale: bool = True + """Scale based on pose bounds. Defaults to True.""" aabb_scale = 4 + """Scene scale, Defaults to 1.0.""" @dataclass @@ -222,9 +275,13 @@ class InstantNGPDataParserConfig(DataParserConfig): """Instant-NGP dataset config""" _target: Type = InstantNGP + """target class to instantiate""" data_directory: Path = Path("data/ours/posterv2") + """directory specifying location of data""" scale_factor: float = 1.0 + """How much to scale the camera origins by.""" scene_scale: float = 0.33 + """How much to scale the scene. Defaults to 0.33""" @dataclass @@ -232,27 +289,43 @@ class Record3DDataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Record3D + """target class to instantiate""" data_directory: Path = Path("data/record3d/garden") + """Location of data""" val_skip: int = 8 + """1/val_skip images to use for validation. Defaults to 8.""" aabb_scale = 4.0 + """Scene scale, Defaults to 4.0.""" max_dataset_size: int = 150 + """Max number of images to train on. If the dataset has + more, images will be sampled approximately evenly. Defaults to 150.""" @dataclass class VanillaDataManagerConfig(InstantiateConfig): - """Configuration for data manager instantiation""" + """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; + After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked + train/eval data at each iteration + """ # Note: eval_dataparser is annotated with Fixed[] to prevent dcargs from trying to # convert Optional[InstantiateConfig] into subcommands for choosing between None and # InstantiateConfig. _target: Type = VanillaDataManager + """target class to instantiate""" train_dataparser: DataParserConfig = BlenderDataParserConfig() + """specifies the dataparser used to unpack the data""" train_num_rays_per_batch: int = 1024 + """number of rays per batch to use per training iteration""" train_num_images_to_sample_from: int = -1 + """number of images to sample during training iteration""" eval_dataparser: dcargs.conf.Fixed[Optional[InstantiateConfig]] = None + """optionally specify different dataparser to use during eval; if None, uses train_dataparser""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) + """specifies the image indices to use during eval; if None, uses all""" eval_num_rays_per_chunk: int = 4096 + """specifies number of rays per chunk during eval""" @dataclass @@ -260,7 +333,9 @@ class FriendsDataManagerConfig(VanillaDataManagerConfig): """Friends data manager config""" _target: Type = VanillaDataManager + """target class to instantiate""" train_dataparser: DataParserConfig = FriendsDataParserConfig() + """specifies the dataparser used to unpack the data""" # Model related configs @@ -269,13 +344,21 @@ class ModelConfig(InstantiateConfig): """Configuration for model instantiation""" _target: Type = Model + """target class to instantiate""" enable_collider: bool = True + """Whether to create a scene collider to filter rays.""" collider_params: Dict[str, float] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) + """parameters to instantiate scene collider with""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0}) + """Loss specific weights.""" num_coarse_samples: int = 64 + """Number of samples in coarse field evaluation. Defaults to 64""" num_importance_samples: int = 128 + """Number of samples in fine field evaluation. Defaults to 128""" field_implementation: Literal["torch", "tcnn"] = "torch" + """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class""" enable_density_field: bool = False + """Whether to create a density field to filter samples.""" density_field_params: Dict[str, Any] = to_immutable_dict( { "center": 0.0, # simply set it as the center of the scene bbox @@ -285,6 +368,7 @@ class ModelConfig(InstantiateConfig): "update_every_num_iters": 16, } ) + """parameters to instantiate density field with""" @dataclass @@ -292,11 +376,17 @@ class InstantNGPModelConfig(ModelConfig): """Instant NGP Model Config""" _target: Type = NGPModel - enable_density_field: bool = True + """target class to instantiate""" enable_collider: bool = False - field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... + """Whether to create a scene collider to filter rays.""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) + """Loss specific weights.""" + field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... + """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" + enable_density_field: bool = True + """Whether to create a density field to filter samples.""" num_samples: int = 1024 # instead of course/fine samples + """Number of samples in field evaluation. Defaults to 1024,""" @dataclass @@ -304,15 +394,26 @@ class NerfWModelConfig(ModelConfig): """NerfW model config""" _target: Type = NerfWModel + """target class to instantiate""" loss_coefficients: Dict[str, float] = to_immutable_dict( {"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0, "uncertainty_loss": 1.0, "density_loss": 0.01} ) + """Loss specific weights.""" num_coarse_samples: int = 64 + """Number of samples in coarse field evaluation. Defaults to 64,""" num_importance_samples: int = 64 + """Number of samples in fine field evaluation. Defaults to 128,""" uncertainty_min: float = 0.03 + """This is added to the end of the uncertainty + rendering operation. It's called 'beta_min' in other repos. + This avoids calling torch.log() on a zero value, which would be undefined. + Defaults to 0.03.""" num_images: int = 10000 # TODO: don't hardcode this + """How many images exist in the dataset.""" appearance_embedding_dim: int = 48 + """Dimension of appearance embedding. Defaults to 48.""" transient_embedding_dim: int = 16 + """Dimension of transient embedding. Defaults to 16.""" @dataclass @@ -320,10 +421,15 @@ class TensoRFModelConfig(ModelConfig): """TensoRF model config""" _target: Type = TensoRFModel + """target class to instantiate""" init_resolution: int = 128 + """initial render resolution""" final_resolution: int = 200 + """final render resolution""" upsampling_iters: Tuple[int, ...] = (5000, 5500, 7000) + """specifies a list of iteration step numbers to perform upsampling""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "feature_loss": 8e-5}) + """Loss specific weights.""" # Pipeline related configs @@ -332,8 +438,11 @@ class PipelineConfig(InstantiateConfig): """Configuration for pipeline instantiation""" _target: Type = Pipeline + """target class to instantiate""" datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() + """specifies the datamanager config""" model: ModelConfig = ModelConfig() + """specifies the model config""" # Viewer related configs @@ -342,14 +451,20 @@ class ViewerConfig(PrintableConfig): """Configuration for viewer instantiation""" log_filename: Optional[Path] = None + """Filename to use for the log file. Defaults to None. If None, no log file is created.""" enable: bool = False - train: bool = True + """whether to enable viewer""" + start_train: bool = True + """whether to immediately start training upon loading viewer + if False, will just visualize dataset but you can toggle training in viewer""" zmq_url: str = "tcp://127.0.0.1:6000" + """the zmq port to connect to for communication""" launch_bridge_server: bool = True + """whether or not to launch the zmq bridge server""" websocket_port: int = 7007 - min_render_image_height: int = 64 - max_render_image_height: int = 1024 + """the default websocket port to connect to""" num_rays_per_chunk: int = 32768 + """number of rays per chunk to render with visualizer""" # Optimizer related configs diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py index e8e842e47..ab7b3162a 100644 --- a/nerfactory/datamanagers/base.py +++ b/nerfactory/datamanagers/base.py @@ -201,7 +201,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method under the hood. We may clean this up a little bit under the hood with more standard dataloading components that can be strung together, but it can be just used as a black box for now since only the constructor is likely to change in the future, or maybe passing in step number to the - next_train and next_eval functions.""" + next_train and next_eval functions. + + Args: + config: the DataManagerConfig used to instantiate class + """ config: cfg.VanillaDataManagerConfig diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py index 9a7716f8c..c58f7e400 100644 --- a/nerfactory/datamanagers/dataloaders.py +++ b/nerfactory/datamanagers/dataloaders.py @@ -35,10 +35,10 @@ class CacheImageDataloader(DataLoader): Creates batches of the InputDataset return type. Args: - dataset (torch.utils.data.Dataset): Dataset to sample from. - num_samples_to_collate (int, optional): How many images to sample rays for each batch. Defaults to all images. - num_times_to_repeat_images (int): How often to collate new images. Defaults to every iteration. - device (Union[torch.device, str]): Device to perform computation. Defaults to "cpu". + dataset: Dataset to sample from. + num_samples_to_collate: How many images to sample rays for each batch. Defaults to all images. + num_times_to_repeat_images: How often to collate new images. Defaults to every iteration. + device: Device to perform computation. Defaults to "cpu". """ def __init__( @@ -104,7 +104,6 @@ class EvalDataloader(DataLoader): Args: input_dataset: InputDataset to load data from - cameras: Cameras to use for generating rays num_rays_per_chunk: Number of camera rays to generate per chunk device: Device to load data to """ @@ -153,8 +152,12 @@ class EvalDataloader(DataLoader): ) return camera - def get_data_from_image_idx(self, image_idx) -> Tuple[RayBundle, Dict]: - """Returns the data for a specific image index.""" + def get_data_from_image_idx(self, image_idx: int) -> Tuple[RayBundle, Dict]: + """Returns the data for a specific image index. + + Args: + image_idx: Camera image index + """ ray_bundle = self.cameras.generate_rays(camera_indices=image_idx) ray_bundle.num_rays_per_chunk = self.num_rays_per_chunk ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int() @@ -164,7 +167,14 @@ class EvalDataloader(DataLoader): class FixedIndicesEvalDataloader(EvalDataloader): - """Dataloader that returns a fixed set of indices.""" + """Dataloader that returns a fixed set of indices. + + Args: + input_dataset: InputDataset to load data from + num_rays_per_chunk: Number of camera rays to generate per chunk + image_indices: List of image indices to load data from. If None, then use all images. + device: Device to load data to + """ def __init__( self, @@ -174,14 +184,6 @@ class FixedIndicesEvalDataloader(EvalDataloader): device: Union[torch.device, str] = "cpu", **kwargs, ): - """ - Args: - input_dataset: InputDataset to load data from - cameras: Cameras to use for generating rays - num_rays_per_chunk: Number of camera rays to generate per chunk - image_indices: List of image indices to load data from. If None, then use all images. - device: Device to load data to - """ super().__init__(input_dataset, num_rays_per_chunk, device, **kwargs) if image_indices is None: self.image_indices = list(range(len(input_dataset))) @@ -209,7 +211,6 @@ class RandIndicesEvalDataloader(EvalDataloader): Args: input_dataset: InputDataset to load data from - cameras: Cameras to use for generating rays num_rays_per_chunk: Number of camera rays to generate per chunk device: Device to load data to """ diff --git a/nerfactory/datamanagers/dataparsers/base.py b/nerfactory/datamanagers/dataparsers/base.py index a3101c896..d8cb43a10 100644 --- a/nerfactory/datamanagers/dataparsers/base.py +++ b/nerfactory/datamanagers/dataparsers/base.py @@ -25,7 +25,11 @@ from nerfactory.datamanagers.structs import DatasetInputs @dataclass class DataParser: - """A dataset.""" + """A dataset. + + Args: + config: datasetparser config containing all information needed to instantiate dataset + """ def __init__(self, config: cfg.DataParserConfig): super().__init__() @@ -33,23 +37,23 @@ class DataParser: @abstractmethod def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: - """Returns the dataset inputs for the given split. + """Abstract method that returns the dataset inputs for the given split. Args: - split: Which dataset split to generate. + split: Which dataset split to generate (train/test). Returns: - DatasetInputs + DatasetInputs containing data for the specified dataset and split """ def get_dataset_inputs(self, split: str = "train") -> DatasetInputs: """Returns the dataset inputs for the given split. Args: - split: Which dataset split to generate. + split: Which dataset split to generate (train/test). Returns: - DatasetInputs + DatasetInputs containing data for the specified dataset and split """ dataset_inputs = self._generate_dataset_inputs(split) return dataset_inputs diff --git a/nerfactory/datamanagers/dataparsers/blender_parser.py b/nerfactory/datamanagers/dataparsers/blender_parser.py index 69a6ee9f9..970c6a1aa 100644 --- a/nerfactory/datamanagers/dataparsers/blender_parser.py +++ b/nerfactory/datamanagers/dataparsers/blender_parser.py @@ -34,11 +34,6 @@ from nerfactory.utils.io import get_absolute_path, load_from_json class Blender(DataParser): """Blender Dataset Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. - - Args: - data_directory: Location of data - alpha_color: Sets transparent regions to specified color, otherwise black. - scale_factor: How much to scale the camera origins by. """ def __init__(self, config: cfg.BlenderDataParserConfig): diff --git a/nerfactory/datamanagers/dataparsers/friends_parser.py b/nerfactory/datamanagers/dataparsers/friends_parser.py index 2dbdf782f..c5ff3be1d 100644 --- a/nerfactory/datamanagers/dataparsers/friends_parser.py +++ b/nerfactory/datamanagers/dataparsers/friends_parser.py @@ -33,7 +33,12 @@ CONSOLE = Console() def get_semantics_and_masks(image_idx: int, semantics: Semantics): - """function to process additional semantics and mask information""" + """function to process additional semantics and mask information + + Args: + image_idx: specific image index to work with + semantics: semantics data + """ # handle mask person_index = semantics.thing_classes.index("person") thing_image_filename = semantics.thing_filenames[image_idx] @@ -55,7 +60,11 @@ class Friends(DataParser): @classmethod def _get_aabb_and_transform(cls, basedir): - """Returns the aabb and pointcloud transform from the threejs.json file.""" + """Returns the aabb and pointcloud transform from the threejs.json file. + + Args: + basedir: base directory to load from + """ filename = basedir / "threejs.json" assert filename.exists() data = load_from_json(filename) diff --git a/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py index 9e22d007e..952b83948 100644 --- a/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py +++ b/nerfactory/datamanagers/dataparsers/instant_ngp_parser.py @@ -33,13 +33,7 @@ from nerfactory.utils.io import get_absolute_path, load_from_json @dataclass class InstantNGP(DataParser): - """Instant NGP Dataset - - Args: - data_directory: Location of data - scale_factor: How much to scale the camera origins by. - scene_scale: How much to scale the scene. Defaults to 0.33 - """ + """Instant NGP Dataset""" config: cfg.InstantNGPDataParserConfig diff --git a/nerfactory/datamanagers/dataparsers/mipnerf_parser.py b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py index ecfe9172f..58c8cb1ec 100644 --- a/nerfactory/datamanagers/dataparsers/mipnerf_parser.py +++ b/nerfactory/datamanagers/dataparsers/mipnerf_parser.py @@ -30,15 +30,7 @@ from nerfactory.utils.io import get_absolute_path @dataclass class Mipnerf360(DataParser): - """MipNeRF 360 Dataset - - Args: - data_directory: Location of data - downscale_factor: How much to downscale images. Defaults to 1. - val_skip: 1/val_skip images to use for validation. Defaults to 8. - auto_scale: Scale based on pose bounds. Defaults to True. - aabb_scale: Scene scale, Defaults to 1.0. - """ + """MipNeRF 360 Dataset""" config: cfg.MipNerf360DataParserConfig diff --git a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py index 50493a358..36ea925e5 100644 --- a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py +++ b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py @@ -31,15 +31,7 @@ from nerfactory.utils.io import get_absolute_path, load_from_json @dataclass class Nerfactory(DataParser): - """Nerfactory Dataset - - Args: - data_directory: Location of data - scale_factor: How much to scale the camera origins by. - downscale_factor: How much to downscale images. - scene_scale: How much to scale the scene. - orientation_method: The method to use for orientation. Either "pca" or "up". - """ + """Nerfactory Dataset""" config: cfg.NerfactoryDataParserConfig diff --git a/nerfactory/datamanagers/dataparsers/record3d_parser.py b/nerfactory/datamanagers/dataparsers/record3d_parser.py index f966142e5..e5ae26410 100644 --- a/nerfactory/datamanagers/dataparsers/record3d_parser.py +++ b/nerfactory/datamanagers/dataparsers/record3d_parser.py @@ -31,15 +31,7 @@ from nerfactory.utils.io import get_absolute_path, load_from_json @dataclass class Record3D(DataParser): - """Record3D Dataset - - Args: - data_directory: Location of data - val_skip: 1/val_skip images to use for validation. Defaults to 8. - aabb_scale: Scene scale, Defaults to 4.0. - max_dataset_size: Max number of images to train on. If the dataset has - more, images will be sampled approximately evenly. Defaults to 150. - """ + """Record3D Dataset""" config: cfg.Record3DDataParserConfig diff --git a/nerfactory/datamanagers/datasets.py b/nerfactory/datamanagers/datasets.py index 50744b7d9..f5f577bd7 100644 --- a/nerfactory/datamanagers/datasets.py +++ b/nerfactory/datamanagers/datasets.py @@ -32,15 +32,13 @@ from nerfactory.utils.misc import is_not_none class InputDataset(Dataset): - """Dataset that returns images.""" + """Dataset that returns images. + + Args: + config: the DataParserConfig used to instantiate class, which in turn is used to load dataset inputs + """ def __init__(self, config: Union[cfg.DataParserConfig, cfg.InstantiateConfig], split: str): - """_summary_ - - Args: - image_filenames (List[str]): List of image filenames - alpha_color (TensorType[3], optional): Sets transparent regions to specified color, otherwise black. - """ super().__init__() self.dataset_inputs: DatasetInputs = config.setup().get_dataset_inputs(split=split) @@ -48,13 +46,10 @@ class InputDataset(Dataset): return len(self.dataset_inputs.image_filenames) def get_numpy_image(self, image_idx: int) -> npt.NDArray[np.uint8]: - """Returns the image. + """Returns the image of shape (H, W, 3 or 4). Args: - image_idx (int): The image index in the dataset. - - Returns: - np.uint8: an image of shape (H, W, 3 or 4) + image_idx: The image index in the dataset. """ image_filename = self.dataset_inputs.image_filenames[image_idx] pil_image = Image.open(image_filename) @@ -65,7 +60,11 @@ class InputDataset(Dataset): return image def get_image(self, image_idx: int) -> TensorType["image_height", "image_width", "num_channels"]: - """Returns a 3 channel image.""" + """Returns a 3 channel image. + + Args: + image_idx: The image index in the dataset. + """ image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0) if self.dataset_inputs.alpha_color is not None and image.shape[-1] == 4: assert image.shape[-1] == 4 @@ -75,7 +74,11 @@ class InputDataset(Dataset): return image def get_data(self, image_idx) -> Dict: - """Returns the ImageDataset data as a dictionary.""" + """Returns the ImageDataset data as a dictionary. + + Args: + image_idx: The image index in the dataset. + """ image = self.get_image(image_idx) data = {"image_idx": image_idx} assert is_not_none(image) diff --git a/nerfactory/datamanagers/pixel_sampler.py b/nerfactory/datamanagers/pixel_sampler.py index 61bd1df95..4d1beaa89 100644 --- a/nerfactory/datamanagers/pixel_sampler.py +++ b/nerfactory/datamanagers/pixel_sampler.py @@ -17,15 +17,21 @@ Code for sampling pixels. """ import random +from typing import Dict import torch -def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool = False): +def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. Returns a collated batch which is input to the Graph. It will sample only within the valid 'mask' if it's specified. + + Args: + batch: batch of images to sample from + num_rays_per_batch: number of rays to sample per batch + keep_full_image: whether or not to include a reference to the full image in returned batch """ device = batch["image"].device num_images, image_height, image_width, _ = batch["image"].shape @@ -70,14 +76,23 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool class PixelSampler: # pylint: disable=too-few-public-methods - """Samples 'pixel_batch's from 'image_batch's.""" + """Samples 'pixel_batch's from 'image_batch's. - def __init__(self, num_rays_per_batch, keep_full_image=False) -> None: + Args: + num_rays_per_batch: number of rays to sample per batch + keep_full_image: whether or not to include a reference to the full image in returned batch + """ + + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: self.num_rays_per_batch = num_rays_per_batch self.keep_full_image = keep_full_image - def sample(self, image_batch): - """Sample an image batch and return a pixel batch.""" + def sample(self, image_batch: Dict): + """Sample an image batch and return a pixel batch. + + Args: + image_batch: batch of images to sample from + """ pixel_batch = collate_image_dataset_batch( image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image ) diff --git a/nerfactory/datamanagers/structs.py b/nerfactory/datamanagers/structs.py index f30222fd4..11e039aff 100644 --- a/nerfactory/datamanagers/structs.py +++ b/nerfactory/datamanagers/structs.py @@ -32,7 +32,9 @@ class PointCloud: """Dataclass for a point cloud.""" xyz: TensorType["num_points", 3] = None + """the 3D points in point cloud""" rgb: TensorType["num_points", 3] = None + """the color of each point in point cloud""" @dataclass @@ -40,24 +42,27 @@ class Semantics: """Dataclass for semantic labels.""" stuff_filenames: List[Path] + """filenames to load "stuff"/background data""" stuff_classes: List[str] - thing_filenames: List[Path] - thing_classes: List[str] + """class labels for "stuff" data""" stuff_colors: torch.Tensor + """color mapping for "stuff" classes""" + thing_filenames: List[Path] + """filenames to load "thing"/foreground data""" + thing_classes: List[str] + """class labels for "thing" data""" thing_colors: torch.Tensor + """color mapping for "thing" classes""" @dataclass class SceneBounds: - """Data to represent the scene bounds. - - Attributes: - aabb: axis-aligned bounding box. - aabb[0] is the minimum (x,y,z) point. - aabb[1] is the maximum (x,y,z) point. - """ + """Data to represent the scene bounds.""" aabb: TensorType[2, 3] = None + """aabb: axis-aligned bounding box. + aabb[0] is the minimum (x,y,z) point. + aabb[1] is the maximum (x,y,z) point.""" def get_diagonal_length(self): """Returns the longest diagonal length.""" @@ -72,7 +77,11 @@ class SceneBounds: 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.""" + about the origin. + + Args: + scale_factor: How much to scale the camera origins by. + """ return SceneBounds(aabb=(self.aabb - self.get_center()) * scale_factor) @staticmethod @@ -82,9 +91,6 @@ class SceneBounds: Args: positions: the xyz positions aabb: the axis-aligned bounding box - - Returns: - positions that are normalized into the range [0, 1]. """ aabb_lengths = aabb[1] - aabb[0] normalized_positions = (positions - aabb[0]) / aabb_lengths @@ -96,14 +102,23 @@ class SceneBounds: @staticmethod def from_json(json_: Dict) -> "SceneBounds": - """Returns the an instance of SceneBounds from a json dictionary.""" + """Returns the an instance of SceneBounds from a json dictionary. + + Args: + json_: the json dictionary containing scene bound information + """ assert json_["type"] == "aabb" aabb = torch.tensor([json_[0], json_[1]]) return SceneBounds(aabb=aabb) @staticmethod def from_camera_poses(poses: TensorType[..., 3, 4], scale_factor: float) -> "SceneBounds": - """Returns the instance of SceneBounds that fully envelopes a set of poses""" + """Returns the instance of SceneBounds that fully envelopes a set of poses + + Args: + poses: tensor of camera pose matrices + scale_factor: How much to scale the camera origins by. + """ xyzs = poses[..., :3, -1] aabb = torch.stack([torch.min(xyzs, dim=0)[0], torch.max(xyzs, dim=0)[0]]) return SceneBounds(aabb=aabb * scale_factor) @@ -111,22 +126,24 @@ class SceneBounds: @dataclass class DatasetInputs: - """Dataset inputs for the image dataset and the ray generator. - - Args: - image_filenames: Filenames for the images. - intrinsics: Tensor of per-image camera intrisics. - camera_to_world: Tensor of per-image c2w matrices, in [R | t] format. - ... - """ + """Dataset inputs for the image dataset and the ray generator.""" image_filenames: List[Path] + """Filenames for the images.""" cameras: Cameras + """Camera object storing collection of camera information in dataset""" alpha_color: Optional[TensorType[3]] = None + """color of dataset background""" scene_bounds: SceneBounds = SceneBounds() - point_cloud: PointCloud = PointCloud() + """scene bounds of dataset""" # we support additional input information/formats including mask/depth/semantics. additional_inputs: Dict[str, Any] = to_immutable_dict({}) + """Dictionary of additional dataset information (e.g. semantics/point clouds/masks). + {input_name: + ... {"func": function to process additional dataset inputs, + ... "kwargs": dictionary of data to pass into "func"} + } + """ def as_dict(self) -> dict: """Returns the dataclass as a dictionary.""" diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index e8119ae16..1da1fe142 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -19,6 +19,7 @@ from __future__ import annotations import functools import logging +import os from pathlib import Path from typing import Any, Dict, List @@ -45,12 +46,9 @@ def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any: """Main training function that sets up and runs the trainer per process Args: - local_rank (int): current rank of process - world_size (int): total number of gpus available - config (Config): config file specifying training regimen - - Returns: - Any: TODO(): determine the return type + local_rank: current rank of process + world_size: total number of gpus available + config: config file specifying training regimen """ trainer = Trainer(config, local_rank, world_size) trainer.setup() @@ -62,9 +60,9 @@ class Trainer: """Trainer class Args: - config (Config): The configuration object. - local_rank (int, optional): Local rank of the process. Defaults to 0. - world_size (int, optional): World size of the process. Defaults to 1. + config: The configuration object. + local_rank: Local rank of the process. Defaults to 0. + world_size: World size of the process. Defaults to 1. """ def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1): @@ -98,7 +96,7 @@ class Trainer: """Setup the Trainer by calling other setup functions. Args: - test_mode (bool, optional): Whether to setup for testing. Defaults to False. + test_mode: Whether to setup for testing. Defaults to False. """ self.pipeline: Pipeline = self.config.pipeline.setup( device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank @@ -117,7 +115,7 @@ class Trainer: assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs" self.visualizer_state.init_scene( - dataset=self.pipeline.datamanager.train_input_dataset, start_train=self.config.viewer.train + dataset=self.pipeline.datamanager.train_input_dataset, start_train=self.config.viewer.start_train ) with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations @@ -172,7 +170,7 @@ class Trainer: """Perform writes only during appropriate time steps Args: - step (int): Current training step. + step: Current training step. """ if ( step % self.config.logging.steps_per_log == 0 @@ -185,8 +183,12 @@ class Trainer: def _load_checkpoint(self) -> None: """Helper function to load pipeline and optimizer from prespecified checkpoint""" load_dir = self.config.trainer.load_dir - load_step = self.config.trainer.load_step - if load_dir is not None and load_step is not None: + if load_dir is not None: + load_step = self.config.trainer.load_step + if load_step is None: + print("Loading latest checkpoint from load_dir") + # NOTE: this is specific to the checkpoint name format + load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1] load_path = load_dir / f"step-{load_step:09d}.ckpt" assert load_path.exists(), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") @@ -199,6 +201,8 @@ class Trainer: else: logging.info("No checkpoints to load, training from scratch") + logging.info("saving checkpoints to: %s", self.config.trainer.model_dir) + @check_main_thread def _save_checkpoint(self, output_dir: Path, step: int) -> None: """Save the model and optimizers @@ -226,13 +230,10 @@ class Trainer: @profiler.time_function def train_iteration(self, step: int) -> Dict[str, torch.Tensor]: - """Run one iteration with a batch of inputs. + """Run one iteration with a batch of inputs. Returns dictionary of model losses. Args: step: Current training step. - - Returns: - Dictionary of model losses. """ self.optimizers.zero_grad_all() cpu_or_cuda_str = self.device.split(":")[0] diff --git a/nerfactory/fields/base.py b/nerfactory/fields/base.py index c4bae0ad2..af8778d76 100644 --- a/nerfactory/fields/base.py +++ b/nerfactory/fields/base.py @@ -30,8 +30,12 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames class Field(nn.Module): """Base class for fields.""" - def density_fn(self, positions): - """Returns only the density. Used primarily with the density grid.""" + def density_fn(self, positions: TensorType["bs":..., 3]): + """Returns only the density. Used primarily with the density grid. + + Args: + positions: the origin of the samples/frustums + """ # Need to figure out a better way to descibe positions with a ray. ray_samples = RaySamples( frustums=Frustums( @@ -47,34 +51,28 @@ class Field(nn.Module): @abstractmethod def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: - """Computes and returns the densities. + """Computes and returns the densities. Returns a tensor of densities and a tensor of features. Args: - ray_samples (RaySamples): Samples locations to compute density. - - Returns: - Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features. + ray_samples: Samples locations to compute density. """ @abstractmethod def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: - """Computes and returns the colors. + """Computes and returns the colors. Returns output field values. Args: - ray_samples (RaySamples): Samples locations to compute outputs. - density_embedding (TensorType, optional): Density embeddings to condition on. - - Returns: - Dict[FieldHeadNames, TensorType]: Output field values. + ray_samples: Samples locations to compute outputs. + density_embedding: Density embeddings to condition on. """ def forward(self, ray_samples: RaySamples): """Evaluates the field at points along the ray. Args: - ray_samples (RaySamples): Samples to evaluate field on. + ray_samples: Samples to evaluate field on. """ valid_mask = ray_samples.valid_mask diff --git a/nerfactory/fields/density_fields/density_grid.py b/nerfactory/fields/density_fields/density_grid.py index 3a4ee3a8e..6f8900b7c 100644 --- a/nerfactory/fields/density_fields/density_grid.py +++ b/nerfactory/fields/density_fields/density_grid.py @@ -29,14 +29,11 @@ from nerfactory.cameras.rays import RaySamples def _create_grid_coords(resolution: int, device: Union[torch.device, str] = "cpu") -> TensorType["n_coords", 3]: - """Create 3D grid coordinates + """Create 3D grid coordinates. Returns all grid coordinates with shape [res * res * res, 3]. Args: - resolution (int): The 3D resolution of the grid. - device (torch.device): Which device you want the returned tensor to live on. - - Returns: - TensorType["n_coords", 3]: All grid coordinates with shape [res * res * res, 3]. + resolution: The 3D resolution of the grid. + device: Which device you want the returned tensor to live on. """ arrange = torch.arange(resolution, device=device) grid_x, grid_y, grid_z = torch.meshgrid([arrange, arrange, arrange], indexing="ij") @@ -73,11 +70,11 @@ class DensityGrid(nn.Module): TODO(ruilongli): support `center` and `base_scale` with 3-dim. Args: - center (float): Center of all the grids in the world space. Defaults to 0. - base_scale (float): Center of the scale of the base level grid. Defaults to 1. - num_cascades (int): Number of the cascaded multi-res levels. Defaults to 1. - resolution (int): Resolution of the grid. Defaults to 128. - update_every_num_iters (int): How frequently to update the grid values. Defaults to 16. + center: Center of all the grids in the world space. Defaults to 0. + base_scale: Center of the scale of the base level grid. Defaults to 1. + num_cascades: Number of the cascaded multi-res levels. Defaults to 1. + resolution: Resolution of the grid. Defaults to 128. + update_every_num_iters: How frequently to update the grid values. Defaults to 16. """ density_grid: TensorType["num_cascades", "resolution_cubed"] @@ -111,24 +108,16 @@ class DensityGrid(nn.Module): @torch.no_grad() def get_all_cells(self) -> List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]: - """Returns all cells of the grid. - - Returns: - List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]: All grid indices and grid - coordinates of all cascade levels. - """ + """Returns a list of all grid indices and grid coordinates of all cascade levels.""" return [(self.grid_indices, self.grid_coords)] * self.num_cascades @torch.no_grad() def sample_uniform_and_occupied_cells(self, n: int) -> List[Tuple[TensorType["n"], TensorType["n", 3]]]: - """Samples both n uniform and occupied cells (per cascade) + """Samples both n uniform and occupied cells (per cascade). + Returns a tuple of (sampled grid indices, grid coordinates of all cascade levels). Args: n: Number of cells to be sampled. - - Returns: - List[Tuple[TensorType["n"], TensorType["n", 3]]]: Sampled grid indices and grid - coordinates of all cascade levels. """ device = self.density_grid.device @@ -164,11 +153,11 @@ class DensityGrid(nn.Module): returns densities (N, 1). step: The current training step. Instant-NGP has a warmup stage where all cells are updated. After certain amount of steps (256), it speeds up by only update sampled cells. - density_threshold (float): The threshold to prune cells. Recommended to calculate it using this rule: + density_threshold: The threshold to prune cells. Recommended to calculate it using this rule: `weight_threshold / dt`. For example for Instant-NGP on Lego, the minimum step size `dt` is `3` (scale of the scene) * `sqrt(3)` / `1024` (number of samples). And if we want to the weight threshold to be `0.01`, we will get `0.01 / (3 * sqrt(3) / 1024) ~= 2.` for density threshold. - decay (float): EMA decay for density updating. + decay: EMA decay for density updating. """ # create temporary grid tmp_grid = -torch.ones_like(self.density_grid) @@ -206,10 +195,9 @@ class DensityGrid(nn.Module): @torch.no_grad() def get_densities(self, ray_samples: RaySamples) -> TensorType[..., 1]: """Get the density values stored in the density grid. + Args: - ray_samples RaySamples: ray samples - Returns: - TensorType[..., 1]: Density values + ray_samples: ray samples to get density for """ positions = ray_samples.frustums.get_positions() deltas = ray_samples.frustums.ends - ray_samples.frustums.starts diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 808edbad0..6f79c928c 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -38,16 +38,35 @@ except ImportError: pass -def get_normalized_directions(directions): - """SH encoding must be in the range [0, 1]""" +def get_normalized_directions(directions: TensorType["bs":..., 3]): + """SH encoding must be in the range [0, 1] + + Args: + directions: batch of directions + """ return (directions + 1.0) / 2.0 class TCNNInstantNGPField(Field): - """NeRF Field""" + """NeRF Field + + Args: + aabb: parameters of scene aabb bounds + num_layers: number of hidden layers + hidden_dim: dimension of hidden layers + geo_feat_dim: output geo feat dimensions + num_layers_color: number of hidden layers for color network + hidden_dim_color: dimension of hidden layers for color network + """ def __init__( - self, aabb, num_layers=2, hidden_dim=64, geo_feat_dim=15, num_layers_color=3, hidden_dim_color=64 + self, + aabb, + num_layers: int = 2, + hidden_dim: int = 64, + geo_feat_dim: int = 15, + num_layers_color: int = 3, + hidden_dim_color: int = 64, ) -> None: super().__init__() @@ -99,7 +118,6 @@ class TCNNInstantNGPField(Field): ) def get_density(self, ray_samples: RaySamples): - """Computes and returns the densities.""" positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) diff --git a/nerfactory/fields/modules/base.py b/nerfactory/fields/modules/base.py index 9f2e6295a..d9fc655bd 100644 --- a/nerfactory/fields/modules/base.py +++ b/nerfactory/fields/modules/base.py @@ -26,8 +26,8 @@ class FieldModule(nn.Module): """Field modules that can be combined to store and compute the fields. Args: - in_dim (int, optional): Input dimension to module. - out_dim (int, optional): Ouput dimension to module. + in_dim: Input dimension to module. + out_dim: Ouput dimension to module. """ def __init__(self, in_dim: Optional[int] = None, out_dim: Optional[int] = None) -> None: @@ -43,18 +43,14 @@ class FieldModule(nn.Module): """Sets input dimension of encoding Args: - in_dim (int): input dimension + in_dim: input dimension """ if in_dim <= 0: raise ValueError("Input dimension should be greater than zero") self.in_dim = in_dim def get_out_dim(self) -> int: - """Calculates output dimension of encoding. - - Returns: - int: output dimension - """ + """Calculates output dimension of encoding.""" if self.out_dim is None: raise ValueError("Output dimension has not been set") return self.out_dim @@ -62,10 +58,9 @@ class FieldModule(nn.Module): @abstractmethod def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: """ + Returns processed tensor + Args: in_tensor: Input tensor to process - - Returns: - Processed tensor """ raise NotImplementedError diff --git a/nerfactory/fields/modules/embedding.py b/nerfactory/fields/modules/embedding.py index 354e49cd5..f3434718f 100644 --- a/nerfactory/fields/modules/embedding.py +++ b/nerfactory/fields/modules/embedding.py @@ -28,8 +28,8 @@ class Embedding(FieldModule): # TODO: add different types of initializations Args: - in_dim (int): Number of embeddings - out_dim (int): Dimension of the embedding vectors + in_dim: Number of embeddings + out_dim: Dimension of the embedding vectors """ def __init__(self, in_dim: int, out_dim: int) -> None: @@ -42,5 +42,9 @@ class Embedding(FieldModule): self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim) def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """Call forward""" + """Call forward + + Args: + in_tensor: input tensor to process + """ return self.embedding(in_tensor) diff --git a/nerfactory/fields/modules/encoding.py b/nerfactory/fields/modules/encoding.py index 6aad364b5..4dd5be992 100644 --- a/nerfactory/fields/modules/encoding.py +++ b/nerfactory/fields/modules/encoding.py @@ -43,6 +43,11 @@ class Encoding(FieldModule): @abstractmethod def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + """Call forward and returns and processed tensor + + Args: + in_tensor: the input tensor to process + """ raise NotImplementedError @@ -243,13 +248,10 @@ class HashEncoding(Encoding): return self.num_levels * self.features_per_level def hash_fn(self, in_tensor: TensorType["bs":..., "num_levels", 3]) -> TensorType["bs":..., "num_levels"]: - """Hash tensor using method described in Instant-NGP + """Returns hash tensor using method described in Instant-NGP Args: in_tensor: Tensor to be hashed - - Returns: - Hashed tensor """ # min_val = torch.min(in_tensor) diff --git a/nerfactory/fields/nerf_field.py b/nerfactory/fields/nerf_field.py index 499e02808..ec0f000c9 100644 --- a/nerfactory/fields/nerf_field.py +++ b/nerfactory/fields/nerf_field.py @@ -38,15 +38,15 @@ class NeRFField(Field): """NeRF Field Args: - position_encoding (Encoding, optional): Position encoder. Defaults to Identity(in_dim=3). - direction_encoding (Encoding, optional): Direction encoder. Defaults to Identity(in_dim=3). - base_mlp_num_layers (int, optional): Number of layers for base MLP. Defaults to 8. - base_mlp_layer_width (int, optional): Width of base MLP layers. Defaults to 256. - head_mlp_num_layers (int, optional): Number of layer for ourput head MLP. Defaults to 2. - head_mlp_layer_width (int, optional): Width of output head MLP layers. Defaults to 128. - skip_connections (Tuple, optional): Where to add skip connection in base MLP. Defaults to (4,). - use_integrated_encoding (bool, optional): Used integrated samples as encoding input, Defaults to False. - spatial_distortion (SpatialDistortion, optional): Spatial distortion. Defaults to None. + position_encoding: Position encoder. Defaults to Identity(in_dim=3). + direction_encoding: Direction encoder. Defaults to Identity(in_dim=3). + base_mlp_num_layers: Number of layers for base MLP. Defaults to 8. + base_mlp_layer_width: Width of base MLP layers. Defaults to 256. + head_mlp_num_layers: Number of layer for ourput head MLP. Defaults to 2. + head_mlp_layer_width: Width of output head MLP layers. Defaults to 128. + skip_connections: Where to add skip connection in base MLP. Defaults to (4,). + use_integrated_encoding: Used integrated samples as encoding input, Defaults to False. + spatial_distortion: Spatial distortion. Defaults to None. """ def __init__( diff --git a/nerfactory/fields/nerfw_field.py b/nerfactory/fields/nerfw_field.py index b5e6fecf4..481977d4e 100644 --- a/nerfactory/fields/nerfw_field.py +++ b/nerfactory/fields/nerfw_field.py @@ -39,16 +39,16 @@ class VanillaNerfWField(Field): """The NeRF-W field which has appearance and transient conditioning. Args: - num_images (int): How many images exist in the dataset. - position_encoding (Encoding, optional): Position encoder. Defaults to Identity(in_dim=3). - direction_encoding (Encoding, optional): Direction encoder. Defaults to Identity(in_dim=3). - base_mlp_num_layers (int, optional): Number of layers for base MLP. Defaults to 8. - base_mlp_layer_width (int, optional): Width of base MLP layers. Defaults to 256. - head_mlp_num_layers (int, optional): Number of layer for ourput head MLP. Defaults to 2. - head_mlp_layer_width (int, optional): Width of output head MLP layers. Defaults to 128. - appearance_embedding_dim: (int, optional): Dimension of appearance embedding. Defaults to 48. - transient_embedding_dim: (int, optional): Dimension of transient embedding. Defaults to 16. - skip_connections (Tuple, optional): Where to add skip connection in base MLP. Defaults to (4,). + num_images: How many images exist in the dataset. + position_encoding: Position encoder. Defaults to Identity(in_dim=3). + direction_encoding: Direction encoder. Defaults to Identity(in_dim=3). + base_mlp_num_layers: Number of layers for base MLP. Defaults to 8. + base_mlp_layer_width: Width of base MLP layers. Defaults to 256. + head_mlp_num_layers: Number of layer for ourput head MLP. Defaults to 2. + head_mlp_layer_width: Width of output head MLP layers. Defaults to 128. + appearance_embedding_dim:: Dimension of appearance embedding. Defaults to 48. + transient_embedding_dim:: Dimension of transient embedding. Defaults to 16. + skip_connections: Where to add skip connection in base MLP. Defaults to (4,). """ def __init__( diff --git a/nerfactory/fields/semantic_nerf_field.py b/nerfactory/fields/semantic_nerf_field.py index e0f6f5342..99edf6e15 100644 --- a/nerfactory/fields/semantic_nerf_field.py +++ b/nerfactory/fields/semantic_nerf_field.py @@ -34,7 +34,18 @@ from nerfactory.fields.modules.mlp import MLP class SemanticNerfField(Field): - """Semantic-NeRF field""" + """Semantic-NeRF field + + Args: + num_semantic_classes: Number of distinct semantic classes. + position_encoding: Position encoder. Defaults to Identity(in_dim=3). + direction_encoding: Direction encoder. Defaults to Identity(in_dim=3). + base_mlp_num_layers: Number of layers for base MLP. Defaults to 8. + base_mlp_layer_width: Width of base MLP layers. Defaults to 256. + head_mlp_num_layers: Number of layer for ourput head MLP. Defaults to 2. + head_mlp_layer_width: Width of output head MLP layers. Defaults to 128. + skip_connections: Where to add skip connection in base MLP. Defaults to (4,). + """ def __init__( self, diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index a95eeea3e..3e9b65a21 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -40,15 +40,9 @@ class Model(nn.Module): Where everything (Fields, Optimizers, Samplers, Visualization, etc) is linked together. This should be subclassed for custom NeRF model. - TODO: - Args: - scene_bounds: Bounds of target scene. - loss_coefficients: Loss specific weights. - enable_collider: Whether to create a scene collider to filter rays. - collider_config: Configuration of scene collider. - enable_density_field: Whether to create a density field to filter samples. - density_field_config: Configuration of density field. + config: configuration for instantiating model + scene_bounds: dataset scene bounds """ config: cfg.ModelConfig @@ -134,14 +128,13 @@ class Model(nn.Module): ) -> Union[ Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]] ]: - """Run forward starting with a ray bundle. + """Run forward starting with a ray bundle. This outputs different things depending on the configuration + of the model and whether or not the batch is provided (whether or not we are training basically) - This takes in raybundles (containing all the information needed to render that ray... - latents included) and the batch (containing all the auxilury things needed to train - like masks and ground truth pixels). - - This outputs different things depending on the configuration of the model and whether or not - the batch is provided (whether or not we are training basically).""" + Args: + ray_bundle: containing all the information needed to render that ray latents included + batch: containing all the auxilury things needed to train like masks and ground truth pixels. + """ if self.collider is not None: intersected_ray_bundle = self.collider(ray_bundle) # pylint: disable=not-callable @@ -169,18 +162,34 @@ class Model(nn.Module): return outputs, loss_dict, metrics_dict def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: - """Compute and returns metrics.""" + """Compute and returns metrics. + + Args: + outputs: the output to compute loss dict to + batch: ground truth batch corresponding to outputs + """ # pylint: disable=unused-argument # pylint: disable=no-self-use return {} @abstractmethod def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: - """Computes and returns the losses dict.""" + """Computes and returns the losses dict. + + Args: + outputs: the output to compute loss dict to + batch: ground truth batch corresponding to outputs + metrics_dict: collection of metrics to compute + loss_coefficients: list of loss coefficients/weightings to apply + """ @torch.no_grad() def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]: - """Takes in camera parameters and computes the output of the model.""" + """Takes in camera parameters and computes the output of the model. + + Args: + camera_ray_bundle: ray bundle to calculate outputs over + """ assert is_not_none(camera_ray_bundle.num_rays_per_chunk) image_height, image_width = camera_ray_bundle.origins.shape[:2] num_rays = len(camera_ray_bundle) @@ -213,6 +222,10 @@ class Model(nn.Module): """ def load_model(self, loaded_state: Dict[str, Any]) -> None: - """Load the checkpoint from the given path""" + """Load the checkpoint from the given path + + Args: + loaded_state: dictionary of pre-trained model states + """ state = {key.replace("module.", ""): value for key, value in loaded_state["model"].items()} self.load_state_dict(state) # type: ignore diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 98f42baf4..eeffd5ac2 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -46,8 +46,7 @@ class NGPModel(Model): """Instant NGP model Args: - field_implementation (str): one of "torch" or "tcnn", or other fields in 'field_implementation_to_class' - kwargs: additional params to pass up to the parent class model + config: instant NGP configuration to instantiate model """ config: cfg.InstantNGPModelConfig diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index bd7e4b428..514673c46 100644 --- a/nerfactory/models/mipnerf.py +++ b/nerfactory/models/mipnerf.py @@ -42,7 +42,11 @@ from nerfactory.utils import colors, misc, visualization, writer class MipNerfModel(Model): - """mip-NeRF model""" + """mip-NeRF model + + Args: + config: MipNerf configuration to instantiate model + """ def __init__( self, diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index 1a138a85a..19dd35452 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/nerfactory/models/mipnerf_360.py @@ -45,7 +45,11 @@ from nerfactory.utils import colors, misc, visualization, writer class MipNerf360Model(Model): - """mip-NeRF model""" + """mip-NeRF model + + Args: + config: MipNerf360 configuration to instantiate model + """ def __init__( self, diff --git a/nerfactory/models/modules/ray_generator.py b/nerfactory/models/modules/ray_generator.py index 8b521648a..c30f40f00 100644 --- a/nerfactory/models/modules/ray_generator.py +++ b/nerfactory/models/modules/ray_generator.py @@ -27,8 +27,7 @@ class RayGenerator(nn.Module): This class will store the intrinsics and extrinsics parameters of the cameras. Args: - intrinsics: The intrinsics parameters. - camera_to_world: Camera to world transformation matrix. + cameras: Camera objects containing camera info """ def __init__(self, cameras: Cameras) -> None: diff --git a/nerfactory/models/modules/ray_losses.py b/nerfactory/models/modules/ray_losses.py index 69ece5132..9724f6112 100644 --- a/nerfactory/models/modules/ray_losses.py +++ b/nerfactory/models/modules/ray_losses.py @@ -23,7 +23,7 @@ from nerfactory.cameras.rays import RaySamples def distortion_loss( ray_samples: RaySamples, densities: TensorType["bs":..., "num_samples", 1] ) -> TensorType["bs":..., 1]: - """Ray baserd distortion loss proposed in MipNeRF-360. + """Ray baserd distortion loss proposed in MipNeRF-360. Returns distortion Loss. .. math:: @@ -36,9 +36,6 @@ def distortion_loss( Args: ray_samples: Ray samples to compute loss over densities: Predicted sample densities - - Returns: - Distortion Loss. """ # Compute the weight at each sample location diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index e4cf53dc9..5ca58c83a 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -29,7 +29,12 @@ from nerfactory.fields.density_fields.density_grid import DensityGrid class Sampler(nn.Module): - """Generate Samples""" + """Generate Samples + + Args: + num_samples: number of samples to take + density_field: density grid specifying weighting of + """ def __init__( self, num_samples: int, density_field: Optional[DensityGrid] = None, weight_threshold: float = 0.01 @@ -56,7 +61,16 @@ class Sampler(nn.Module): class SpacedSampler(Sampler): - """Sample points according to a function.""" + """Sample points according to a function. + + Args: + num_samples: Number of samples per ray + spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform). + spacing_fn_inv: The inverse of spacing_fn. + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ def __init__( self, @@ -67,15 +81,6 @@ class SpacedSampler(Sampler): density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-2, ) -> None: - """ - Args: - num_samples: Number of samples per ray - spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform). - spacing_fn_inv: The inverse of spacing_fn. - train_stratified: Use stratified sampling during training. Defults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - """ super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified self.spacing_fn = spacing_fn @@ -121,7 +126,14 @@ class SpacedSampler(Sampler): class UniformSampler(SpacedSampler): - """Sample uniformly along a ray""" + """Sample uniformly along a ray + + Args: + num_samples: Number of samples per ray + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ def __init__( self, @@ -130,13 +142,6 @@ class UniformSampler(SpacedSampler): density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-2, ) -> None: - """ - Args: - num_samples: Number of samples per ray - train_stratified: Use stratified sampling during training. Defults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - """ super().__init__( num_samples=num_samples, spacing_fn=lambda x: x, @@ -148,7 +153,14 @@ class UniformSampler(SpacedSampler): class LinearDisparitySampler(SpacedSampler): - """Sample linearly in disparity along a ray""" + """Sample linearly in disparity along a ray + + Args: + num_samples: Number of samples per ray + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ def __init__( self, @@ -157,13 +169,6 @@ class LinearDisparitySampler(SpacedSampler): density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-2, ) -> None: - """ - Args: - num_samples: Number of samples per ray - train_stratified: Use stratified sampling during training. Defults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - """ super().__init__( num_samples=num_samples, spacing_fn=lambda x: 1 / x, @@ -175,7 +180,14 @@ class LinearDisparitySampler(SpacedSampler): class SqrtSampler(SpacedSampler): - """Square root sampler along a ray""" + """Square root sampler along a ray + + Args: + num_samples: Number of samples per ray + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ def __init__( self, @@ -184,13 +196,6 @@ class SqrtSampler(SpacedSampler): density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-2, ) -> None: - """ - Args: - num_samples: Number of samples per ray - train_stratified: Use stratified sampling during training. Defults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - """ super().__init__( num_samples=num_samples, spacing_fn=torch.sqrt, @@ -202,7 +207,14 @@ class SqrtSampler(SpacedSampler): class LogSampler(SpacedSampler): - """Log sampler along a ray""" + """Log sampler along a ray + + Args: + num_samples: Number of samples per ray + train_stratified: Use stratified sampling during training. Defults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + """ def __init__( self, @@ -211,13 +223,6 @@ class LogSampler(SpacedSampler): density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-2, ) -> None: - """ - Args: - num_samples: Number of samples per ray - train_stratified: Use stratified sampling during training. Defults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - """ super().__init__( num_samples=num_samples, spacing_fn=torch.log, @@ -229,7 +234,16 @@ class LogSampler(SpacedSampler): class PDFSampler(Sampler): - """Sample based on probability distribution""" + """Sample based on probability distribution + + Args: + num_samples: Number of samples per ray + train_stratified: Randomize location within each bin during training. Defaults to True + include_original: Add original samples to ray. Defaults to True + density_field: Density grid. If provides, samples below weight_threshold as set as invalid. + weight_threshold: Removes samples below threshold weight. Only used if density field is provided. + histogram_padding: Amount to weights prior to computing PDF. Defaults to 0.01. + """ def __init__( self, @@ -240,15 +254,6 @@ class PDFSampler(Sampler): weight_threshold: float = 1e-2, histogram_padding: float = 0.01, ) -> None: - """ - Args: - num_samples: Number of samples per ray - train_stratified: Randomize location within each bin during training. Defaults to True - include_original: Add original samples to ray. Defaults to True - density_field: Density grid. If provides, samples below weight_threshold as set as invalid. - weight_threshold: Removes samples below threshold weight. Only used if density field is provided. - histogram_padding: Amount to weights prior to computing PDF. Defaults to 0.01. - """ super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified self.include_original = include_original diff --git a/nerfactory/models/modules/scene_colliders.py b/nerfactory/models/modules/scene_colliders.py index 8c316f470..5a18a6db6 100644 --- a/nerfactory/models/modules/scene_colliders.py +++ b/nerfactory/models/modules/scene_colliders.py @@ -39,7 +39,11 @@ class SceneBoundsCollider(nn.Module): class AABBBoxCollider(SceneBoundsCollider): - """Module for colliding rays with the scene bounds to compute near and far values.""" + """Module for colliding rays with the scene bounds to compute near and far values. + + Args: + scene_bounds: scene bounds to apply to dataset + """ def __init__(self, scene_bounds: SceneBounds, **kwargs) -> None: super().__init__(**kwargs) @@ -49,12 +53,13 @@ class AABBBoxCollider(SceneBoundsCollider): def intersect_with_aabb( cls, rays_o: TensorType["num_rays", 3], rays_d: TensorType["num_rays", 3], aabb: TensorType[2, 3] ): - """_summary_ + """Returns collection of valid rays within a specified near/far bounding box along with a mask + specifying which rays are valid Args: - rays_o (torch.tensor): (num_rays, 3) - rays_d (torch.tensor): (num_rays, 3) - aabb (torch.tensor): (2, 3) This is [min point (x,y,z), max point (x,y,z)] + rays_o: (num_rays, 3) ray origins + rays_d: (num_rays, 3) ray directions + aabb: (2, 3) This is [min point (x,y,z), max point (x,y,z)] """ # avoid divide by zero dir_fraction = 1.0 / (rays_d + 1e-6) @@ -87,6 +92,9 @@ class AABBBoxCollider(SceneBoundsCollider): def forward(self, ray_bundle: RayBundle) -> RayBundle: """Intersects the rays with the scene bounds and updates the near and far values. Populates nears and fars fields and returns the ray_bundle. + + Args: + ray_bundle: specified ray bundle to operate on """ aabb = self.scene_bounds.aabb nears, fars, valid_mask = self.intersect_with_aabb(ray_bundle.origins, ray_bundle.directions, aabb) @@ -97,7 +105,12 @@ class AABBBoxCollider(SceneBoundsCollider): class NearFarCollider(SceneBoundsCollider): - """Sets the nears and fars with fixed values.""" + """Sets the nears and fars with fixed values. + + Args: + near_plane: distance to near plane + far_plane: distance to far plane + """ def __init__(self, near_plane: float, far_plane: float, **kwargs) -> None: self.near_plane = near_plane diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index 3b3d6f5cc..176e2b001 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -43,7 +43,11 @@ from nerfactory.utils import colors, misc, visualization, writer class NerfWModel(Model): - """NeRF-W model""" + """NeRF-W model + + Args: + config: NerfW configuration to instantiate model + """ config: cfg.NerfWModelConfig diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index b175d8711..b43d06d30 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/nerfactory/models/semantic_nerf.py @@ -34,7 +34,12 @@ from nerfactory.utils import misc, writer class SemanticNerfModel(NeRFModel): - """Semantic-NeRF model""" + """Semantic-NeRF model + + Args: + config: SemanticNeRF configuration to instantiate model + semantics: additional semantics data info + """ def __init__(self, config: cfg.ModelConfig, semantics: Semantics, **kwargs) -> None: self.stuff_classes = semantics.stuff_classes diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index b054c16c2..5817200cf 100644 --- a/nerfactory/models/tensorf.py +++ b/nerfactory/models/tensorf.py @@ -50,8 +50,10 @@ from nerfactory.utils.callbacks import ( class TensoRFModel(Model): - """ - TensoRF Model + """TensoRF Model + + Args: + config: TensoRF configuration to instantiate model """ def __init__( diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index a0abb5b2f..be2cdd431 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -51,12 +51,7 @@ class NeRFModel(Model): """Vanilla NeRF model Args: - near_plane: Where to start sampling points. Defaults to 2.0. - far_plane: Where to stop sampling points. Defaults to 6.0. - num_coarse_samples: Number of samples in coarse field evaluation. Defaults to 64, - num_importance_samples: Number of samples in fine field evaluation. Defaults to 64, - enable_density_field: Whether to create a density field to filter samples. Defaults to False. - density_field_config: Configuration of density field. Defaults to None. + config: Basic NeRF configuration to instantiate model """ def __init__( diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 9e8099db9..85790cdcd 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -66,13 +66,15 @@ class Pipeline(nn.Module): Args: - model: The model to be used in the pipeline. - datamanager: The datamanager to be used in the pipeline. - loss_coefficients: A dictionary of loss coefficients that will be used + config: configuration to instantiate pipeline + device: location to place model and data + test_mode: if True, loads test datset. if False, loads train/eval datasets + world_size: total number of machines available + local_rank: rank of current machine Attributes: - self.datamanager (DataManager): The data manager that will be used - self.model (Model): The model that will be used + self.datamanager: The data manager that will be used + self.model: The model that will be used """ def __init__( @@ -107,7 +109,11 @@ class Pipeline(nn.Module): def get_train_loss_dict(self, step: Optional[int] = None): """This function gets your training loss dict. This will be responsible for getting the next batch of data from the DataManager and interfacing with the - Model class, feeding the data to the model's forward function.""" + Model class, feeding the data to the model's forward function. + + Args: + step: current iteration step to update sampler if using DDP (distributed) + """ if self.world_size > 1: self.datamanager.sampler.set_epoch(step) ray_bundle, batch = self.datamanager.next_train() @@ -118,7 +124,11 @@ class Pipeline(nn.Module): @profiler.time_function def get_eval_loss_dict(self, step: Optional[int] = None): """This function gets your evaluation loss dict. It needs to get the data - from the DataManager and feed it to the model's forward function""" + from the DataManager and feed it to the model's forward function + + Args: + step: current iteration step + """ self.eval() # NOTE(ethan): next_eval() is not being used right now assert self.datamanager.eval_dataloader is not None @@ -147,7 +157,11 @@ class Pipeline(nn.Module): """Log the test image outputs""" def load_pipeline(self, loaded_state: Dict[str, Any]) -> None: - """Load the checkpoint from the given path""" + """Load the checkpoint from the given path + + Args: + loaded_state: pre-trained model state dict + """ state = {key.replace("module.", ""): value for key, value in loaded_state.items()} self.load_state_dict(state) # type: ignore diff --git a/nerfactory/utils/callbacks.py b/nerfactory/utils/callbacks.py index 874c6aa17..165a7ed47 100644 --- a/nerfactory/utils/callbacks.py +++ b/nerfactory/utils/callbacks.py @@ -31,8 +31,11 @@ class TrainingCallbackAttributes: # TODO(ethan): type this without circular imports optimizers: Optional[InitVar] + """optimizers for training""" grad_scaler: Optional[InitVar] + """gradient scalers""" pipeline: Optional[InitVar] + """reference to training pipeline""" class TrainingCallbackLocation(Enum): @@ -48,8 +51,10 @@ class TrainingCallback: including at iteration 0. The function is called after the training iteration. Args: - update_every_num_iters: How often to call the function `func`. + where_to_run: List of locations for when to run callbak (before/after iteration) func: The function that will be called. + update_every_num_iters: How often to call the function `func`. + iters: Tuple of iteration steps to perform callback args: args for the function 'func'. kwargs: kwargs for the function 'func'. """ @@ -77,7 +82,7 @@ class TrainingCallback: """Callback to run after training step Args: - step (int): current iteration step + step: current iteration step """ if self.update_every_num_iters is not None: if step % self.update_every_num_iters == 0: @@ -87,6 +92,11 @@ class TrainingCallback: self.func(*self.args, **self.kwargs, step=step) def run_callback_at_location(self, step: int, location: TrainingCallbackLocation): - """Runs the callback if it's supposed to be run at the given location.""" + """Runs the callback if it's supposed to be run at the given location. + + Args: + step: current iteration step + location: when to run callback (before/after iteration) + """ if location in self.where_to_run: self.run_callback(step=step) diff --git a/nerfactory/utils/colmap_utils.py b/nerfactory/utils/colmap_utils.py index 3f567cbb1..e58e1448b 100644 --- a/nerfactory/utils/colmap_utils.py +++ b/nerfactory/utils/colmap_utils.py @@ -47,80 +47,68 @@ import numpy as np @dataclass class CameraModel: - """Camera model. - - Attributes: - model_id: Model identifier. - model_name: Model name. - num_params: Number of parameters. - """ + """Camera model""" model_id: int + """Model identifier""" model_name: str + """Model name""" num_params: int + """Number of parameters""" @dataclass class Camera: - """Camera - - Attributes: - camera_id: Camera identifier. - model: Camera model. - width: Image width. - height: Image height. - params: Camera parameters. - """ + """Camera""" id: int + """Camera identifier""" model: str + """Camera model""" width: int + """Image width""" height: int + """Image height""" params: np.ndarray + """Camera parameters""" @dataclass class Image: - """Data the corresponds to a single image. - - Attributes: - id: Image identifier. - qvec: Quaternion vector. - tvec: Translation vector. - camera_id: Camera identifier. - name: Image name. - xys: 2D points. - point3D_ids: Point3D identifiers. - """ + """Data the corresponds to a single image""" id: int + """Image identifier""" qvec: np.ndarray + """Quaternion vector""" tvec: np.ndarray + """Translation vector""" camera_id: int + """Camera identifier""" name: str + """Image name""" xys: np.ndarray + """2D points""" point3d_ids: np.ndarray + """Point3D identifiers""" @dataclass class Point3D: - """Data that corresponds to a single 3D point. - - Attributes: - id: Point3D identifier. - xyz: 3D point. - rgb: Color. - error: Reconstruction error. - image_ids: Image identifiers. - point2d_idxs: Point2D indices. - """ + """Data that corresponds to a single 3D point""" id: int + """Point3D identifier""" xyz: np.ndarray + """3D point""" rgb: np.ndarray + """Color""" error: float + """Reconstruction error""" image_ids: np.ndarray + """Image identifiers""" point2d_idxs: np.ndarray + """Point2D indices""" CAMERA_MODELS = [ diff --git a/nerfactory/utils/decorators.py b/nerfactory/utils/decorators.py index 99adfa79d..4558b51bd 100644 --- a/nerfactory/utils/decorators.py +++ b/nerfactory/utils/decorators.py @@ -21,7 +21,11 @@ from nerfactory.utils import comms def decorate_all(decorators: List[Callable]) -> Callable: - """A decorator to decorate all member functions of a class""" + """A decorator to decorate all member functions of a class + + Args: + decorators: list of decorators to add to all functions in the class + """ def decorate(cls): for attr in cls.__dict__: diff --git a/nerfactory/utils/math.py b/nerfactory/utils/math.py index 68126532c..68d163cf7 100644 --- a/nerfactory/utils/math.py +++ b/nerfactory/utils/math.py @@ -22,12 +22,11 @@ from torchtyping import TensorType def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]: """ - Args: - levels (int): Number of spherical harmonic levels to compute. - directions (TensorType[..., 3]): Spherical hamonic coefficients + Returns value for each component of spherical harmonics. - Returns: - TensorType[..., "components"]: Value for each component. + Args: + levels: Number of spherical harmonic levels to compute. + directions: Spherical hamonic coefficients """ num_components = levels**2 components = torch.zeros((*directions.shape[:-1], num_components)).to(directions.device) @@ -90,8 +89,8 @@ class Gaussians: """Stores Gaussians Args: - mean (TensorType[..., "dim"]): Mean of multivariate Gaussian - cov (TensorType[..., "dim", "dim"]): Covariance of multivariate Gaussian. + mean: Mean of multivariate Gaussian + cov: Covariance of multivariate Gaussian. """ mean: TensorType[..., "dim"] @@ -107,10 +106,10 @@ def compute_3d_gaussian( """Compute guassian along ray. Args: - directions (TensorType[..., 3]): Axis of Gaussian. - means (TensorType[..., 3]): Mean of Gaussian. - dir_variance (TensorType[..., 3]): Variance along direction axis. - radius_variance (TensorType[..., 3]): Variance tangent to direction axis. + directions: Axis of Gaussian. + means: Mean of Gaussian. + dir_variance: Variance along direction axis. + radius_variance: Variance tangent to direction axis. Returns: Gaussians: Oriented 3D gaussian. @@ -136,11 +135,11 @@ def cylinder_to_gaussian( """Approximates cylinders with a Gaussian distributions. Args: - origins (TensorType[..., 3]): Origins of cylinders. - directions (TensorType[..., 3]): Direction (axis) of cylinders. - starts (TensorType[..., 1]): Start of cylinders. - ends (TensorType[..., 1]): End of cylinders. - radius (TensorType[..., 1]): Radii of cylinders. + origins: Origins of cylinders. + directions: Direction (axis) of cylinders. + starts: Start of cylinders. + ends: End of cylinders. + radius: Radii of cylinders. Returns: Gaussians: Approximation of cylinders @@ -163,11 +162,11 @@ def conical_frustum_to_gaussian( Uses stable parameterization described in mip-NeRF publication. Args: - origins (TensorType[..., 3]): Origins of cones. - directions (TensorType[..., 3]): Direction (axis) of frustums. - starts (TensorType[..., 1]): Start of conical frustums. - ends (TensorType[..., 1]): End of conical frustums. - radius (TensorType[..., 1]): Radii of cone a distance of 1 from the origin. + origins: Origins of cones. + directions: Direction (axis) of frustums. + starts: Start of conical frustums. + ends: End of conical frustums. + radius: Radii of cone a distance of 1 from the origin. Returns: Gaussians: Approximation of conical frustums @@ -184,8 +183,8 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor: """Computes the expected value of sin(y) where y ~ N(x_means, x_vars) Args: - x_means (torch.Tensor): Mean values. - x_vars (torch.Tensor): Variance of values. + x_means: Mean values. + x_vars: Variance of values. Returns: torch.Tensor: The expected value of sin. diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index 5f4a5d0e9..0271e4d12 100644 --- a/nerfactory/utils/misc.py +++ b/nerfactory/utils/misc.py @@ -24,8 +24,12 @@ from typing import Any, Callable, Dict, Union import torch -def get_dict_to_torch(stuff, device: Union[torch.device, str] = "cpu"): - """Set everything in the dict to the specified torch device.""" +def get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = "cpu"): + """Set everything in the dict to the specified torch device. + + Args: + stuff: things to convert to torch + """ if isinstance(stuff, dict): for k, v in stuff.items(): stuff[k] = get_dict_to_torch(v, device) @@ -35,8 +39,12 @@ def get_dict_to_torch(stuff, device: Union[torch.device, str] = "cpu"): return stuff -def get_dict_to_cpu(stuff): - """Set everything in the dict to CPU.""" +def get_dict_to_cpu(stuff: Any): + """Set everything in the dict to CPU. + + Args: + stuff: things to place onto cpu + """ if isinstance(stuff, dict): for k, v in stuff.items(): stuff[k] = get_dict_to_cpu(v) @@ -46,14 +54,19 @@ def get_dict_to_cpu(stuff): return stuff -def is_not_none(var): +def is_not_none(var: Any) -> bool: """Return True if the variable var is None.""" return not isinstance(var, type(None)) def get_masked_dict(d, mask): """Return a masked dictionary. - TODO(ethan): add more asserts/checks so this doesn't have unpredictable behavior.""" + TODO(ethan): add more asserts/checks so this doesn't have unpredictable behavior. + + Args: + d: dict to process + mask: mask to apply to values in dictionary + """ masked_dict = {} for key, value in d.items(): masked_dict[key] = value[mask] @@ -62,7 +75,11 @@ def get_masked_dict(d, mask): def get_hash_str_from_dict(dictionary: Dict[str, Any]) -> str: """MD5 hash of a dictionary. Based on - https://www.doc.ic.ac.uk/~nuric/coding/how-to-hash-a-dictionary-in-python.html""" + https://www.doc.ic.ac.uk/~nuric/coding/how-to-hash-a-dictionary-in-python.html + + Args: + dictionary: dict to process + """ dhash = hashlib.md5() encoded = json.dumps(dictionary, sort_keys=True).encode() dhash.update(encoded) @@ -79,16 +96,16 @@ class IterableWrapper: # pylint: disable=too-few-public-methods for. Args: - new_iter (callable): function that will be called instead as the __iter__() function - new_next (callable): function that will be called instead as the __next__() function - length (int): length of the iterable. If -1, the iterable will be infinite. + new_iter: function that will be called instead as the __iter__() function + new_next: function that will be called instead as the __next__() function + length: length of the iterable. If -1, the iterable will be infinite. Attributes: - new_iter (callable): object's pointer to the function we are calling for __iter__() - new_next (callable): object's pointer to the function we are calling for __next__() - length (int): length of the iterable. If -1, the iterable will be infinite. - i (int): current index of the iterable. + new_iter: object's pointer to the function we are calling for __iter__() + new_next: object's pointer to the function we are calling for __next__() + length: length of the iterable. If -1, the iterable will be infinite. + i: current index of the iterable. """ @@ -112,7 +129,11 @@ class IterableWrapper: # pylint: disable=too-few-public-methods def human_format(num): - """Format a number in a more human readable way""" + """Format a number in a more human readable way + + Args: + num: number to format + """ units = ["", "K", "M", "B", "T", "P"] k = 1000.0 magnitude = int(floor(log(num, k))) diff --git a/nerfactory/utils/poses.py b/nerfactory/utils/poses.py index 56e7f9d5e..628cbaf35 100644 --- a/nerfactory/utils/poses.py +++ b/nerfactory/utils/poses.py @@ -24,10 +24,10 @@ def to4x4(pose: TensorType[..., 3, 4]) -> TensorType[..., 4, 4]: """Convert 3x4 pose matrices to a 4x4 with the addition of a homogeneous coordinate. Args: - pose (TensorType[..., 3, 4]): Camera pose without homogenous coordinate. + pose: Camera pose without homogenous coordinate. Returns: - TensorType[..., 4, 4]: Camera poses with additional homogenous coordinate added. + Camera poses with additional homogenous coordinate added. """ constants = torch.zeros_like(pose[..., :1, :], device=pose.device) constants[..., :, 3] = 1 @@ -38,10 +38,10 @@ def inverse(pose: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: """Invert provided pose matrix. Args: - pose (TensorType[..., 3, 4]): Camera pose without homogenous coordinate. + pose: Camera pose without homogenous coordinate. Returns: - TensorType[..., 3, 4]: Inverse of pose. + Inverse of pose. """ R = pose[..., :3, :3] t = pose[..., :3, 3:] @@ -54,11 +54,11 @@ def multiply(pose_a: TensorType[..., 3, 4], pose_b: TensorType[..., 3, 4]) -> Te """Multiply two pose matrices, A @ B. Args: - pose_a (TensorType[..., 3, 4]): Left pose matrix, usually a transformation applied to the right. - pose_b (TensorType[..., 3, 4]): Right pose matrix, usually a camera pose that will be tranformed by pose_a. + pose_a: Left pose matrix, usually a transformation applied to the right. + pose_b: Right pose matrix, usually a camera pose that will be tranformed by pose_a. Returns: - TensorType[..., 3, 4]: Camera pose matrix where pose_a was applied to pose_b. + Camera pose matrix where pose_a was applied to pose_b. """ R1, t1 = pose_a[..., :3, :3], pose_a[..., :3, 3:] R2, t2 = pose_b[..., :3, :3], pose_b[..., :3, 3:] @@ -71,10 +71,10 @@ def normalize(poses: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: """Normalize the XYZs of poses to fit within a unit cube ([-1, 1]). Note: This operation is not in-place. Args: - poses (TensorType[..., 3, 4]): A collection of poses to be normalized. + poses: A collection of poses to be normalized. Returns; - TensorType[..., 3, 4]: Normalized collection of poses. + Normalized collection of poses. """ pose_copy = torch.clone(poses) pose_copy[..., :3, 3] /= torch.max(torch.abs(poses[..., :3, 3])) diff --git a/nerfactory/utils/profiler.py b/nerfactory/utils/profiler.py index ee69c2bf5..15009466c 100644 --- a/nerfactory/utils/profiler.py +++ b/nerfactory/utils/profiler.py @@ -70,9 +70,9 @@ class Profiler: """update the profiler dictionary with running averages of durations Args: - func_name (str): the function name that is being profiled - start_time (float): the start time when function is called - end_time (float): the end time when function terminated + func_name: the function name that is being profiled + start_time: the start time when function is called + end_time: the end time when function terminated """ val = end_time - start_time func_dict = self.profiler_dict.get(func_name, {"val": 0, "step": 0}) diff --git a/nerfactory/utils/tensor_dataclass.py b/nerfactory/utils/tensor_dataclass.py index 863a2792d..298b115b5 100644 --- a/nerfactory/utils/tensor_dataclass.py +++ b/nerfactory/utils/tensor_dataclass.py @@ -60,7 +60,8 @@ class TensorDataclass: """Finishes setting up the TensorDataclass This will 1) find the broadcasted shape and 2) broadcast all fields to this shape 3) - set _shape to be the broadcasted shape.""" + set _shape to be the broadcasted shape. + """ if not dataclasses.is_dataclass(self): raise TypeError("TensorDataclass must be a dataclass") @@ -83,10 +84,11 @@ class TensorDataclass: """Returns batch shapes of all tensors in a dictionary Args: - dict_ (Dict): The dictionary to get the batch shapes of. + dict_: The dictionary to get the batch shapes of. Returns: - list: The batch shapes of all tensors in the dictionary.""" + The batch shapes of all tensors in the dictionary. + """ batch_shapes = [] for v in dict_.values(): if isinstance(v, torch.Tensor): @@ -101,10 +103,11 @@ class TensorDataclass: """Broadcasts all tensors in a dictionary according to batch_shape Args: - dict_ (Dict): The dictionary to broadcast. + dict_: The dictionary to broadcast. Returns: - Dict: The broadcasted dictionary.""" + The broadcasted dictionary. + """ new_dict = {} for k, v in dict_.items(): if isinstance(v, torch.Tensor): @@ -161,10 +164,10 @@ class TensorDataclass: This should deepcopy as well. Args: - shape (Tuple[int, ...]): The new shape of the tensor dataclass. + shape: The new shape of the tensor dataclass. Returns: - TensorDataclass: A new TensorDataclass with the same data but with a new shape. + A new TensorDataclass with the same data but with a new shape. """ if isinstance(shape, int): shape = (shape,) @@ -187,10 +190,10 @@ class TensorDataclass: meaning it is NOT a deepcopy, and they are still linked. Args: - shape (Tuple[int, ...]): The new shape of the tensor dataclass. + shape: The new shape of the tensor dataclass. Returns: - TensorDataclass: A new TensorDataclass with the same data but with a new shape. + A new TensorDataclass with the same data but with a new shape. """ return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1]))) @@ -201,7 +204,7 @@ class TensorDataclass: device: The device to place the tensor dataclass. Returns: - TensorDataclass: A new TensorDataclass with the same data but on the specified device. + A new TensorDataclass with the same data but on the specified device. """ return self._apply_fn_to_fields(lambda x: x.to(device)) @@ -220,11 +223,11 @@ class TensorDataclass: from dataclasses.fields(self) as we do below and in other places. Args: - fn (Callable): The function to apply to tensor fields. - dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. + fn: The function to apply to tensor fields. + dataclass_fn: The function to apply to TensorDataclass fields. Returns: - TensorDataclass: A new TensorDataclass with the same data but with a new shape. + A new TensorDataclass with the same data but with a new shape. """ new_fields = self._apply_fn_to_dict( @@ -237,12 +240,12 @@ class TensorDataclass: """A helper function for _apply_fn_to_fields, applying a function to all fields of dict_ Args: - dict_ (Dict): The dictionary to apply the function to. - fn (Callable): The function to apply to tensor fields. - dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. + dict_: The dictionary to apply the function to. + fn: The function to apply to tensor fields. + dataclass_fn: The function to apply to TensorDataclass fields. Returns: - Dict: A new dictionary with the same data but with a new shape. Will deep copy""" + A new dictionary with the same data but with a new shape. Will deep copy""" field_names = dict_.keys() new_dict = {} diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 14cfd95db..2de5e7665 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -65,7 +65,12 @@ class EventType(enum.Enum): @check_main_thread def put_image(name, image: TensorType["H", "W", "C"], step: int): - """Setter function to place images into the queue to be written out""" + """Setter function to place images into the queue to be written out + + Args: + image: image to write out + step: step associated with image + """ if isinstance(name, EventName): name = name.value @@ -74,7 +79,13 @@ def put_image(name, image: TensorType["H", "W", "C"], step: int): @check_main_thread def put_scalar(name: str, scalar: float, step: int): - """Setter function to place scalars into the queue to be written out""" + """Setter function to place scalars into the queue to be written out + + Args: + name: name of scalar + scalar: value + step: step associated with scalar + """ if isinstance(name, EventName): name = name.value @@ -83,17 +94,27 @@ def put_scalar(name: str, scalar: float, step: int): @check_main_thread def put_dict(name: str, scalar_dict: float, step: int): - """Setter function to place a dictionary of scalars into the queue to be written out""" + """Setter function to place a dictionary of scalars into the queue to be written out + + Args: + name: name of scalar dictionary + scalar_dict: values to write out + step: step associated with dict + """ EVENT_STORAGE.append({"name": name, "write_type": EventType.DICT, "event": scalar_dict, "step": step}) @check_main_thread def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, update_eta: bool = False): - """Setter function to place a time element into the queue to be written out + """Setter function to place a time element into the queue to be written out. + Processes the time info according to the options. - Processes the time info according to the options: - avg_over_steps (bool): if True, calculate and record a running average of the times - update_eta (bool): if True, update the ETA. should only be set for the training iterations/s + Args: + name: name of time item + duration: value + step: step associated with value + avg_over_steps: if True, calculate and record a running average of the times + update_eta: if True, update the ETA. should only be set for the training iterations/s """ if isinstance(name, EventName): name = name.value @@ -139,7 +160,13 @@ def write_out_storage(): @check_main_thread def setup_event_writers(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None: - """Initialization of all event writers specified in config""" + """Initialization of all event writers specified in config + + Args: + config: configuration to instantiate loggers + max_iter: maximum number of train iterations + banner_messages: list of messages to always display at bottom of screen + """ for writer_type_config in config.writer: if writer_type_config.enable: if isinstance(writer_type_config, cfg.LocalWriterConfig): @@ -164,12 +191,12 @@ class Writer: @abstractmethod def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: - """_summary_ + """method to write out image Args: - name (str): data identifier - image (TensorType["H", "W", "C"]): rendered image to write - step (int): the time step to log + name: data identifier + image: rendered image to write + step: the time step to log """ raise NotImplementedError @@ -178,8 +205,9 @@ class Writer: """Required method to write a single scalar value to the logger Args: - name (str): data identifier - step (int): the time step to log + name: data identifier + scalar: value to write out + step: the time step to log """ raise NotImplementedError @@ -188,8 +216,8 @@ class Writer: """Function that writes out all scalars from a given dictionary to the logger Args: - scalar_dict (dict): dictionary containing all scalar values with key names and quantities - step (int): the time step to log + scalar_dict: dictionary containing all scalar values with key names and quantities + step: the time step to log """ for name, scalar in scalar_dict.items(): self.write_scalar(name, scalar, step) @@ -233,22 +261,10 @@ class WandbWriter(Writer): wandb.init(dir=config.log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: - """_summary_ - - Args: - name (str): data identifier - image (TensorType["H", "W", "C"]): rendered image to write - """ image = torch.permute(image, (2, 0, 1)) wandb.log({name: wandb.Image(image)}, step=step) def write_scalar(self, name: str, scalar: float, step: int) -> None: - """Wandb method to write a single scalar value to the logger - - Args: - name (str): data identifier - scalar (float): scalar value to write - """ wandb.log({name: scalar}, step=step) @@ -261,22 +277,10 @@ class TensorboardWriter(Writer): self.tb_writer = SummaryWriter(log_dir=config.log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: - """_summary_ - - Args: - name (str): data identifier - image (TensorType["H", "W", "C"]): rendered image to write - """ image = to8b(image) self.tb_writer.add_image(name, image, step, dataformats="HWC") def write_scalar(self, name: str, scalar: float, step: int) -> None: - """Tensorboard method to write a single scalar value to the logger - - Args: - name (str): data identifier - scalar (float): scalar value to write - """ self.tb_writer.add_scalar(name, scalar, step) @@ -284,7 +288,7 @@ def _cursorup(x: int): """utility tool to move the cursor up on the terminal Args: - x (int): amount of lines to move cursor upward + x: amount of lines to move cursor upward """ print(f"\r\033[{x}A", end="\x1b[1K\r") @@ -311,14 +315,15 @@ def _format_time(seconds): @decorate_all([check_main_thread]) class LocalWriter(Writer): - """Local Writer Class""" + """Local Writer Class + TODO: migrate to prettyprint + + Args: + config: configuration to instatiate class + banner_messages: list of messages to always display at bottom of screen + """ def __init__(self, config: cfg.LocalWriterConfig, banner_messages: Optional[List[str]] = None): - """ - Args: - stats_to_track (ListConfig): the names of stats that should be logged. - max_log size (int): max number of lines that will be logged to teminal. - """ super().__init__(config.log_dir) self.stats_to_track = [name.value for name in config.stats_to_track] self.max_log_size = config.max_log_size @@ -360,7 +365,12 @@ class LocalWriter(Writer): return latest_map, new_key def _update_header(self, latest_map, new_key): - """helper to handle the printing of the header labels""" + """helper to handle the printing of the header labels + + Args: + latest_map: the most recent dictionary of stats that have been recorded + new_key: indicator whether or not there is a new key added to logger + """ full_log_cond = not self.max_log_size and GLOBAL_BUFFER["step"] <= GLOBAL_BUFFER["steps_per_log"] capped_log_cond = self.max_log_size and (len(self.past_mssgs) - self.banner_len <= 2 or new_key) if full_log_cond or capped_log_cond: @@ -376,7 +386,12 @@ class LocalWriter(Writer): # self.has_printed = True def _print_stats(self, latest_map, padding=" "): - """helper to print out the stats in a readable format""" + """helper to print out the stats in a readable format + + Args: + latest_map: the most recent dictionary of stats that have been recorded + padding: type of characters to print to pad open space + """ step = GLOBAL_BUFFER["step"] fraction_done = step / GLOBAL_BUFFER["max_iter"] curr_mssg = f"{step} ({fraction_done*100:.02f}%)" diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 4794c8cf3..aeeb9833b 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -239,6 +239,8 @@ class VisualizerState: Args: dataset: dataset to render in the scene + start_train: whether to start train when viewer init; + if False, only displays dataset until resume train is toggled """ # clear the current scene self.vis["sceneState/sceneBounds"].delete() @@ -303,7 +305,13 @@ class VisualizerState: return self.check_interrupt def _is_render_step(self, step: int, default_steps: int = 2, max_steps: int = 10) -> bool: - """dynamically calculate when to render grapic based on resolution of image""" + """dynamically calculate when to render grapic based on resolution of image + + Args: + step: current train iteration step + default_steps: base multiple number of steps + max_steps: maximum number of steps in between renders + """ if step != 0: if self.res_upscale_factor == 1: return True @@ -325,6 +333,7 @@ class VisualizerState: Args: outputs: the output tensors for which to apply colormaps on stuff_colors: is only set if colormap is for semantics. Defaults to None. + eps: epsilon to handle floating point comparisons """ # default for rgb images if self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].shape[-1] == 3: @@ -369,6 +378,7 @@ class VisualizerState: Args: outputs: the dictionary of outputs to choose from, from the graph stuff_colors: is only set if colormap is for semantics. Defaults to None. + eps: epsilon to handle floating point comparisons """ if not self.outputs_set: self.vis["renderingState/output_options"].write(list(outputs.keys())) @@ -399,6 +409,7 @@ class VisualizerState: Args: render_time: total time spent rendering current view + num_rays: number of rays rendered image_height: resolution of the current view """ writer.put_time( @@ -429,6 +440,9 @@ class VisualizerState: """ Draw an image using the current camera pose from the viewer. The image is sent of a TCP connection and then uses WebRTC to send it to the viewer. + + Args: + graph: current checkpoint of model """ # check and perform camera updates data = self.vis["renderingState/camera"].read()