diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index 030f1a129..8a032d962 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -31,19 +31,19 @@ from nerfactory.utils.tensor_dataclass import TensorDataclass class Frustums(TensorDataclass): """Describes region of space as a frustum. - Args: - origins (TensorType[..., 3]): xyz coordinate for ray origin. - directions (TensorType[..., 3]): Direction of ray. - starts (TensorType[..., num_samples, 1]): Where the frustum starts along a ray. - ends (TensorType[..., num_samples, 1]): Where the frustum ends along a ray. - pixel_area (TensorType[..., 1]): Projected area of pixel a distance 1 away from origin. + 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. """ - origins: TensorType[..., 3] - directions: TensorType[..., 3] - starts: TensorType[..., 1] - ends: TensorType[..., 1] - pixel_area: TensorType[..., 1] + origins: TensorType["bs":..., 3] + directions: TensorType["bs":..., 3] + starts: TensorType["bs":..., 1] + ends: TensorType["bs":..., 1] + pixel_area: TensorType["bs":..., 1] def get_positions(self) -> TensorType[..., 3]: """Calulates "center" position of frustum. Not weighted by mass. @@ -90,17 +90,17 @@ class RaySamples(TensorDataclass): """Samples along a ray Args: - frustums (Frustums): Frustums along ray. - camera_indices (TensorType[..., 1]): Camera index. - valid_mask (TensorType[..., 1]): Rays that are valid. - deltas (TensorType[..., 1]): "width" of each sample. - ts (TensorType[..., 1]): sample-camera distance of each sample. + 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. """ frustums: Frustums - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[..., 1] = None - deltas: TensorType[..., 1] = None + camera_indices: TensorType["bs":..., 1] = None + valid_mask: TensorType["bs":..., 1] = None + deltas: TensorType["bs":..., 1] = None def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]: """Return weights based on predicted densities @@ -152,13 +152,13 @@ class RayBundle(TensorDataclass): """A bundle of ray parameters. Args: - 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 (TensorType[..., 1]): Camera indices - nears (TensorType[..., 1]): Distance along ray to start sampling - fars (TensorType[..., 1]): Rays Distance along ray to stop sampling - valid_mask (TensorType[..., 1]): Rays that are valid + 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 """ origins: TensorType["num_rays", 3] @@ -167,7 +167,7 @@ class RayBundle(TensorDataclass): camera_indices: Optional[TensorType["num_rays", 1]] = None nears: Optional[TensorType["num_rays", 1]] = None fars: Optional[TensorType["num_rays", 1]] = None - valid_mask: Optional[TensorType["num_rays", 1]] = None + valid_mask: Optional[TensorType["num_rays", 1, bool]] = None num_rays_per_chunk: int = None def move_to_device(self, device: torch.device) -> None: @@ -220,7 +220,7 @@ class RayBundle(TensorDataclass): return self.flatten()[start_idx:end_idx] def get_ray_samples( - self, bin_starts: TensorType[..., "num_samples", 1], bin_ends: TensorType[..., "num_samples", 1] + 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. diff --git a/nerfactory/fields/modules/base.py b/nerfactory/fields/modules/base.py index b50f953c3..5cd14fc77 100644 --- a/nerfactory/fields/modules/base.py +++ b/nerfactory/fields/modules/base.py @@ -59,12 +59,12 @@ class FieldModule(nn.Module): return self.out_dim @abstractmethod - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: """ Args: - in_tensor (TensorType[..., "input_dim"]): Input tensor to process + in_tensor: Input tensor to process Returns: - TensorType[..., "output_dim"]: Processed tensor + Processed tensor """ raise NotImplementedError diff --git a/nerfactory/fields/modules/encoding.py b/nerfactory/fields/modules/encoding.py index 3e6bddcbc..b7086c0ad 100644 --- a/nerfactory/fields/modules/encoding.py +++ b/nerfactory/fields/modules/encoding.py @@ -33,7 +33,7 @@ class Encoding(FieldModule): """Encode an input tensor. Intended to be subclassed Args: - in_dim (int): Input dimension of tensor + in_dim: Input dimension of tensor """ def __init__(self, in_dim: int) -> None: @@ -42,7 +42,7 @@ class Encoding(FieldModule): super().__init__(in_dim=in_dim) @abstractmethod - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: raise NotImplementedError @@ -52,7 +52,7 @@ class Identity(Encoding): def get_out_dim(self) -> int: return self.in_dim - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: return in_tensor @@ -60,9 +60,9 @@ class ScalingAndOffset(Encoding): """Simple scaling and offet to input Args: - in_dim (int): Input dimension of tensor - scaling (float, optional): Scaling applied to tensor. Defaults to 1.0. - offset (float, optional): Offset applied to tensor. Defaults to 0.0. + in_dim: Input dimension of tensor + scaling: Scaling applied to tensor. Defaults to 1.0. + offset: Offset applied to tensor. Defaults to 0.0. """ def __init__(self, in_dim: int, scaling: float = 1.0, offset: float = 0.0) -> None: @@ -74,7 +74,7 @@ class ScalingAndOffset(Encoding): def get_out_dim(self) -> int: return self.in_dim - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: return self.scaling * in_tensor + self.offset @@ -83,11 +83,11 @@ class NeRFEncoding(Encoding): Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp. Args: - in_dim (int): Input dimension of tensor - num_frequencies (int): Number of encoded frequencies per axis - min_freq_exp (float): Minimum frequency exponent - max_freq_exp (float): Maximum frequency exponent - include_input (float): Append the input coordinate to the encoding + in_dim: Input dimension of tensor + num_frequencies: Number of encoded frequencies per axis + min_freq_exp: Minimum frequency exponent + max_freq_exp: Maximum frequency exponent + include_input: Append the input coordinate to the encoding """ def __init__( @@ -107,16 +107,18 @@ class NeRFEncoding(Encoding): return out_dim def forward( - self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None - ) -> TensorType[..., "output_dim"]: + self, + in_tensor: TensorType["bs":..., "input_dim"], + covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None, + ) -> TensorType["bs":..., "output_dim"]: """Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed in mip-NeRF. Args: - in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1. - covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None. + in_tensor: For best performance, the input tensor should be between 0 and 1. + covs: Covariances of input points. Defaults to None. Returns: - TensorType[..., "output_dim"]: Output values will be between -1 and 1 + Output values will be between -1 and 1 """ in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi] freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(in_tensor.device) @@ -141,10 +143,10 @@ class RFFEncoding(Encoding): """Random Fourier Feature encoding. Supports integrated encodings. Args: - in_dim (int): Input dimension of tensor - num_frequencies (int): Number of encoding frequencies - scale (float): Std of Gaussian to sample frequencies. Must be greater than zero - include_input (float): Append the input coordinate to the encoding + in_dim: Input dimension of tensor + num_frequencies: Number of encoding frequencies + scale: Std of Gaussian to sample frequencies. Must be greater than zero + include_input: Append the input coordinate to the encoding """ def __init__(self, in_dim: int, num_frequencies: int, scale: float, include_input: bool = False) -> None: @@ -162,17 +164,19 @@ class RFFEncoding(Encoding): return self.num_frequencies * 2 def forward( - self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None - ) -> TensorType[..., "output_dim"]: + self, + in_tensor: TensorType["bs":..., "input_dim"], + covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None, + ) -> TensorType["bs":..., "output_dim"]: """Calculates RFF encoding. If covariances are provided the encodings will be integrated as proposed in mip-NeRF. Args: - in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1. - covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None. + in_tensor: For best performance, the input tensor should be between 0 and 1. + covs: Covariances of input points. Defaults to None. Returns: - TensorType[..., "output_dim"]: Output values will be between -1 and 1 + Output values will be between -1 and 1 """ in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi] scaled_inputs = in_tensor @ self.b_matrix # [..., "num_frequencies"] @@ -195,12 +199,12 @@ class HashEncoding(Encoding): """Hash encoding Args: - num_levels (int, optional): Number of feature grids. Defaults to 16. - min_res (int, optional): Resolution of smallest feature grid. Defaults to 16. - max_res (int, optional): Resolution of largest feature grid. Defaults to 1024. - hash_table_size (int, optional): Size of hash table. Defaults to 2**19. - features_per_level (int, optional): Number of features per level. Defaults to 2. - hash_init_scale (float, optional): Value to initialize hash grid. Defaults to 0.001. + num_levels: Number of feature grids. Defaults to 16. + min_res: Resolution of smallest feature grid. Defaults to 16. + max_res: Resolution of largest feature grid. Defaults to 1024. + hash_table_size: Size of hash table. Defaults to 2**19. + features_per_level: Number of features per level. Defaults to 2. + hash_init_scale: Value to initialize hash grid. Defaults to 0.001. """ def __init__( @@ -230,14 +234,14 @@ class HashEncoding(Encoding): def get_out_dim(self) -> int: return self.num_levels * self.features_per_level - def hash_fn(self, in_tensor: TensorType[..., "num_levels", 3]) -> TensorType[..., "num_levels"]: + def hash_fn(self, in_tensor: TensorType["bs":..., "num_levels", 3]) -> TensorType["bs":..., "num_levels"]: """Hash tensor using method described in Instant-NGP Args: - in_tensor (TensorType[..., "num_levels", 3]): Tensor to be hashed + in_tensor: Tensor to be hashed Returns: - TensorType[..., "num_levels"]: Hashed tensor + Hashed tensor """ # min_val = torch.min(in_tensor) @@ -252,7 +256,7 @@ class HashEncoding(Encoding): x += self.hash_offset.to(x.device) return x - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: assert in_tensor.shape[-1] == 3 in_tensor = in_tensor[..., None, :] # [..., 1, 3] @@ -299,9 +303,9 @@ class TensorCPEncoding(Encoding): """Learned CANDECOMP/PARFAC (CP) decomposition encoding used in TensoRF Args: - resolution (int, optional): Resolution of grid. Defaults to 256. - num_components (int, optional): Number of components per dimension. Defaults to 24. - init_scale (float, optional): Initialization scale. Defaults to 0.1. + resolution: Resolution of grid. Defaults to 256. + num_components: Number of components per dimension. Defaults to 24. + init_scale: Initialization scale. Defaults to 0.1. """ def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None: @@ -316,7 +320,7 @@ class TensorCPEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] @@ -335,7 +339,7 @@ class TensorCPEncoding(Encoding): """Upsamples underyling feature grid Args: - resolution (int): Target resolution. + resolution: Target resolution. """ self.line_coef.data = F.interpolate( @@ -349,9 +353,9 @@ class TensorVMEncoding(Encoding): """Learned vector-matrix encoding proposed by TensoRF Args: - resolution (int, optional): Resolution of grid. Defaults to 256. - num_components (int, optional): Number of components per dimension. Defaults to 24. - init_scale (float, optional): Initialization scale. Defaults to 0.1. + resolution: Resolution of grid. Defaults to 256. + num_components: Number of components per dimension. Defaults to 24. + init_scale: Initialization scale. Defaults to 0.1. """ def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None: @@ -367,7 +371,7 @@ class TensorVMEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components * 3 - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: in_tensor = in_tensor / 4.0 + 0.5 plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]]) # [3,...,2] line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] @@ -390,7 +394,7 @@ class TensorVMEncoding(Encoding): """Upsamples underyling feature grid Args: - resolution (int): Target resolution. + resolution: Target resolution. """ self.plane_coef.data = F.interpolate( @@ -407,7 +411,7 @@ class SHEncoding(Encoding): """Spherical harmonic encoding Args: - levels (int, optional): Number of spherical hamonic levels to encode. Defaults to 4. + levels: Number of spherical hamonic levels to encode. Defaults to 4. """ def __init__(self, levels: int = 4) -> None: @@ -422,5 +426,5 @@ class SHEncoding(Encoding): return self.levels**2 @torch.no_grad() - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor) diff --git a/nerfactory/fields/modules/field_heads.py b/nerfactory/fields/modules/field_heads.py index a99fbed55..9f816d2ac 100644 --- a/nerfactory/fields/modules/field_heads.py +++ b/nerfactory/fields/modules/field_heads.py @@ -40,10 +40,10 @@ class FieldHead(FieldModule): """Base field output Args: - out_dim (int): output dimension for renderer - field_head_name (FieldHeadNames): Field type - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + out_dim: output dimension for renderer + field_head_name: Field type + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__( @@ -71,14 +71,14 @@ class FieldHead(FieldModule): def _construct_net(self): self.net = nn.Linear(self.in_dim, self.out_dim) - def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType: + def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: """Process network output for renderer Args: - in_tensor (TensorType[..., "in_dim]): Network input + in_tensor: Network input Returns: - TensorType[..., "out_dim"]: Render head output + Render head output """ if not self.net: raise SystemError("in_dim not set. Must be provided to construtor, or set_in_dim() should be called.") @@ -92,8 +92,8 @@ class DensityFieldHead(FieldHead): """Density output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: @@ -104,8 +104,8 @@ class RGBFieldHead(FieldHead): """RGB output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: @@ -116,10 +116,10 @@ class SHFieldHead(FieldHead): """Spherical harmonics output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - levels (int, optional): Number of spherical harmonics layers. Defaults to 3. - channels (int, optional): Number of channels. Defaults to 3 (ie RGB). - activation (Optional[nn.Module], optional): Output activation. Defaults to None. + in_dim: input dimension. If not defined in constructor, it must be set later. + levels: Number of spherical harmonics layers. Defaults to 3. + channels: Number of channels. Defaults to 3 (ie RGB). + activation: Output activation. Defaults to None. """ def __init__( @@ -134,8 +134,8 @@ class UncertaintyFieldHead(FieldHead): """Uncertainty output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: @@ -146,8 +146,8 @@ class TransientRGBFieldHead(FieldHead): """Transient RGB output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: @@ -158,8 +158,8 @@ class TransientDensityFieldHead(FieldHead): """Transient density output Args: - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Softplus()) -> None: @@ -172,9 +172,9 @@ class SemanticFieldHead(FieldHead): """Semantic stuff output Args: - num_classes (int): Number of semantic classes - in_dim (int, optional): input dimension. If not defined in constructor, it must be set later. - activation (Optional[nn.Module]): output head activation + num_classes: Number of semantic classes + in_dim: input dimension. If not defined in constructor, it must be set later. + activation: output head activation """ def __init__(self, num_classes: int, in_dim: Optional[int] = None) -> None: diff --git a/nerfactory/fields/modules/mlp.py b/nerfactory/fields/modules/mlp.py index f61ba7451..a55de10b1 100644 --- a/nerfactory/fields/modules/mlp.py +++ b/nerfactory/fields/modules/mlp.py @@ -28,12 +28,12 @@ class MLP(FieldModule): """Multilayer perceptron Args: - in_dim (int): Input layer dimension - num_layers (int): Number of network layers - layer_width (int): Width of each MLP layer - out_dim (int, optional): Ouput layer dimension. Defaults to layer_width - activation (Optional[nn.Module], optional): intermediate layer activation function. Defaults to nn.ReLU. - out_activation (Optional[nn.Module], optional): output activation function. Defaults to None. + in_dim: Input layer dimension + num_layers: Number of network layers + layer_width: Width of each MLP layer + out_dim: Ouput layer dimension. Defaults to layer_width + activation: intermediate layer activation function. Defaults to nn.ReLU. + out_activation: output activation function. Defaults to None. """ def __init__( @@ -76,14 +76,14 @@ class MLP(FieldModule): layers.append(nn.Linear(self.layer_width, self.out_dim)) self.layers = nn.ModuleList(layers) - def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]: + def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: """Process input with a multilayer perceptron. Args: - in_tensor (TensorType[..., "in_dim]): Network input + in_tensor: Network input Returns: - TensorType[..., "out_dim"]: MLP network output + MLP network output """ x = in_tensor for i, layer in enumerate(self.layers): diff --git a/nerfactory/fields/modules/spatial_distortions.py b/nerfactory/fields/modules/spatial_distortions.py index c225dc0ea..41a54002e 100644 --- a/nerfactory/fields/modules/spatial_distortions.py +++ b/nerfactory/fields/modules/spatial_distortions.py @@ -27,13 +27,15 @@ from nerfactory.utils.math import Gaussians class SpatialDistortion(nn.Module): """Apply spatial distortions""" - def forward(self, positions: Union[TensorType[..., 3], Gaussians]) -> Union[TensorType[..., 3], Gaussians]: + def forward( + self, positions: Union[TensorType["bs":..., 3], Gaussians] + ) -> Union[TensorType["bs":..., 3], Gaussians]: """ Args: - positions (Union[TensorType[..., 3], Gaussians]): Sample to distort + positions: Sample to distort Returns: - Union[TensorType[..., 3], Gaussians]: distorted sample + Union: distorted sample """ diff --git a/nerfactory/graphs/base.py b/nerfactory/graphs/base.py index a13e89645..77cbe44eb 100644 --- a/nerfactory/graphs/base.py +++ b/nerfactory/graphs/base.py @@ -43,7 +43,7 @@ def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) """Setup the graph. The dataset inputs should be set with the training data. Args: - dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters. + dataset_inputs: The inputs which will be used to define the camera parameters. """ graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict()) graph.to(device) @@ -77,14 +77,14 @@ class Graph(AbstractGraph): subclassed for custom NeRF model. Args: - intrinsics (torch.Tensor): Camera intrinsics. - camera_to_world (torch.Tensor): Camera to world transformation. - loss_coefficients (DictConfig): Loss specific weights. - scene_bounds (SceneBounds): Bounds of target scene. - enable_collider (bool): Whether to create a scene collider to filter rays. - collider_config (DictConfig): Configuration of scene collider. - enable_density_field (bool): Whether to create a density field to filter samples. - density_field_config (DictConfig): Configuration of density field. + intrinsics: Camera intrinsics. + camera_to_world: Camera to world transformation. + loss_coefficients: Loss specific weights. + scene_bounds: Bounds of target scene. + 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. """ def __init__( @@ -145,7 +145,7 @@ class Graph(AbstractGraph): """Obtain the parameter groups for the optimizers Returns: - Dict[str, List[Parameter]]: Mapping of different parameter groups + Mapping of different parameter groups """ @abstractmethod @@ -153,10 +153,10 @@ class Graph(AbstractGraph): """Takes in a Ray Bundle and returns a dictionary of outputs. Args: - ray_bundle (RayBundle): Input bundle of rays. + ray_bundle: Input bundle of rays. Returns: - dict: Outputs of graph. (ie. rendered colors) + Outputs of graph. (ie. rendered colors) """ def process_outputs_as_images(self, outputs): # pylint:disable=no-self-use diff --git a/nerfactory/graphs/modules/ray_generator.py b/nerfactory/graphs/modules/ray_generator.py index 18000b1d7..a6393da27 100644 --- a/nerfactory/graphs/modules/ray_generator.py +++ b/nerfactory/graphs/modules/ray_generator.py @@ -27,9 +27,8 @@ class RayGenerator(nn.Module): This class will store the intrinsics and extrinsics parameters of the cameras. Args: - intrinsics (TensorType["num_cameras", "num_intrinsics_params"]): - The intrinsics parameters. - camera_to_world (TensorType["num_cameras", 3, 4]): Camera to world transformation matrix. + intrinsics: The intrinsics parameters. + camera_to_world: Camera to world transformation matrix. """ def __init__( @@ -54,7 +53,7 @@ class RayGenerator(nn.Module): """Index into the cameras to generate the rays. Args: - ray_indices (TensorType["num_rays", 3]): Contains camera, row, and col indicies for target rays. + ray_indices: Contains camera, row, and col indicies for target rays. """ c = ray_indices[:, 0] # camera indices y = ray_indices[:, 1] # row indices diff --git a/nerfactory/graphs/modules/ray_losses.py b/nerfactory/graphs/modules/ray_losses.py index 4bf70e32f..69ece5132 100644 --- a/nerfactory/graphs/modules/ray_losses.py +++ b/nerfactory/graphs/modules/ray_losses.py @@ -20,7 +20,9 @@ from torchtyping import TensorType from nerfactory.cameras.rays import RaySamples -def distortion_loss(ray_samples: RaySamples, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., 1]: +def distortion_loss( + ray_samples: RaySamples, densities: TensorType["bs":..., "num_samples", 1] +) -> TensorType["bs":..., 1]: """Ray baserd distortion loss proposed in MipNeRF-360. .. math:: @@ -32,11 +34,11 @@ def distortion_loss(ray_samples: RaySamples, densities: TensorType[..., "num_sam is the weight at location :math:`u` between bin locations :math:`s_i` and :math:`s_{i+1}`. Args: - ray_samples (RaySamples): Ray samples to compute loss over - densities (TensorType[..., "num_samples", 1]): Predicted sample densities + ray_samples: Ray samples to compute loss over + densities: Predicted sample densities Returns: - TensorType[..., 1]: Distortion Loss. + Distortion Loss. """ # Compute the weight at each sample location diff --git a/nerfactory/graphs/modules/ray_sampler.py b/nerfactory/graphs/modules/ray_sampler.py index 6ac0c9efc..b293e0b60 100644 --- a/nerfactory/graphs/modules/ray_sampler.py +++ b/nerfactory/graphs/modules/ray_sampler.py @@ -69,13 +69,12 @@ class SpacedSampler(Sampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - spacing_fn (Callable): Function that dictates sample spacing (ie `lambda x : x` is uniform). - spacing_fn_inv (Callable): The inverse of spacing_fn. - train_stratified (bool): Use stratified sampling during training. Defults to True - density_field (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. + 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 @@ -91,11 +90,11 @@ class SpacedSampler(Sampler): """Generates position samples accoring to spacing function. Args: - ray_bundle (RayBundle): Rays to generate samples for - num_samples (Optional[int]): Number of samples per ray + ray_bundle: Rays to generate samples for + num_samples: Number of samples per ray Returns: - RaySamples: Positions and deltas for samples along a ray + Positions and deltas for samples along a ray """ assert ray_bundle is not None assert ray_bundle.nears is not None @@ -133,11 +132,10 @@ class UniformSampler(SpacedSampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - train_stratified (bool): Use stratified sampling during training. Defults to True - density_field (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. + 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, @@ -161,11 +159,10 @@ class LinearDisparitySampler(SpacedSampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - train_stratified (bool): Use stratified sampling during training. Defults to True - density_field (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. + 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, @@ -189,11 +186,10 @@ class SqrtSampler(SpacedSampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - train_stratified (bool): Use stratified sampling during training. Defults to True - density_field (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. + 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, @@ -217,11 +213,10 @@ class LogSampler(SpacedSampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - train_stratified (bool): Use stratified sampling during training. Defults to True - density_field (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. + 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, @@ -247,13 +242,12 @@ class PDFSampler(Sampler): ) -> None: """ Args: - num_samples (int): Number of samples per ray - train_stratified: boolean: Randomize location within each bin during training. Defaults to True + 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 (DensityGrid, optional): Density grid. If provides, - samples below weight_threshold as set as invalid. - weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. - histogram_padding (float): Amount to weights prior to computing PDF. Defaults to 0.01. + 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 @@ -272,14 +266,14 @@ class PDFSampler(Sampler): """Generates position samples given a distribution. Args: - ray_bundle (RayBundle): Rays to generate samples for - ray_samples (RaySamples): Existing ray samples - weights: (TensorType[..., "num_samples", 1]): Weights for each bin - num_samples (Optional[int]): Number of samples per ray - eps: float: Small value to prevent numerical issues. Defaults to 1e-5 + ray_bundle: Rays to generate samples for + ray_samples: Existing ray samples + weights: Weights for each bin + num_samples: Number of samples per ray + eps: Small value to prevent numerical issues. Defaults to 1e-5 Returns: - RaySamples: Positions and deltas for samples along a ray + Positions and deltas for samples along a ray """ num_samples = num_samples or self.num_samples num_bins = num_samples + 1 @@ -367,12 +361,11 @@ class NGPSpacedSampler(Sampler): TODO(ruilongli): maybe move aabb test to the collision detector? Args: - ray_bundle (RayBundle): Rays to generate samples for - aabb (TensorType[2, 3]): Bounding box of the scene. - num_samples (Optional[int]): Number of samples per ray + ray_bundle: Rays to generate samples for + aabb: Bounding box of the scene. + num_samples: Number of samples per ray Returns: - Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples"], TensorType["total_samples"]]: First return is ray samples in a packed way where only the valid samples are kept. Second return contains all the information to recover packed samples into unpacked mode for rendering. The last two returns are t_min and t_max from ray-aabb test.