From 1c1a58c4e55d8817d1f7db7cfeb472721fc25ed8 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 13 May 2022 17:44:16 -0700 Subject: [PATCH] nerf-w and semantic-nerf initial implementations (#73) * initial nerfw code * nerfw should be implmented correctly now * semantic nerf in progress. data loader is too slow and not clean right now * added implementation of semantic nerf * mend --- .vscode/launch.json | 32 +++- configs/graph/default_graph.yaml | 2 +- configs/graph/nerfw_graph.yaml | 5 + configs/graph/semantic_nerf_graph.yaml | 10 + configs/nerfw.yaml | 14 ++ configs/semantic_nerf.yaml | 14 ++ mattport/nerf/dataset/friends.py | 36 +++- mattport/nerf/dataset/image_dataset.py | 37 +++- mattport/nerf/dataset/structs.py | 27 ++- mattport/nerf/dataset/utils.py | 24 ++- mattport/nerf/field_modules/embedding.py | 30 +++ mattport/nerf/field_modules/field_heads.py | 54 +++++- mattport/nerf/graph/base.py | 4 +- mattport/nerf/graph/instant_ngp.py | 44 +---- mattport/nerf/graph/nerfw.py | 205 ++++++++++++++++++++- mattport/nerf/graph/semantic_nerf.py | 158 +++++++++++++++- mattport/nerf/graph/vanilla_nerf.py | 25 ++- mattport/nerf/renderers.py | 36 ++++ mattport/nerf/trainer.py | 12 +- mattport/structures/rays.py | 31 ++-- scripts/run_train.py | 3 +- tests/nerf/field_modules/test_embedding.py | 14 ++ 22 files changed, 725 insertions(+), 92 deletions(-) create mode 100644 configs/graph/semantic_nerf_graph.yaml create mode 100644 configs/nerfw.yaml create mode 100644 configs/semantic_nerf.yaml create mode 100644 mattport/nerf/field_modules/embedding.py create mode 100644 tests/nerf/field_modules/test_embedding.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 1054dacb4..4fc7507dd 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -17,9 +17,9 @@ "type": "python", "request": "launch", "program": "scripts/run_train.py", - // "env": { - // "CUDA_VISIBLE_DEVICES": "8" - // }, + "env": { + "CUDA_VISIBLE_DEVICES": "8" + }, "console": "integratedTerminal" }, { @@ -42,5 +42,31 @@ "--config-name=instant_ngp.yaml" ] }, + { + "name": "Python: run_train nerfw (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "env": { + "CUDA_VISIBLE_DEVICES": "8" + }, + "console": "integratedTerminal", + "args": [ + "--config-name=nerfw.yaml" + ] + }, + { + "name": "Python: run_train semantic_nerf (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "env": { + "CUDA_VISIBLE_DEVICES": "4" + }, + "console": "integratedTerminal", + "args": [ + "--config-name=semantic_nerf.yaml" + ] + }, ] } \ No newline at end of file diff --git a/configs/graph/default_graph.yaml b/configs/graph/default_graph.yaml index 819b333e4..eb29d08c6 100644 --- a/configs/graph/default_graph.yaml +++ b/configs/graph/default_graph.yaml @@ -1,7 +1,7 @@ # saving model specifications model_dir: /tmp/mattport_models/ steps_per_save: null -steps_per_test: 1000 +steps_per_test: 100 max_num_iterations: 100000 # pretrained loading options diff --git a/configs/graph/nerfw_graph.yaml b/configs/graph/nerfw_graph.yaml index e89c04c80..5036f4cf8 100644 --- a/configs/graph/nerfw_graph.yaml +++ b/configs/graph/nerfw_graph.yaml @@ -4,3 +4,8 @@ defaults: network: _target_: mattport.nerf.graph.nerfw.NerfWGraph + loss_coefficients: # loss coefficients + rgb_loss_coarse: 1.0 + rgb_loss_fine: 1.0 + uncertainty_loss: 1.0 + density_loss: 0.01 diff --git a/configs/graph/semantic_nerf_graph.yaml b/configs/graph/semantic_nerf_graph.yaml new file mode 100644 index 000000000..74c55b8d2 --- /dev/null +++ b/configs/graph/semantic_nerf_graph.yaml @@ -0,0 +1,10 @@ +defaults: + - default_graph.yaml + - _self_ + +network: + _target_: mattport.nerf.graph.semantic_nerf.SemanticNerfGraph + loss_coefficients: # loss coefficients + rgb_loss_coarse: 1.0 + rgb_loss_fine: 1.0 + semantic_loss_fine: 0.05 \ No newline at end of file diff --git a/configs/nerfw.yaml b/configs/nerfw.yaml new file mode 100644 index 000000000..45a3cfe64 --- /dev/null +++ b/configs/nerfw.yaml @@ -0,0 +1,14 @@ +defaults: + - default + - override data/dataset: friends_TBBT-big_living_room + - override graph: nerfw_graph + - _self_ + +experiment_name: friends_TBBT-big_living_room +method_name: nerfw + +network: + near_plane: 0.1 + far_plane: 10.0 + num_coarse_samples: 64 + num_importance_samples: 64 diff --git a/configs/semantic_nerf.yaml b/configs/semantic_nerf.yaml new file mode 100644 index 000000000..885ae6db4 --- /dev/null +++ b/configs/semantic_nerf.yaml @@ -0,0 +1,14 @@ +defaults: + - default + - override data/dataset: friends_TBBT-big_living_room + - override graph: semantic_nerf_graph + - _self_ + +experiment_name: friends_TBBT-big_living_room +method_name: semantic_nerf + +network: + near_plane: 0.1 + far_plane: 10.0 + num_coarse_samples: 64 + num_importance_samples: 64 diff --git a/mattport/nerf/dataset/friends.py b/mattport/nerf/dataset/friends.py index 33bfce485..552bdb858 100644 --- a/mattport/nerf/dataset/friends.py +++ b/mattport/nerf/dataset/friends.py @@ -8,10 +8,13 @@ import numpy as np import torch from mattport.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary -from mattport.nerf.dataset.structs import DatasetInputs, PointCloud +from mattport.nerf.dataset.structs import DatasetInputs, PointCloud, Semantics +from mattport.utils import profiler +from mattport.utils.io import load_from_json -def load_friends_data(basedir, downscale_factor=1, split="train", include_point_cloud=False): +@profiler.time_function +def load_friends_data(basedir, downscale_factor=1, split="train", include_semantics=True, include_point_cloud=False): """_summary_ Args: @@ -40,7 +43,6 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_point_ intrinsics = [] for image_path in image_paths: cam = cameras_data[image_path_to_image_id[image_path]] - print(cam.params) assert len(cam.params) == 3 focal_length = cam.params[0] # f (fx and fy) cx = cam.params[1] # cx @@ -63,7 +65,27 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_point_ camera_to_world[..., 1:3] *= -1 # --- masks to mask out things (e.g., people) --- - # TODO: + + # --- semantics --- + semantics = Semantics() + if include_semantics: + thing_filenames = [ + image_filename.replace("/images/", "/segmentations/thing/").replace(".jpg", ".png") + for image_filename in image_filenames + ] + stuff_filenames = [ + image_filename.replace("/images/", "/segmentations/stuff/").replace(".jpg", ".png") + for image_filename in image_filenames + ] + panoptic_classes = load_from_json(os.path.join(basedir, "panoptic_classes.json")) + stuff_classes = panoptic_classes["stuff"] + thing_classes = panoptic_classes["thing"] + semantics = Semantics( + stuff_classes=stuff_classes, + stuff_filenames=stuff_filenames, + thing_classes=thing_classes, + thing_filenames=thing_filenames, + ) # --- possibly transform --- # TODO: @@ -73,18 +95,20 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_point_ # Possibly include the sparse point cloud from COLMAP in the dataset inputs. # NOTE(ethan): this will be common across the different splits. - point_cloud = None + point_cloud = PointCloud() if include_point_cloud: points_3d = read_pointsTD_binary(os.path.join(basedir, "colmap", "points3D.bin")) xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])) rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])) - point_cloud = PointCloud(xyz=xyz, rgb=rgb) + point_cloud.xyz = xyz + point_cloud.rgb = rgb dataset_inputs = DatasetInputs( image_filenames=image_filenames, downscale_factor=downscale_factor, intrinsics=intrinsics / downscale_factor, camera_to_world=camera_to_world, + semantics=semantics, point_cloud=point_cloud, ) return dataset_inputs diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index dfa59f299..907d45e7b 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -12,6 +12,8 @@ from PIL import Image from torch.utils.data import default_collate from torchtyping import TensorType +from mattport.nerf.dataset.structs import Semantics + def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False): """_summary_ @@ -31,6 +33,8 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False) c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] pixels = batch["image"][c, y, x] mask = batch["mask"][c, y, x] + if "stuff_image" in batch: + stuff_image = batch["stuff_image"][c, y, x] assert pixels.shape == (num_rays_per_batch, 3), pixels.shape # Needed to correct the random indices to their actual camera idx locations. @@ -42,6 +46,10 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False) "pixels": pixels, "mask": mask, } + + if "stuff_image" in batch: + collated_batch["stuff_image"] = stuff_image + if keep_full_image: collated_batch["image"] = batch["image"] @@ -52,7 +60,12 @@ class ImageDataset(torch.utils.data.Dataset): """Dataset that returns images.""" def __init__( - self, image_filenames: List[str], downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None + self, + image_filenames: List[str], + downscale_factor: int = 1, + alpha_color: Optional[TensorType[3]] = None, + semantics: Semantics = Semantics(), + **kwargs, ): """_summary_ @@ -66,6 +79,8 @@ class ImageDataset(torch.utils.data.Dataset): self.image_filenames = image_filenames self.downscale_factor = downscale_factor self.alpha_color = alpha_color + self.semantics = semantics + self.kwargs = kwargs def __len__(self): return len(self.image_filenames) @@ -123,4 +138,24 @@ class ImageDataset(torch.utils.data.Dataset): "image": image, # the pixels "mask": mask, } + + if self.semantics.stuff_filenames: + stuff_image_filename = self.semantics.stuff_filenames[image_idx] + pil_image = Image.open(stuff_image_filename) + if self.downscale_factor != 1.0: + image_width, image_height = pil_image.size + if image_width % self.downscale_factor != 0: + raise ValueError( + f"Image width {image_width} is not divisible by downscale_factor {self.downscale_factor}" + ) + if image_height % self.downscale_factor != 0: + raise ValueError( + f"Image height {image_height} is not divisible by downscale_factor {self.downscale_factor}" + ) + # NOTE(ethan): the use of NEAREST is important for semantic classes + pil_image = pil_image.resize( + (image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.NEAREST + ) + stuff_image = torch.from_numpy(np.array(pil_image, dtype="int32"))[..., None] + data["stuff_image"] = stuff_image return data diff --git a/mattport/nerf/dataset/structs.py b/mattport/nerf/dataset/structs.py index 39ca9d371..d7c689169 100644 --- a/mattport/nerf/dataset/structs.py +++ b/mattport/nerf/dataset/structs.py @@ -13,8 +13,18 @@ from torchtyping import TensorType class PointCloud: """_summary_""" - xyz: None - rgb: None + xyz = None + rgb = None + + +@dataclass +class Semantics: + """_summary_""" + + stuff_classes: List[str] = None + stuff_filenames: List[str] = None + thing_classes: List[str] = None + thing_filenames: List[str] = None @dataclass @@ -31,12 +41,21 @@ class DatasetInputs: intrinsics: torch.tensor = None camera_to_world: torch.tensor = None bounding_box: torch.tensor = None + mask_filenames: List[str] = None depth_filenames: List[str] = None - semantic_filenames: List[str] = None - point_cloud: PointCloud = None # TODO(ethan): specify the type here + semantics: Semantics = Semantics() + point_cloud: PointCloud = PointCloud() alpha_color: Optional[TensorType[3]] = None def check_inputs(self): """Check the inputs to make sure everything is okay.""" assert self.intrinsics.dtype == torch.float32 assert self.camera_to_world.dtype == torch.float32 + + def save_to_folder_name(self, data_directory: str): + """Save the dataset inputs.""" + raise NotImplementedError + + def load_from_folder_name(self, data_directory: str): + """Load the saved dataset inputs.""" + raise NotImplementedError diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 627cd7f93..97b0d8754 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -19,6 +19,7 @@ def get_dataset_inputs_dict( downscale_factor: int = 1, alpha_color: Optional[Union[str, list, ListConfig]] = None, splits: Tuple[str] = ("train", "val"), + load_dataset_inputs_from_cache: bool = False, ) -> Dict[str, DatasetInputs]: """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator. # TODO: implement the `test` split, which will have depths and normals, etc. @@ -28,12 +29,16 @@ def get_dataset_inputs_dict( dataset_type (str): Name of dataset type downscale_factor (int, optional): How much to downscale images. Defaults to 1.0. alpha_color (str, list, optional): Sets transparent regions to specified color, otherwise black. - + load_from_cache (bool) Returns: Dict[str, DatasetInputs]: The inputs needed for generating rays. """ dataset_inputs_dict = {} + if load_dataset_inputs_from_cache: + print("Loading from cache.") + return None + if alpha_color is not None: alpha_color = get_color(alpha_color) @@ -47,12 +52,23 @@ def get_dataset_inputs_dict( ) dataset_inputs_dict[split] = dataset_inputs elif dataset_type == "friends": + # TODO(ethan): change this hack, and do proper splits for the friends dataset + # currently we assume that there is only the training set of images + dataset_inputs = load_friends_data( + get_absolute_path(data_directory), + downscale_factor=downscale_factor, + split="train", + include_point_cloud=False, + ) for split in splits: - dataset_inputs = load_friends_data( - get_absolute_path(data_directory), downscale_factor=downscale_factor, split=split - ) dataset_inputs_dict[split] = dataset_inputs else: raise NotImplementedError(f"{dataset_type} is not a valid dataset type") + for split in splits: + # save the dataset inputs to cache + # TODO: + # dataset_inputs_dict[splits].save_to_folder_name() + pass + return dataset_inputs_dict diff --git a/mattport/nerf/field_modules/embedding.py b/mattport/nerf/field_modules/embedding.py new file mode 100644 index 000000000..ae22d1fc4 --- /dev/null +++ b/mattport/nerf/field_modules/embedding.py @@ -0,0 +1,30 @@ +""" +Code for embeddings. +""" + + +import torch +from torchtyping import TensorType +from mattport.nerf.field_modules.base import FieldModule + + +class Embedding(FieldModule): + """Index into embeddings.""" + + def __init__(self, in_dim: int, out_dim: int) -> None: + """ + Args: + in_dim (int): Number of embeddings + out_dim (int): Dimension of the embedding vectors + """ + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.build_nn_modules() + + def build_nn_modules(self) -> None: + self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim) + + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """Call forward""" + return self.embedding(in_tensor) diff --git a/mattport/nerf/field_modules/field_heads.py b/mattport/nerf/field_modules/field_heads.py index 5af4ded2a..0357f0bd7 100644 --- a/mattport/nerf/field_modules/field_heads.py +++ b/mattport/nerf/field_modules/field_heads.py @@ -15,6 +15,11 @@ class FieldHeadNames(Enum): RGB = "rgb" DENSITY = "density" + UNCERTAINTY = "uncertainty" + TRANSIENT_RGB = "transient_rgb" + TRANSIENT_DENSITY = "transient_density" + SEMANTICS_STUFF = "semantics_stuff" + SEMANTICS_THING = "semantics_thing" class FieldHead(FieldModule): @@ -66,12 +71,53 @@ class FieldHead(FieldModule): class DensityFieldHead(FieldHead): """Density output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None: - super().__init__(in_dim, out_dim=1, field_head_name=FieldHeadNames.DENSITY, activation=activation) + def __init__( + self, in_dim: int, field_head_name=FieldHeadNames.DENSITY, activation: Optional[nn.Module] = nn.Softplus() + ) -> None: + super().__init__(in_dim, out_dim=1, field_head_name=field_head_name, activation=activation) class RGBFieldHead(FieldHead): """RGB output""" - def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: - super().__init__(in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation) + def __init__( + self, in_dim: int, field_head_name=FieldHeadNames.RGB, activation: Optional[nn.Module] = nn.Sigmoid() + ) -> None: + super().__init__(in_dim, out_dim=3, field_head_name=field_head_name, activation=activation) + + +class UncertaintyFieldHead(FieldHead): + """Uncertainty output""" + + def __init__( + self, in_dim: int, field_head_name=FieldHeadNames.UNCERTAINTY, activation: Optional[nn.Module] = nn.Softplus() + ) -> None: + super().__init__(in_dim, out_dim=1, field_head_name=field_head_name, activation=activation) + + +class TransientRGBHead(RGBFieldHead): + """Transient RGB output""" + + def __init__( + self, in_dim: int, field_head_name=FieldHeadNames.TRANSIENT_RGB, activation: Optional[nn.Module] = nn.Sigmoid() + ) -> None: + super().__init__(in_dim, field_head_name, activation) + + +class TransientDensityHead(DensityFieldHead): + """Transient density output""" + + def __init__( + self, + in_dim: int, + field_head_name=FieldHeadNames.TRANSIENT_DENSITY, + activation: Optional[nn.Module] = nn.Softplus(), + ) -> None: + super().__init__(in_dim, field_head_name, activation) + + +class SemanticStuffHead(FieldHead): + """Semantic stuff output""" + + def __init__(self, in_dim: int, num_classes: int, field_head_name=FieldHeadNames.SEMANTICS_STUFF) -> None: + super().__init__(in_dim, out_dim=num_classes, field_head_name=field_head_name, activation=None) diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index 32e4b874e..045fae1ee 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -111,13 +111,15 @@ class Graph(nn.Module): return aggregated_loss @torch.no_grad() - def get_outputs_for_camera(self, intrinsics, camera_to_world, chunk_size=1024): + def get_outputs_for_camera(self, intrinsics, camera_to_world, chunk_size=1024, training_camera_index=0): """Takes in camera parameters and computes the output of the graph.""" assert len(intrinsics.shape) == 1 num_intrinsics_params = len(intrinsics) camera_class = get_camera_model(num_intrinsics_params) camera = camera_class(*intrinsics.tolist(), camera_to_world=camera_to_world) camera_ray_bundle = camera.generate_camera_rays() + # TODO(ethan): decide how to properly handle the image indices for validation images + camera_ray_bundle.set_camera_indices(camera_index=training_camera_index) image_height, image_width = camera_ray_bundle.origins.shape[:2] device = self.get_device() diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py index 8bb9fd6ca..7a9bcd340 100644 --- a/mattport/nerf/graph/instant_ngp.py +++ b/mattport/nerf/graph/instant_ngp.py @@ -2,28 +2,26 @@ Implementation of Instant NGP. """ -import torch from torch import nn from mattport.nerf.field_modules.encoding import HashEncoding, NeRFEncoding -from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP -from mattport.nerf.graph.vanilla_nerf import NeRFGraph -from mattport.structures.rays import RaySamples +from mattport.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph -class NGPField(nn.Module): +class NGPField(NeRFField): """NeRF module""" def __init__(self, num_layers=3, layer_width=64) -> None: - super().__init__() - self.num_layers = num_layers - self.layer_width = layer_width + super().__init__(num_layers=num_layers, layer_width=layer_width) + def build_encodings(self): self.encoding_xyz = HashEncoding() self.encoding_dir = NeRFEncoding( in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True ) + + def build_mlp_base(self): self.mlp_base = MLP( in_dim=self.encoding_xyz.get_out_dim(), out_dim=self.layer_width, @@ -31,36 +29,6 @@ class NGPField(nn.Module): layer_width=self.layer_width, activation=nn.ReLU(), ) - self.mlp_rgb = MLP( - in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(), - out_dim=self.layer_width // 2, - num_layers=1, - activation=nn.ReLU(), - ) - self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) - self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) - - def forward(self, ray_samples: RaySamples): - """Evaluates the field at points along the ray - Args: - xyz: () - # TODO(ethan): change the input to be something more abstracted - e.g., a FieldInput structure - """ - positions = ray_samples.positions - directions = ray_samples.directions - encoded_xyz = self.encoding_xyz(positions) - encoded_dir = self.encoding_dir(directions) - base_mlp_out = self.mlp_base(encoded_xyz) - rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1)) - - field_rgb_output = self.field_output_rgb(rgb_mlp_out) - field_density_out = self.field_output_density(base_mlp_out) - - field_outputs = {} - field_outputs.update(field_rgb_output) - field_outputs.update(field_density_out) - return field_outputs class NGPGraph(NeRFGraph): diff --git a/mattport/nerf/graph/nerfw.py b/mattport/nerf/graph/nerfw.py index ee6971944..80fba0402 100644 --- a/mattport/nerf/graph/nerfw.py +++ b/mattport/nerf/graph/nerfw.py @@ -3,18 +3,215 @@ NeRF-W (NeRF in the wild) implementation. TODO: """ -from mattport.nerf.graph.vanilla_nerf import NeRFGraph -from mattport.structures.rays import RayBundle +from typing import Tuple + +import torch +from torch import nn + +from mattport.nerf.field_modules.embedding import Embedding +from mattport.nerf.field_modules.field_heads import ( + FieldHeadNames, + TransientDensityHead, + TransientRGBHead, + UncertaintyFieldHead, +) +from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph +from mattport.nerf.renderers import UncertaintyRenderer +from mattport.structures.rays import RayBundle, RaySamples +from mattport.utils import visualization, writer + + +class NerfWField(NeRFField): + """The NeRF-W field which has appearance and transient conditioning.""" + + def __init__( + self, + num_layers=8, + layer_width=256, + skip_connections: Tuple = (4,), + num_images: int = 0, + appearance_embedding_dim: int = 48, + transient_embedding_dim: int = 16, + ) -> None: + assert num_images > 0 + self.num_images = num_images + self.appearance_embedding_dim = appearance_embedding_dim + self.transient_embedding_dim = transient_embedding_dim + super().__init__(num_layers, layer_width, skip_connections) + + def build_mlp_base(self): + self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim) + self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) + super().build_mlp_base() + self.mlp_transient = MLP( + in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(), + out_dim=self.layer_width // 2, + num_layers=4, + layer_width=self.layer_width, + activation=nn.ReLU(), + ) + + def build_mlp_rgb(self): + self.mlp_rgb = MLP( + in_dim=self.mlp_base.get_out_dim() + + self.encoding_dir.get_out_dim() + + self.embedding_appearance.get_out_dim(), + out_dim=self.layer_width // 2, + num_layers=1, + activation=nn.ReLU(), + ) + + def build_heads(self): + super().build_heads() + self.field_output_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.get_out_dim()) + self.field_output_transient_rgb = TransientRGBHead(in_dim=self.mlp_transient.get_out_dim()) + self.field_output_transient_density = TransientDensityHead(in_dim=self.mlp_transient.get_out_dim()) + + def forward(self, ray_samples: RaySamples): + positions = ray_samples.positions + directions = ray_samples.directions + camera_indices = ray_samples.get_camera_indices() + encoded_xyz = self.encoding_xyz(positions) + encoded_dir = self.encoding_dir(directions) + embedded_appearance = self.embedding_appearance(camera_indices) + embedded_transient = self.embedding_transient(camera_indices) + + base_mlp_out = self.mlp_base(encoded_xyz) + rgb_mlp_out = self.mlp_rgb(torch.cat([base_mlp_out, encoded_dir, embedded_appearance], dim=-1)) + transient_mlp_out = self.mlp_transient(torch.cat([base_mlp_out, embedded_transient], dim=-1)) + + field_rgb_output = self.field_output_rgb(rgb_mlp_out) + field_density_out = self.field_output_density(base_mlp_out) + field_transient_uncertainty_out = self.field_output_transient_uncertainty(transient_mlp_out) + field_transient_rgb_out = self.field_output_transient_rgb(transient_mlp_out) + field_transient_density_out = self.field_output_transient_density(transient_mlp_out) + + field_outputs = {} + field_outputs.update(field_rgb_output) + field_outputs.update(field_density_out) + field_outputs.update(field_transient_uncertainty_out) + field_outputs.update(field_transient_rgb_out) + field_outputs.update(field_transient_density_out) + return field_outputs class NerfWGraph(NeRFGraph): """NeRF-W graph""" def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None: + self.num_images = len(intrinsics) + self.appearance_embedding_dim = 48 + self.transient_embedding_dim = 16 super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + self.renderer_uncertainty = UncertaintyRenderer() + + def populate_fields(self): + """Set the fields.""" + self.field_coarse = NeRFField() + self.field_fine = NerfWField( + num_images=self.num_images, + appearance_embedding_dim=self.appearance_embedding_dim, + transient_embedding_dim=self.transient_embedding_dim, + ) def get_outputs(self, ray_bundle: RayBundle): - raise NotImplementedError + # coarse network + uniform_ray_samples = self.sampler_uniform(ray_bundle) + coarse_field_outputs = self.field_coarse(uniform_ray_samples) + coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) + rgb_coarse = self.renderer_rgb( + rgb=coarse_field_outputs[FieldHeadNames.RGB], + weights=coarse_weights, + ).rgb + depth_coarse = self.renderer_depth(coarse_weights, uniform_ray_samples.ts).depth + + # fine network + pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) + fine_field_outputs = self.field_fine(pdf_ray_samples) + fine_weights = pdf_ray_samples.get_weights( + fine_field_outputs[FieldHeadNames.DENSITY] + fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY] + ) + fine_weights_static = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) + fine_weights_transient = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY]) + + rgb_fine_static_component = self.renderer_rgb( + rgb=fine_field_outputs[FieldHeadNames.RGB], + weights=fine_weights, + ).rgb + rgb_fine_transient_component = self.renderer_rgb( + rgb=fine_field_outputs[FieldHeadNames.TRANSIENT_RGB], + weights=fine_weights, + ).rgb + rgb_fine = rgb_fine_static_component + rgb_fine_transient_component + + rgb_fine_static = self.renderer_rgb( + rgb=fine_field_outputs[FieldHeadNames.RGB], + weights=fine_weights_static, + ).rgb + + # fine_renderer_accumulation_static = self.renderer_accumulation(fine_weights_static) + depth_fine = self.renderer_depth(fine_weights, pdf_ray_samples.ts).depth + depth_fine_static = self.renderer_depth(fine_weights_static, pdf_ray_samples.ts).depth + uncertainty = self.renderer_uncertainty( + fine_field_outputs[FieldHeadNames.UNCERTAINTY], fine_weights_transient + ).uncertainty + + density_transient = fine_field_outputs[FieldHeadNames.TRANSIENT_DENSITY] + + # outputs: + outputs = { + "rgb_coarse": rgb_coarse, # (num_rays, 3) + "rgb_fine": rgb_fine, + "rgb_fine_static": rgb_fine_static, + "depth_coarse": depth_coarse, + "depth_fine": depth_fine, + "depth_fine_static": depth_fine_static, + "density_transient": density_transient, # (num_rays, num_samples, 1) + "uncertainty": uncertainty, # (num_rays, 1) + } + return outputs def get_loss_dict(self, outputs, batch): - raise NotImplementedError + device = outputs["rgb_coarse"].device + pixels = batch["pixels"].to(device) + rgb_coarse = outputs["rgb_coarse"] + rgb_fine = outputs["rgb_fine"] + density_transient = outputs["density_transient"] + betas = outputs["uncertainty"] + rgb_loss_coarse = 0.5 * ((pixels - rgb_coarse) ** 2).sum(-1).mean() + rgb_loss_fine = 0.5 * (((pixels - rgb_fine) ** 2).sum(-1) / (betas[..., 0] ** 2)).mean() + uncertainty_loss = 0.5 * (3 + torch.log(betas)).mean() + density_loss = density_transient.mean() + + loss_dict = { + "rgb_loss_coarse": rgb_loss_coarse, + "rgb_loss_fine": rgb_loss_fine, + "uncertainty_loss": uncertainty_loss, + "density_loss": density_loss, + } + loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) + return loss_dict + + def log_test_image_outputs(self, image_idx, step, image, outputs): + rgb_coarse = outputs["rgb_coarse"] + rgb_fine = outputs["rgb_fine"] + rgb_fine_static = outputs["rgb_fine_static"] + depth_coarse = outputs["depth_coarse"] + depth_fine = outputs["depth_fine"] + depth_fine_static = outputs["depth_fine_static"] + uncertainty = outputs["uncertainty"] + + depth_coarse = visualization.apply_depth_colormap(depth_coarse) + depth_fine = visualization.apply_depth_colormap(depth_fine) + depth_fine_static = visualization.apply_depth_colormap(depth_fine_static) + uncertainty = visualization.apply_depth_colormap(uncertainty) + + row0 = torch.cat([image, uncertainty, torch.ones_like(rgb_fine)], dim=-2) + row1 = torch.cat([rgb_fine, rgb_fine_static, rgb_coarse], dim=-2) + row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) + combined_image = torch.cat([row0, row1, row2], dim=-3) + + writer.write_event( + {"name": f"image_idx_{image_idx}-nerfw", "x": combined_image, "step": step, "group": "val_img"} + ) diff --git a/mattport/nerf/graph/semantic_nerf.py b/mattport/nerf/graph/semantic_nerf.py index 91e6779c8..9b32c6376 100644 --- a/mattport/nerf/graph/semantic_nerf.py +++ b/mattport/nerf/graph/semantic_nerf.py @@ -1,4 +1,160 @@ """ Semantic NeRF implementation. -TODO: """ + +from typing import Tuple + +import torch +from torch import nn + +from mattport.nerf.field_modules.field_heads import FieldHeadNames, SemanticStuffHead +from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph +from mattport.nerf.renderers import SemanticRenderer +from mattport.structures.rays import RayBundle, RaySamples +from mattport.utils import visualization, writer + + +class SemanticNerfField(NeRFField): + """Semantic-NeRF field""" + + def __init__( + self, num_layers=8, layer_width=256, skip_connections: Tuple = (4,), num_stuff_classes: int = None + ) -> None: + assert num_stuff_classes is not None + self.num_stuff_classes = num_stuff_classes + super().__init__(num_layers=num_layers, layer_width=layer_width, skip_connections=skip_connections) + self.semantic_mlp = MLP( + in_dim=self.mlp_base.get_out_dim(), + out_dim=self.layer_width // 2, + num_layers=1, + activation=nn.ReLU(), + ) + self.field_output_semantic = SemanticStuffHead( + in_dim=self.semantic_mlp.get_out_dim(), num_classes=num_stuff_classes + ) + self.semantic_sequential = torch.nn.Sequential(self.semantic_mlp, self.field_output_semantic) + + def forward(self, ray_samples: RaySamples): + """Evaluates the field at points along the ray.""" + positions = ray_samples.positions + directions = ray_samples.directions + encoded_xyz = self.encoding_xyz(positions) + encoded_dir = self.encoding_dir(directions) + base_mlp_out = self.mlp_base(encoded_xyz) + rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1)) + + field_rgb_output = self.field_output_rgb(rgb_mlp_out) + field_density_out = self.field_output_density(base_mlp_out) + field_semantic_output = self.semantic_sequential(base_mlp_out) + + field_outputs = {} + field_outputs.update(field_rgb_output) + field_outputs.update(field_density_out) + field_outputs.update(field_semantic_output) + return field_outputs + + +class SemanticNerfGraph(NeRFGraph): + """Semantic-NeRF graph""" + + def __init__(self, intrinsics=None, camera_to_world=None, stuff_classes: int = None, **kwargs) -> None: + assert stuff_classes is not None + self.stuff_classes = stuff_classes + super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + self.renderer_semantic = SemanticRenderer() + self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="mean") + + def populate_fields(self): + """Set the fields.""" + self.field_coarse = NeRFField() + self.field_fine = SemanticNerfField(num_stuff_classes=len(self.stuff_classes)) + + def get_outputs(self, ray_bundle: RayBundle): + # coarse network: + uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples + + coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs + + coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) + + coarse_renderer_outputs = self.renderer_rgb( + rgb=coarse_field_outputs[FieldHeadNames.RGB], + weights=coarse_weights, + ) # RendererOutputs + coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights) # RendererOutputs + coarse_renderer_depth = self.renderer_depth(coarse_weights, uniform_ray_samples.ts) + + # fine network: + pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples + fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs + + fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) + + fine_renderer_outputs = self.renderer_rgb( + rgb=fine_field_outputs[FieldHeadNames.RGB], + weights=fine_weights, + ) # RendererOutputs + fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs + fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts) + + # TODO refactor this into "vis" section. Doesn't need to be run during training. + coarse_renderer_depth = visualization.apply_depth_colormap( + coarse_renderer_depth.depth, + accumulation=coarse_renderer_accumulation.accumulation, + near_plane=self.near_plane, + far_plane=self.far_plane, + ) + fine_renderer_depth = visualization.apply_depth_colormap( + fine_renderer_depth.depth, + accumulation=fine_renderer_accumulation.accumulation, + near_plane=self.near_plane, + far_plane=self.far_plane, + ) + coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation) + fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation) + + semantic_logits = self.renderer_semantic( + fine_field_outputs[FieldHeadNames.SEMANTICS_STUFF], weights=fine_weights + ).semantics + + # outputs: + outputs = { + "rgb_coarse": coarse_renderer_outputs.rgb, + "rgb_fine": fine_renderer_outputs.rgb, + "accumulation_coarse": coarse_renderer_accumulation, + "accumulation_fine": fine_renderer_accumulation, + "depth_coarse": coarse_renderer_depth, + "depth_fine": fine_renderer_depth, + "semantic_logits": semantic_logits, + } + return outputs + + def get_loss_dict(self, outputs, batch): + # TODO(ethan): batch has "stuff_image". use it + device = outputs["rgb_coarse"].device + pixels = batch["pixels"].to(device) + rgb_loss_coarse = self.rgb_loss(pixels, outputs["rgb_coarse"]) + rgb_loss_fine = self.rgb_loss(pixels, outputs["rgb_fine"]) + + semantic_logits = outputs["semantic_logits"] + semantic_classes = batch["stuff_image"][..., 0].to(device).long() + semantic_loss_fine = self.cross_entropy_loss(semantic_logits, semantic_classes) + loss_dict = { + "rgb_loss_coarse": rgb_loss_coarse, + "rgb_loss_fine": rgb_loss_fine, + "semantic_loss_fine": semantic_loss_fine, + } + + loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) + return loss_dict + + def log_test_image_outputs(self, image_idx, step, image, outputs): + super().log_test_image_outputs(image_idx, step, image, outputs) + semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantic_logits"], dim=-1), dim=-1) / len( + self.stuff_classes + ) + semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None]) + writer.write_event( + {"name": f"semantic_image_idx_{image_idx}", "x": semantic_labels_image, "step": step, "group": "val_sem"} + ) diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index b35d737f4..9e841dcdc 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -32,12 +32,22 @@ class NeRFField(nn.Module): self.layer_width = layer_width self.skip_connections = skip_connections + self.build_encodings() + self.build_mlp_base() + self.build_mlp_rgb() + self.build_heads() + + def build_encodings(self): + """Build the encodings.""" self.encoding_xyz = NeRFEncoding( in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True ) self.encoding_dir = NeRFEncoding( in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True ) + + def build_mlp_base(self): + """Build the MLP base.""" self.mlp_base = MLP( in_dim=self.encoding_xyz.get_out_dim(), out_dim=self.layer_width, @@ -46,6 +56,9 @@ class NeRFField(nn.Module): skip_connections=self.skip_connections, activation=nn.ReLU(), ) + + def build_mlp_rgb(self): + """Build the MLP for RGB.""" self.mlp_rgb = MLP( in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(), out_dim=self.layer_width // 2, @@ -53,16 +66,14 @@ class NeRFField(nn.Module): layer_width=self.layer_width // 2, activation=nn.ReLU(), ) + + def build_heads(self): + """Build the heads.""" self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) def forward(self, ray_samples: RaySamples): - """Evaluates the field at points along the ray - Args: - xyz: () - # TODO(ethan): change the input to be something more abstracted - e.g., a FieldInput structure - """ + """Evaluates the field at points along the ray.""" positions = ray_samples.positions directions = ray_samples.directions encoded_xyz = self.encoding_xyz(positions) @@ -144,9 +155,7 @@ class NeRFGraph(Graph): # coarse network: uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples - # time coarse network # xyz -> asd coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs - # time end coarse network coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index f746b5942..8b5474872 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -17,6 +17,8 @@ class RendererOutputs: density: TensorType["num_rays", 1] = None accumulation: TensorType["num_rays", 1] = None depth: TensorType["num_rays", 1] = None + uncertainty: TensorType["num_rays", 1] = None + semantics: TensorType["num_rays", "num_classes"] = None class RGBRenderer(nn.Module): @@ -116,3 +118,37 @@ class DepthRenderer(nn.Module): return RendererOutputs(depth=depth[..., None]) raise NotImplementedError(f"Method {self.method} not implemented") + + +class UncertaintyRenderer(nn.Module): + """Calculate uncertainty along the ray.""" + + @classmethod + def forward( + cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples"] + ) -> RendererOutputs: + """_summary_ + + Args: + betas (TensorType[..., "num_samples", 1]): _description_ + weights (TensorType[..., "num_samples"]): _description_ + + Returns: + RendererOutputs: _description_ + """ + uncertainty = torch.sum(weights[..., None] * betas, dim=-2) + renderer_outputs = RendererOutputs(uncertainty=uncertainty) + return renderer_outputs + + +class SemanticRenderer(nn.Module): + """Calculate semantics along the ray.""" + + @classmethod + def forward( + cls, semantics: TensorType[..., "num_samples", "num_classes"], weights: TensorType[..., "num_samples"] + ) -> RendererOutputs: + """_summary_""" + sem = torch.sum(weights[..., None] * semantics, dim=-2) + renderer_outputs = RendererOutputs(semantics=sem) + return renderer_outputs diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index cf2fa468c..0d9bd9c76 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -88,9 +88,9 @@ class Trainer: self.train_image_dataset = ImageDataset( image_filenames=dataset_inputs_dict["train"].image_filenames, downscale_factor=dataset_inputs_dict["train"].downscale_factor, + semantics=dataset_inputs_dict["train"].semantics, alpha_color=dataset_inputs_dict["train"].alpha_color, ) - self.train_dataset = CollateIterDataset( self.train_image_dataset, collate_fn=self.collate_fn, @@ -117,6 +117,7 @@ class Trainer: self.config.graph.network, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world, + stuff_classes=dataset_inputs.semantics.stuff_classes, ).to(self.device) self.setup_optimizers() # NOTE(ethan): can this be before DDP? @@ -129,7 +130,6 @@ class Trainer: def setup_optimizers(self): """_summary_""" - # TODO(ethan): handle different world sizes self.optimizers = Optimizers(self.config.graph.param_groups, self.graph.get_param_groups()) def load_checkpoint(self, load_config: DictConfig) -> int: @@ -222,7 +222,7 @@ class Trainer: # TODO: add the learning rates to tensorboard/logging if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: self.save_checkpoint(self.config.graph.model_dir, step) - if step % self.config.graph.steps_per_test == 0: + if step % self.config.graph.steps_per_test == 0: # NOTE(ethan): we should still run this in dry-run mode! for image_idx in self.config.data.validation_image_indices: _ = self.test_image(image_idx=image_idx, step=step) stats_tracker.print_stats(i / num_iterations) @@ -251,7 +251,11 @@ class Trainer: self.graph.eval() intrinsics = self.val_image_intrinsics[image_idx] camera_to_world = self.val_image_camera_to_world[image_idx] - outputs = self.graph.get_outputs_for_camera(intrinsics, camera_to_world) + chunk_size = 1024 + training_camera_index = image_idx # TODO(ethan): change this because training and test should not be the same + outputs = self.graph.get_outputs_for_camera( + intrinsics, camera_to_world, chunk_size=chunk_size, training_camera_index=training_camera_index + ) image = self.val_image_dataset[image_idx]["image"].to(self.device) psnr = self.graph.log_test_image_outputs(image_idx, step, image, outputs) self.graph.train() diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 753b889cc..d567cc2f2 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -22,18 +22,13 @@ class RayBundle: camera_indices: Optional[TensorType["num_rays"]] = None def to_camera_ray_bundle(self, image_height, image_width) -> "CameraRayBundle": - """_summary_ - - Args: - image_height (_type_): _description_ - image_width (_type_): _description_ - - Returns: - CameraRayBundle: _description_ - """ + """Returns a CameraRayBundle from this object.""" camera_ray_bundle = CameraRayBundle( origins=self.origins.view(image_height, image_width, 3), directions=self.directions.view(image_height, image_width, 3), + camera_indices=self.camera_indices.view(image_height, image_width) + if not isinstance(self.camera_indices, type(None)) + else None, ) return camera_ray_bundle @@ -77,7 +72,12 @@ class CameraRayBundle: def __post_init__(self): if not isinstance(self.camera_index, type(None)): - self.camera_indices = torch.ones_like(self.origins[:, :]) * self.camera_index + self.set_camera_indices(self.camera_index) + + def set_camera_indices(self, camera_index: int): + """Sets the camera indices for a specific camera index.""" + self.camera_index = camera_index + self.camera_indices = torch.ones_like(self.origins[..., 0]).long() * camera_index def get_num_rays(self): """Return the number of rays in this bundle.""" @@ -98,7 +98,7 @@ class CameraRayBundle: def get_row_major_sliced_ray_bundle(self, start_idx, end_idx): """Return a RayBundle""" camera_indices = ( - self.camera_indices.view(-1)[start_idx:end_idx] if not isinstance(self.camera_index, type(None)) else None + self.camera_indices.view(-1)[start_idx:end_idx] if not isinstance(self.camera_indices, type(None)) else None ) return RayBundle( origins=self.origins.view(-1, 3)[start_idx:end_idx], @@ -121,6 +121,12 @@ class RaySamples: self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1) self.deltas = self.get_deltas() + def get_camera_indices(self): + """Returns camera indices.""" + assert not isinstance(self.ray_bundle.camera_indices, type(None)), "Camera indices cannot be None here." + camera_indices = self.ray_bundle.camera_indices.unsqueeze(1).repeat(1, self.positions.shape[1]) + return camera_indices + def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]: """Returns positions.""" return ray_bundle.origins[:, None] + self.ts[:, :, None] * ray_bundle.directions[:, None] @@ -128,7 +134,8 @@ class RaySamples: def get_deltas(self) -> TensorType[..., "num_samples"]: """Returns deltas.""" dists = self.ts[..., 1:] - self.ts[..., :-1] - dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] + dist_inf = 1e10 * torch.ones_like(dists[..., -1:]) + dists = torch.cat([dists, dist_inf], -1) # [N_rays, N_samples] deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1) return deltas diff --git a/scripts/run_train.py b/scripts/run_train.py index 654936a3d..5c05e28a8 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -6,6 +6,7 @@ import logging import random import socket from datetime import timedelta +import traceback from typing import Any, Callable import hydra @@ -153,7 +154,7 @@ def launch( try: main_func(local_rank=0, world_size=1, config=config) except KeyboardInterrupt: - pass + print(traceback.format_exc()) finally: profiler.flush_profiler(config.logging) elif world_size > 1: diff --git a/tests/nerf/field_modules/test_embedding.py b/tests/nerf/field_modules/test_embedding.py new file mode 100644 index 000000000..43b5340f6 --- /dev/null +++ b/tests/nerf/field_modules/test_embedding.py @@ -0,0 +1,14 @@ +""" +Embedding tests +""" +from mattport.nerf.field_modules.embedding import Embedding + + +def test_indexing(): + """Test embedding indexing""" + in_dim = 100 + out_dim = 64 + + embedding = Embedding(in_dim, out_dim) + assert embedding + # TODO