diff --git a/README.md b/README.md index e0bc7e199..b871711bc 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ We provide many logging functionalities for timing and/or tracking losses during 1. **Writer**: Logs losses and generated images during training to a specified output stream. Specify the type of writer (Tensorboard, Local Writer, Weights and Biases), and how often to log in the config. -3. **Profiler**: Computes the average total time of execution for any function with the `@profiler.time_function` decorator. Prints out the full profile at the termination of training or the program. +2. **Profiler**: Computes the average total time of execution for any function with the `@profiler.time_function` decorator. Prints out the full profile at the termination of training or the program. # Setting up Jupyter @@ -119,3 +119,29 @@ Modify the `BENCH` variable to specify which jobs ("ckpt_dir") and methods ("met ``` python scripts/benchmarking/run_benchmark.py ``` + +# Running other repos with our data + +``` +# nerf-pytorch +cd external +python run_nerf.py --config configs/chair.txt --datadir /path/to/mattport/data/blender/chair + +# jaxnerf +cd external +conda activate jaxnerf +python -m jaxnerf.train --data_dir=/path/to/mattport/data/blender/chair --train_dir=/path/to/mattport/outputs/blender_chair_jaxnerf --config=/path/to/mattport/external/jaxnerf/configs/demo --render_every 100 +``` + +# Speeding up the code + +Documentation for running the code with CUDA. + +``` +This is how to install tiny-cuda-nn. +# https://github.com/NVlabs/tiny-cuda-nn +pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch + +To run instant-ngp with tcnn, you can do the following. This is with the fox dataset. +python scripts/run_train.py --config-name=instant_ngp_tcnn.yaml data/dataset=instant_ngp_fox +``` \ No newline at end of file diff --git a/configs/data/dataset/instant_ngp_fox.yaml b/configs/data/dataset/instant_ngp_fox.yaml new file mode 100644 index 000000000..90674b3d6 --- /dev/null +++ b/configs/data/dataset/instant_ngp_fox.yaml @@ -0,0 +1,3 @@ +data_directory: data/instant_ngp/fox +dataset_format: instant_ngp +downscale_factor: 1 diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index eba1a6036..a547545c5 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -3,8 +3,9 @@ defaults: - dataset: blender_lego # validation options -val_image_indices: [0, 10, 20] # Which images to evaluate on. +val_image_indices: [0, 10, 20] # which images to evaluate on val_downscale_factor: 8 +val_num_rays_per_chunk: 4096 # dataloader options dataloader: @@ -12,4 +13,4 @@ dataloader: shuffle: true num_images_to_sample_from: 8 num_rays_per_batch: 1024 - num_times_to_repeat_images: 40 + num_times_to_repeat_images: 100 diff --git a/configs/default.yaml b/configs/default.yaml index 192a5ac79..c71b434d7 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -7,6 +7,8 @@ defaults: - override hydra/hydra_logging: colorlog - _self_ +seed: 0 + hydra: run: # Configure output dir of each experiment programmatically from the arguments diff --git a/configs/graph/instant_ngp_graph.yaml b/configs/graph/instant_ngp_graph.yaml index cda7fa4a9..5d0a72f14 100644 --- a/configs/graph/instant_ngp_graph.yaml +++ b/configs/graph/instant_ngp_graph.yaml @@ -4,6 +4,9 @@ defaults: network: _target_: mattport.nerf.graph.instant_ngp.NGPGraph + field_implementation: torch # torch, tcnn, ... + loss_coefficients: # loss coefficients + rgb_loss: 1.0 # optimizer options for network param_groups: diff --git a/configs/instant_ngp.yaml b/configs/instant_ngp.yaml index 6a0617b63..b9c0262df 100644 --- a/configs/instant_ngp.yaml +++ b/configs/instant_ngp.yaml @@ -3,3 +3,5 @@ defaults: - override graph: instant_ngp_graph - _self_ method_name: instant_ngp +data.dataloader.num_rays_per_batch: 4096 +data.val_num_rays_per_chunk: 16384 \ No newline at end of file diff --git a/configs/instant_ngp_tcnn.yaml b/configs/instant_ngp_tcnn.yaml new file mode 100644 index 000000000..e5778354c --- /dev/null +++ b/configs/instant_ngp_tcnn.yaml @@ -0,0 +1,8 @@ +defaults: + - default + - override graph: instant_ngp_graph + - _self_ +method_name: instant_ngp +data.dataloader.num_rays_per_batch: 4096 +data.val_num_rays_per_chunk: 16384 +graph.network.field_implementation: tcnn \ No newline at end of file diff --git a/mattport/nerf/dataset/collate.py b/mattport/nerf/dataset/collate.py index 8e241d1e4..d200dbeb4 100644 --- a/mattport/nerf/dataset/collate.py +++ b/mattport/nerf/dataset/collate.py @@ -31,15 +31,15 @@ class CollateIterDataset(torch.utils.data.IterableDataset): dataset: torch.utils.data.Dataset, collate_fn: Callable, num_samples_to_collate: int, - num_times_to_repeat: int = 0, + num_times_to_repeat_images: int = 0, ): super().__init__() self.dataset = dataset self.collate_fn = collate_fn self.num_samples_to_collate = num_samples_to_collate - self.num_times_to_repeat = num_times_to_repeat + self.num_times_to_repeat_images = num_times_to_repeat_images - self.num_repeated = self.num_times_to_repeat # starting value + self.num_repeated = self.num_times_to_repeat_images # starting value self.cached_batch_list = None def __getitem__(self, idx): @@ -54,12 +54,12 @@ class CollateIterDataset(torch.utils.data.IterableDataset): def __iter__(self): while True: - if self.num_repeated >= self.num_times_to_repeat: + if self.num_repeated >= self.num_times_to_repeat_images: # trigger a reset self.num_repeated = 0 batch_list = self.get_batch_list() # possibly save a cached item - self.cached_batch_list = batch_list if self.num_times_to_repeat != 0 else None + self.cached_batch_list = batch_list if self.num_times_to_repeat_images != 0 else None else: batch_list = self.cached_batch_list self.num_repeated += 1 diff --git a/mattport/nerf/dataset/format/instant_ngp.py b/mattport/nerf/dataset/format/instant_ngp.py index 2440ed7ae..a00126956 100644 --- a/mattport/nerf/dataset/format/instant_ngp.py +++ b/mattport/nerf/dataset/format/instant_ngp.py @@ -2,3 +2,79 @@ Code for loading instant-ngp dataset formats. # TODO """ + +import logging +import os + +import imageio +import numpy as np +import torch + +from mattport.nerf.dataset.structs import DatasetInputs, SceneBounds +from mattport.utils.io import load_from_json + + +def load_instant_ngp_data( + basedir: str, downscale_factor: int = 1, split: str = "train", camera_translation_scalar=0.33 +) -> DatasetInputs: + """Processes the a blender dataset directory. + Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. + + Args: + data_directory (str): Location of data + downscale_factor (int, optional): How much to downscale images. Defaults to 1.0. + split (str, optional): Which dataset split to generate. + + Returns: + DatasetInputs + """ + # pylint: disable=unused-argument + # pylint: disable-msg=too-many-locals + meta = load_from_json(os.path.join(basedir, "transforms.json")) + image_filenames = [] + poses = [] + num_skipped_image_filenames = 0 + for frame in meta["frames"]: + fname = os.path.join(basedir, frame["file_path"]) + if not os.path.exists(fname): + num_skipped_image_filenames += 1 + else: + image_filenames.append(fname) + poses.append(np.array(frame["transform_matrix"])) + if num_skipped_image_filenames >= 0: + logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split) + poses = np.array(poses).astype(np.float32) + poses[:3, 3] *= camera_translation_scalar + + img_0 = imageio.imread(image_filenames[0]) + image_height, image_width = img_0.shape[:2] + camera_angle_x = float(meta["camera_angle_x"]) + focal_length = 0.5 * image_width / np.tan(0.5 * camera_angle_x) + + cx = image_width / 2.0 + cy = image_height / 2.0 + camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform + num_cameras = len(image_filenames) + num_intrinsics_params = 3 + intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) + intrinsics *= torch.tensor([cx, cy, focal_length]) + + # in x,y,z order + # assumes that the scene is centered at the origin + aabb_scale = meta["aabb_scale"] + scene_bounds = SceneBounds( + aabb=torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ) + ) + + # TODO(ethan): add alpha background color + dataset_inputs = DatasetInputs( + image_filenames=image_filenames, + downscale_factor=downscale_factor, + intrinsics=intrinsics * 1.0 / downscale_factor, # downscaling the intrinsics here + camera_to_world=camera_to_world, + scene_bounds=scene_bounds, + ) + + return dataset_inputs diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 1888626a4..e2aef9989 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -8,6 +8,7 @@ from omegaconf import ListConfig from mattport.nerf.dataset.format.blender import load_blender_data from mattport.nerf.dataset.format.friends import load_friends_data +from mattport.nerf.dataset.format.instant_ngp import load_instant_ngp_data from mattport.nerf.dataset.structs import DatasetInputs from mattport.structures.colors import get_color from mattport.utils.io import get_absolute_path @@ -57,6 +58,12 @@ def get_dataset_inputs( split="train", include_point_cloud=False, ) + elif dataset_format == "instant_ngp": + dataset_inputs = load_instant_ngp_data( + get_absolute_path(data_directory), + downscale_factor=downscale_factor, + split="train", + ) else: raise NotImplementedError(f"{dataset_format} is not a valid dataset type") diff --git a/mattport/nerf/fields/instant_ngp_field.py b/mattport/nerf/fields/instant_ngp_field.py new file mode 100644 index 000000000..deae732b2 --- /dev/null +++ b/mattport/nerf/fields/instant_ngp_field.py @@ -0,0 +1,151 @@ +""" +Instant-NGP field implementations using tiny-cuda-nn, torch, .... +""" + + +from typing import Tuple + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from mattport.nerf.field_modules.encoding import Encoding, HashEncoding, SHEncoding +from mattport.nerf.fields.base import Field +from mattport.nerf.fields.nerf_field import NeRFField +from mattport.structures.rays import PointSamples +from mattport.utils.misc import is_not_none + +try: + import tinycudann as tcnn +except ImportError as e: + # tinycudann module doesn't exist + pass + + +def get_normalized_positions(positions, aabb): + """Return normalized positions in range [0, 1] based on the aabb axis-aligned bounding box.""" + aabb_lengths = aabb[1] - aabb[0] + positions = (positions - aabb[0]) / aabb_lengths + return positions + + +class TCNNInstantNGPField(Field): + """NeRF Field""" + + def __init__( + self, aabb, num_layers=2, hidden_dim=64, geo_feat_dim=15, num_layers_color=3, hidden_dim_color=64 + ) -> None: + super().__init__() + + # TODO: make this a parameter that is put to the correct device + self.aabb = Parameter(aabb, requires_grad=False) + + self.geo_feat_dim = geo_feat_dim + + # TODO: set this properly based on the aabb + per_level_scale = 1.4472692012786865 + + self.position_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "HashGrid", + "n_levels": 16, + "n_features_per_level": 2, + "log2_hashmap_size": 19, + "base_resolution": 16, + "per_level_scale": per_level_scale, + }, + ) + + self.direction_encoding = tcnn.Encoding( + n_input_dims=3, + encoding_config={ + "otype": "SphericalHarmonics", + "degree": 4, + }, + ) + + self.mlp_base = tcnn.Network( + n_input_dims=32, + n_output_dims=1 + self.geo_feat_dim, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim, + "n_hidden_layers": num_layers - 1, + }, + ) + + self.mlp_head = tcnn.Network( + n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim, + n_output_dims=3, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": hidden_dim_color, + "n_hidden_layers": num_layers_color - 1, + }, + ) + + def get_density(self, point_samples: PointSamples): + """Computes and returns the densities.""" + # TODO: add valid_mask masking! + # remap positions to range 0 to 1 based on the aabb + positions = get_normalized_positions(point_samples.positions, self.aabb) + positions_flat = positions.view(-1, 3) + dtype = positions_flat.dtype + x = self.position_encoding(positions_flat) + h = self.mlp_base(x).view(*point_samples.positions.shape[:-1], -1).to(dtype) + density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) + density = F.softplus(density_before_activation) + return density, base_mlp_out + + def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): + # TODO: add valid_mask masking! + assert is_not_none(density_embedding) + directions = point_samples.directions + directions_flat = directions.view(-1, 3) + dtype = directions_flat.dtype + d = self.direction_encoding(directions_flat) + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + h = self.mlp_head(h).view(*point_samples.directions.shape[:-1], -1).to(dtype) + rgb = torch.sigmoid(h) + return {"rgb": rgb} + + +class TorchInstantNGPField(NeRFField): + """ + PyTorch implementation of the instant-ngp field. + """ + + def __init__( + self, + aabb, + position_encoding: Encoding = HashEncoding(), + direction_encoding: Encoding = SHEncoding(), + base_mlp_num_layers: int = 3, + base_mlp_layer_width: int = 64, + head_mlp_num_layers: int = 2, + head_mlp_layer_width: int = 32, + skip_connections: Tuple = (4,), + ) -> None: + super().__init__( + position_encoding, + direction_encoding, + base_mlp_num_layers, + base_mlp_layer_width, + head_mlp_num_layers, + head_mlp_layer_width, + skip_connections, + ) + self.aabb = Parameter(aabb, requires_grad=False) + + def get_density(self, point_samples: PointSamples): + normalized_point_samples = point_samples + normalized_point_samples.positions = get_normalized_positions(normalized_point_samples.positions, self.aabb) + return super().get_density(normalized_point_samples) + + +field_implementation_to_class = {"tcnn": TCNNInstantNGPField, "torch": TorchInstantNGPField} diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py index 908d8df4b..783f06de1 100644 --- a/mattport/nerf/graph/instant_ngp.py +++ b/mattport/nerf/graph/instant_ngp.py @@ -2,35 +2,151 @@ Implementation of Instant NGP. """ -from mattport.nerf.field_modules.encoding import HashEncoding, SHEncoding -from mattport.nerf.graph.vanilla_nerf import NeRFField, NeRFGraph +from typing import Dict, List + +import torch +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + +from mattport.nerf.colliders import AABBBoxCollider +from mattport.nerf.field_modules.field_heads import FieldHeadNames +from mattport.nerf.fields.instant_ngp_field import field_implementation_to_class +from mattport.nerf.graph.base import Graph +from mattport.nerf.loss import MSELoss +from mattport.nerf.occupancy_grid import OccupancyGrid +from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer +from mattport.nerf.sampler import PDFSampler, UniformSampler +from mattport.structures import colors +from mattport.structures.rays import RayBundle +from mattport.utils import visualization, writer -class NGPGraph(NeRFGraph): +class NGPGraph(Graph): + """NeRF-W graph""" - def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None: + def __init__(self, field_implementation="torch", intrinsics=None, camera_to_world=None, **kwargs) -> None: + assert field_implementation in field_implementation_to_class + self.field_implementation = field_implementation + self.field = None super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + def populate_collider(self): + self.collider = AABBBoxCollider(self.scene_bounds) + def populate_fields(self): """Set the fields.""" + # torch or tiny-cuda-nn version + self.field = field_implementation_to_class[self.field_implementation](self.scene_bounds.aabb) - position_encoding = HashEncoding() - direction_encoding = SHEncoding() + def populate_misc_modules(self): + # occupancy grid + self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb) - self.field_coarse = NeRFField( - position_encoding=position_encoding, - direction_encoding=direction_encoding, - base_mlp_num_layers=3, - base_mlp_layer_width=64, - head_mlp_num_layers=2, - head_mlp_layer_width=32, + # samplers + self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid) + # NOTE(ethan): are we sure we want the include_original flag used like this? + # it could be easily forgotten that it's by default True...? + self.sampler_pdf = PDFSampler(num_samples=128, include_original=False) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = StructuralSimilarityIndexMeasure() + self.lpips = LearnedPerceptualImagePatchSimilarity() + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + param_groups = {} + param_groups["fields"] = list(self.field.parameters()) + return param_groups + + def get_outputs(self, ray_bundle: RayBundle): + + # uniform sampling + ray_samples_uniform = self.sampler_occupancy_grid(ray_bundle) + field_outputs_uniform = self.field.forward(ray_samples_uniform.to_point_samples()) + weights_uniform = ray_samples_uniform.get_weights(field_outputs_uniform[FieldHeadNames.DENSITY]) + + # pdf sampling + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_uniform) + field_outputs_pdf = self.field.forward(ray_samples_pdf.to_point_samples()) + + ts, indices = torch.sort(torch.cat([ray_samples_uniform.ts, ray_samples_pdf.ts], -1), -1) + ray_samples = ray_bundle.get_ray_samples(ts) + field_outputs = {} + for fo_name, _ in field_outputs_pdf.items(): + fo_uniform = field_outputs_uniform[fo_name] + fo_pdf = field_outputs_pdf[fo_name] + fo_uniform_pdf = torch.cat([fo_uniform, fo_pdf], 1) + index = indices.view(fo_uniform_pdf[..., :1].shape) + index = index.expand(-1, -1, fo_uniform_pdf.shape[-1]) # TODO: don't hardcode this + field_outputs[fo_name] = torch.gather(fo_uniform_pdf, dim=1, index=index) + + weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) + rgb = self.renderer_rgb( + rgb=field_outputs[FieldHeadNames.RGB], + weights=weights, ) - self.field_fine = NeRFField( - position_encoding=position_encoding, - direction_encoding=direction_encoding, - base_mlp_num_layers=3, - base_mlp_layer_width=64, - head_mlp_num_layers=2, - head_mlp_layer_width=32, + accumulation = self.renderer_accumulation(weights) + depth = self.renderer_depth(weights, ray_samples.ts) + + densities_occupancy_grid = self.occupancy_grid.get_densities(ray_samples.positions) + weights_occupancy_grid = ray_samples.get_weights(densities_occupancy_grid) + depth_occupancy_grid = self.renderer_depth(weights_occupancy_grid, ray_samples.ts) + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + "depth_occupancy_grid": depth_occupancy_grid, + } + return outputs + + def get_loss_dict(self, outputs, batch): + device = self.get_device() + pixels = batch["pixels"].to(device) + rgb_loss = self.rgb_loss(pixels, outputs["rgb"]) + loss_dict = {"rgb_loss": rgb_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 = outputs["rgb"] + acc = visualization.apply_colormap(outputs["accumulation"]) + depth = visualization.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], ) + + combined_rgb = torch.cat([image, rgb], dim=1) + combined_acc = torch.cat([acc], dim=1) + combined_depth = torch.cat([depth], dim=1) + + writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") + + depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"]) + writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth_occupancy_grid") + + # Switch images from [H, W, C] to [1, C, H, W] for metrics computations + image = torch.moveaxis(image, -1, 0)[None, ...] + rgb = torch.moveaxis(rgb, -1, 0)[None, ...] + + psnr = self.psnr(image, rgb) + ssim = self.ssim(image, rgb) + lpips = self.lpips(image, rgb) + + writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(psnr), step=step, group="psnr") + writer.put_scalar(name=f"val_{image_idx}", scalar=float(ssim), step=step, group="ssim") + writer.put_scalar(name=f"val_{image_idx}", scalar=float(lpips), step=step, group="lpips") + + return psnr.item() diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 28a13ff58..b77d14a41 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -99,7 +99,7 @@ class Trainer: self.train_image_dataset, collate_fn=self.collate_fn, num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from, - num_times_to_repeat=self.config.data.dataloader.num_times_to_repeat_images, + num_times_to_repeat_images=self.config.data.dataloader.num_times_to_repeat_images, ) self.train_dataloader = DataLoader( self.train_dataset, @@ -276,7 +276,7 @@ class Trainer: self.graph.eval() intrinsics = self.val_image_intrinsics[image_idx] camera_to_world = self.val_image_camera_to_world[image_idx] - chunk_size = 1024 + chunk_size = self.config.data.val_num_rays_per_chunk 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 diff --git a/mattport/utils/visualization.py b/mattport/utils/visualization.py index a8fac0487..aceb8c024 100644 --- a/mattport/utils/visualization.py +++ b/mattport/utils/visualization.py @@ -19,15 +19,12 @@ def apply_colormap(image: TensorType[..., 1], cmap="viridis") -> TensorType[..., colormap = cm.get_cmap(cmap) colormap = torch.tensor(colormap.colors).to(image.device) - - image2 = (image * 255).long() - - if torch.min(image2) < 0 or torch.max(image2) > 255: - print("hi") - - image = (image * 255).long() - - return colormap[image[..., 0]] + image_long = (image * 255).long() + image_long_min = torch.min(image_long) + image_long_max = torch.max(image_long) + assert image_long_min >= 0, f"the min value is {image_long_min}" + assert image_long_max <= 255, f"the max value is {image_long_max}" + return colormap[image_long[..., 0]] def apply_depth_colormap( diff --git a/scripts/run_train.py b/scripts/run_train.py index bdaea1e80..1590c7668 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -100,7 +100,7 @@ def _distributed_worker( assert num_gpus_per_machine <= torch.cuda.device_count() torch.cuda.set_device(local_rank) - _set_random_seed(1234 + global_rank) + _set_random_seed(config.seed + global_rank) comms.synchronize(world_size) output = main_func(local_rank, world_size, config) @@ -149,11 +149,11 @@ def launch( world_size = num_machines * num_gpus_per_machine if world_size == 0: # Using only CPU and one process. - _set_random_seed(0) + _set_random_seed(config.seed) main_func(local_rank=0, world_size=0, config=config) elif world_size == 1: # Using one gpu and one process. - _set_random_seed(0) + _set_random_seed(config.seed) try: main_func(local_rank=0, world_size=1, config=config) except KeyboardInterrupt: diff --git a/tests/nerf/dataloader/test_ray_generator.py b/tests/nerf/dataloader/test_ray_generator.py index 022a169b3..98682960c 100644 --- a/tests/nerf/dataloader/test_ray_generator.py +++ b/tests/nerf/dataloader/test_ray_generator.py @@ -52,7 +52,7 @@ def test_dataloader(visualize=False): image_dataset, collate_fn=lambda batch_list: collate_batch(batch_list, num_rays_per_batch, keep_full_image=True), num_samples_to_collate=num_images_to_sample_from, - num_times_to_repeat=num_times_to_repeat_images, + num_times_to_repeat_images=num_times_to_repeat_images, ) dataloader = DataLoader( iter_dataset, diff --git a/tests/nerf/test_fields.py b/tests/nerf/test_fields.py new file mode 100644 index 000000000..65421d958 --- /dev/null +++ b/tests/nerf/test_fields.py @@ -0,0 +1,42 @@ +""" +Test the fields +""" +import torch + +from mattport.nerf.field_modules.field_heads import FieldHeadNames +from mattport.nerf.fields.instant_ngp_field import TCNNInstantNGPField +from mattport.structures.rays import PointSamples + + +def test_tcnn_instant_ngp_field(): + """Test the tiny-cuda-nn field""" + # pylint: disable=import-outside-toplevel + # pylint: disable=unused-import + try: + import tinycudann as tcnn + except ImportError as e: + # tinycudann module doesn't exist + print(e) + return + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + aabb_scale = 1.0 + aabb = torch.tensor( + [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 + ).to(device) + field = TCNNInstantNGPField(aabb) + num_rays = 1024 + num_samples = 256 + positions = torch.rand((num_rays, num_samples, 3), dtype=torch.float32, device=device) + directions = torch.rand_like(positions) + point_samples = PointSamples(positions=positions, directions=directions) + field_outputs = field.forward(point_samples) + + positions_dtype = positions.dtype + field_outputs_density_dtype = field_outputs[FieldHeadNames.DENSITY].dtype + assert ( + positions_dtype == field_outputs_density_dtype + ), f"The input and output dtypes do not match: {positions_dtype} vs. {field_outputs_density_dtype}" + + +if __name__ == "__main__": + test_tcnn_instant_ngp_field()