diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index b83fa2ecf..b850c9a51 100644 --- a/.github/workflows/code_checks.yml +++ b/.github/workflows/code_checks.yml @@ -40,4 +40,4 @@ jobs: --disable=too-many-instance-attributes - name: Test with pytest run: | - pytest tests --ignore tests/nerf/dataloader/test_ray_generator.py --typeguard-packages=mattport --torchtyping-patch-typeguard + pytest tests --typeguard-packages=mattport --torchtyping-patch-typeguard --disable-warnings diff --git a/.vscode/launch.json b/.vscode/launch.json index 0cae34962..9e87d43a5 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -3,9 +3,16 @@ // Hover to view descriptions of existing attributes. // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", - "configurations": [ + "configurations": [{ + "name": "Python: Current File", + "type": "python", + "request": "launch", + "program": "${file}", + "console": "integratedTerminal", + "justMyCode": true + }, { - "name": "Python: run_train_nerf", + "name": "Python: run_train", "type": "python", "request": "launch", "program": "scripts/run_train.py", diff --git a/.vscode/settings.json b/.vscode/settings.json index ed3e7382f..885849bfb 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -8,7 +8,7 @@ "**/runs*": true, "**/temp*": true, }, - // "editor.formatOnSave": true, + "editor.formatOnSave": true, "python.formatting.provider": "black", "python.formatting.blackArgs": [ "--line-length=120" diff --git a/data/blender/lego_test/train/r_0.png b/data/blender/lego_test/train/r_0.png new file mode 100755 index 000000000..b292e16fa Binary files /dev/null and b/data/blender/lego_test/train/r_0.png differ diff --git a/data/blender/lego_test/transforms_train.json b/data/blender/lego_test/transforms_train.json new file mode 100755 index 000000000..c398a6b58 --- /dev/null +++ b/data/blender/lego_test/transforms_train.json @@ -0,0 +1,35 @@ +{ + "camera_angle_x": 0.6911112070083618, + "frames": [ + { + "file_path": "./train/r_0", + "rotation": 0.012566370614359171, + "transform_matrix": [ + [ + -0.9999021887779236, + 0.004192245192825794, + -0.013345719315111637, + -0.05379832163453102 + ], + [ + -0.013988681137561798, + -0.2996590733528137, + 0.95394366979599, + 3.845470428466797 + ], + [ + -4.656612873077393e-10, + 0.9540371894836426, + 0.29968830943107605, + 1.2080823183059692 + ], + [ + 0.0, + 0.0, + 0.0, + 1.0 + ] + ] + } + ] +} \ No newline at end of file diff --git a/data/blender/lego_test/transforms_val.json b/data/blender/lego_test/transforms_val.json new file mode 100644 index 000000000..b41732ec0 --- /dev/null +++ b/data/blender/lego_test/transforms_val.json @@ -0,0 +1,35 @@ +{ + "camera_angle_x": 0.6911112070083618, + "frames": [ + { + "file_path": "./val/r_0", + "rotation": 0.012566370614359171, + "transform_matrix": [ + [ + -0.963964581489563, + -0.2611401677131653, + 0.0507759265601635, + 0.2046843022108078 + ], + [ + 0.26603081822395325, + -0.9462433457374573, + 0.18398693203926086, + 0.7416750192642212 + ], + [ + 7.450580596923828e-09, + 0.1908649355173111, + 0.9816163182258606, + 3.957021951675415 + ], + [ + 0.0, + 0.0, + 0.0, + 1.0 + ] + ] + } + ] +} \ No newline at end of file diff --git a/data/blender/lego_test/val/r_0.png b/data/blender/lego_test/val/r_0.png new file mode 100755 index 000000000..3f9c6d503 Binary files /dev/null and b/data/blender/lego_test/val/r_0.png differ diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index cc6313350..791546bda 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -5,8 +5,9 @@ Some dataset code. from typing import List import numpy as np -import torch +import numpy.typing as npt import PIL +import torch from PIL import Image from torch.utils.data import default_collate @@ -58,7 +59,7 @@ class ImageDataset(torch.utils.data.Dataset): def __len__(self): return len(self.image_filenames) - def get_image(self, image_idx: int) -> np.uint8: + def get_image(self, image_idx: int) -> npt.NDArray[np.uint8]: """Returns the image. Args: @@ -82,8 +83,9 @@ class ImageDataset(torch.utils.data.Dataset): pil_image = pil_image.resize( (image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.BILINEAR ) - image = np.array(pil_image) # shape is (h, w, 3 or 4) + image = np.array(pil_image, dtype="uint8") # shape is (h, w, 3 or 4) assert len(image.shape) == 3 + assert image.dtype == np.uint8 assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." return image diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index 872f67fad..50f5216ea 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -4,7 +4,7 @@ For loading the blender dataset format. import os from dataclasses import dataclass -from typing import List +from typing import Dict, List, Tuple import imageio import numpy as np @@ -69,8 +69,11 @@ def load_blender_data(basedir, downscale_factor=1.0, split="train"): return dataset_inputs -def get_dataset_inputs(data_directory: str, dataset_type: str, downscale_factor: float = 1.0): +def get_dataset_inputs_dict( + data_directory: str, dataset_type: str, downscale_factor: float = 1.0, splits: Tuple[str] = ("train", "val") +) -> 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. Args: data_directory (str): _description_ @@ -78,11 +81,14 @@ def get_dataset_inputs(data_directory: str, dataset_type: str, downscale_factor: downscale_factor (float, optional): _description_. Defaults to 1.0. Returns: - DatasetInputs: The inputs needed for generating rays. + Dict[str, DatasetInputs]: The inputs needed for generating rays. """ + dataset_inputs_dict = {} if dataset_type == "blender": - dataset_inputs = load_blender_data(get_absolute_path(data_directory), downscale_factor=downscale_factor) + for split in splits: + dataset_inputs = load_blender_data(get_absolute_path(data_directory), downscale_factor=downscale_factor) + dataset_inputs_dict[split] = dataset_inputs else: raise NotImplementedError() - return dataset_inputs + return dataset_inputs_dict diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index d081f117d..483da5228 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -53,7 +53,7 @@ class MLP(FieldModule): else: for i in range(self.num_layers - 1): if i == 0: - assert i not in self.skip_connections, "Skip connection at layer 0 doesn't make sense." + assert i not in list(self.skip_connections), "Skip connection at layer 0 doesn't make sense." layers.append(nn.Linear(self.in_dim, self.layer_width)) elif i in self.skip_connections: layers.append(nn.Linear(self.layer_width + self.in_dim, self.layer_width)) diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py new file mode 100644 index 000000000..5461d0b40 --- /dev/null +++ b/mattport/nerf/graph/instant_ngp.py @@ -0,0 +1,4 @@ +""" +Instant NGP implementation. +TODO: +""" diff --git a/mattport/nerf/graph/tensorf.py b/mattport/nerf/graph/tensorf.py new file mode 100644 index 000000000..b9ef32bfc --- /dev/null +++ b/mattport/nerf/graph/tensorf.py @@ -0,0 +1,4 @@ +""" +TensorRF implementation. +TODO: +""" diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index cd8ceabfe..2212c09af 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -16,7 +16,7 @@ from torch.utils.data import DataLoader import mattport.utils.writer from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch -from mattport.nerf.dataset.utils import get_dataset_inputs +from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs_dict from mattport.nerf.metrics import get_psnr from mattport.nerf.optimizers import Optimizers from mattport.utils import profiler @@ -37,8 +37,7 @@ class Trainer: self.train_image_dataset = None self.train_dataset = None self.train_dataloader = None - self.test_dataset = None - self.test_dataloader = None + self.val_image_dataset = None # model variables self.graph = None self.optimizers = None @@ -49,13 +48,21 @@ class Trainer: self.writer = writer(local_rank, world_size, save_dir=self.config.logging_configs.writer.save_dir) self.stats = StatsTracker(config, self.is_main_thread) profiler.PROFILER = profiler.Profiler(config, self.is_main_thread) + self.dry_run = self.config.get("dry_run", False) @profiler.time_function - def setup_dataset(self): + def setup(self): + """Setup the Trainer by calling other setup functions.""" + dataset_inputs_dict = get_dataset_inputs_dict(**self.config.dataset) + self.setup_datasets(dataset_inputs_dict) + self.setup_graph(dataset_inputs_dict["train"]) + + @profiler.time_function + def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs]): """_summary_""" - dataset_inputs = get_dataset_inputs(**self.config.dataset) self.train_image_dataset = ImageDataset( - image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor + image_filenames=dataset_inputs_dict["train"].image_filenames, + downscale_factor=dataset_inputs_dict["train"].downscale_factor, ) self.train_dataset = CollateIterDataset( self.train_image_dataset, @@ -72,12 +79,18 @@ class Trainer: collate_fn=collate_batch_size_one, pin_memory=True, ) - # TODO(ethan): implement the test dataset + self.val_image_dataset = ImageDataset( + image_filenames=dataset_inputs_dict["val"].image_filenames, + downscale_factor=dataset_inputs_dict["val"].downscale_factor, + ) @profiler.time_function - def setup_graph(self): - """_summary_""" - dataset_inputs = get_dataset_inputs(**self.config.dataset) + def setup_graph(self, dataset_inputs: DatasetInputs): + """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. + """ self.graph = instantiate( self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world ).to(f"cuda:{self.local_rank}") @@ -170,10 +183,9 @@ class Trainer: # TODO: add the learning rates to tensorboard/logging if step != 0 and self.config.steps_per_save and step % self.config.steps_per_save == 0: self.save_checkpoint(self.config.model_dir, step) - if step != 0 and step % self.config.steps_per_test == 0: + if step % self.config.steps_per_test == 0: self.test_image(image_idx=0, step=step) - self.test_image(image_idx=10, step=step) - self.test_image(image_idx=20, step=step) + _ = self.test_image(image_idx=10, step=step) if not self.dry_run else None self.stats.print_stats(i / num_iterations) self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1) @@ -198,7 +210,7 @@ class Trainer: @profiler.time_function def test_image(self, image_idx, step): """Test a specific image.""" - image = self.train_image_dataset[image_idx]["image"] # ground truth + image = self.val_image_dataset[image_idx]["image"] # ground truth image_height, image_width, _ = image.shape pixel_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij") pixel_coords = torch.stack(pixel_coords, dim=-1).long() @@ -219,8 +231,8 @@ class Trainer: rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu() combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1) - self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test") + self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="val") fine_psnr = get_psnr(image, rgb_fine) - self.stats.update_value(Stats.CURR_TEST_PSNR, fine_psnr, step) - self.writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test") + self.stats.update_value(Stats.CURR_TEST_PSNR, float(fine_psnr), step) + self.writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", float(fine_psnr), step, group="val") diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 50920d753..69d5f0e41 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -18,6 +18,16 @@ from mattport.utils.decorators import check_main_thread, decorate_all to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) +def get_tensorboard_name(name: str, group: str = None, prefix: str = None): + """Returns a string for tensorboard with an optional group and prefix. + Where tensorboard_name has the form `group/prefix-name`. + """ + group_string = f"{group}/" if group else "" + prefix_string = f"{prefix}" if prefix else "" + tensorboard_name = f"{group_string}{prefix_string}{name}" + return tensorboard_name + + class Writer: """Writer class""" @@ -88,7 +98,8 @@ class TensorboardWriter(Writer): x (TensorType["H", "W", 3]): rendered image to write """ x = to8b(x) - self.tb_writer.add_image(name, x, step, dataformats="HWC") + tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix) + self.tb_writer.add_image(tensorboard_name, x, step, dataformats="HWC") def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: """Tensorboard method to write a single scalar value to the logger @@ -99,9 +110,8 @@ class TensorboardWriter(Writer): y (float): y value to write group (str)): a prefix to group tensorboard scalars """ - group_s = f"{group}/" if group else "" - prefix_s = f"{prefix}" if prefix else "" - self.tb_writer.add_scalar(f"{group_s}{prefix_s}{name}", scalar, step) + tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix) + self.tb_writer.add_scalar(tensorboard_name, scalar, step) @decorate_all([check_main_thread]) diff --git a/notebooks/visualize_dataset.ipynb b/notebooks/visualize_dataset.ipynb index 1012caefa..1e512fe71 100644 --- a/notebooks/visualize_dataset.ipynb +++ b/notebooks/visualize_dataset.ipynb @@ -22,7 +22,7 @@ "\n", "from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch\n", "from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one\n", - "from mattport.nerf.dataset.utils import get_dataset_inputs\n", + "from mattport.nerf.dataset.utils import get_dataset_inputs_dict\n", "from mattport.nerf.field_modules.ray_generator import RayGenerator\n", "from mattport.structures.rays import RayBundle\n", "from mattport.utils.io import get_absolute_path\n", @@ -49,7 +49,7 @@ " cfg.dataset.downscale_factor = 4\n", " \n", "cfg.dataset.data_directory = get_absolute_path(cfg.dataset.data_directory)\n", - "dataset_inputs = get_dataset_inputs(**cfg.dataset)\n", + "dataset_inputs = get_dataset_inputs(**cfg.dataset)[\"train\"]\n", "image_dataset = ImageDataset(\n", " image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n", " )\n", diff --git a/scripts/run_train.py b/scripts/run_train.py index a93f51c08..009634f11 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -116,8 +116,7 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any: Any: TODO(): determine the return type """ trainer = Trainer(config, local_rank, world_size) - trainer.setup_graph() - trainer.setup_dataset() + trainer.setup() trainer.train() return 0 diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml index 7b30f9e57..22932e573 100644 --- a/tests/configs/test_default.yml +++ b/tests/configs/test_default.yml @@ -1,3 +1,5 @@ +dry_run: true + defaults: # This makes the logging module logs colorful - override hydra/job_logging: colorlog @@ -39,7 +41,7 @@ model_dir: /tmp/mattport_models/ # dataset options dataset: - data_directory: data/blender/lego + data_directory: data/blender/lego_test dataset_type: blender downscale_factor: 4 @@ -47,7 +49,7 @@ dataset: dataloader: num_workers: 4 shuffle: true - num_images_to_sample_from: 8 + num_images_to_sample_from: 1 num_rays_per_batch: 1024 num_times_to_repeat_images: 40 diff --git a/tests/nerf/dataloader/test_ray_generator.py b/tests/nerf/dataloader/test_ray_generator.py index 108a34a7f..d1b7b3fda 100644 --- a/tests/nerf/dataloader/test_ray_generator.py +++ b/tests/nerf/dataloader/test_ray_generator.py @@ -9,7 +9,7 @@ from tqdm import tqdm from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch -from mattport.nerf.dataset.utils import get_dataset_inputs +from mattport.nerf.dataset.utils import get_dataset_inputs_dict from mattport.nerf.field_modules.ray_generator import RayGenerator from mattport.utils.io import get_absolute_path from mattport.viewer.plotly import visualize_dataset @@ -30,21 +30,21 @@ def visualize_batch(batch): return image -def test_dataloader(): +def test_dataloader(visualize=False): """Testing for the dataloader from input dataset parameters to rays.""" - data_directory = "data/blender/lego" + data_directory = "data/blender/lego_test" dataset_type = "blender" downscale_factor = 1 - num_images_to_sample_from = 8 + num_images_to_sample_from = 1 num_times_to_repeat_images = 40 num_rays_per_batch = 1024 num_workers = 0 data_directory = get_absolute_path(data_directory) - dataset_inputs = get_dataset_inputs( + dataset_inputs = get_dataset_inputs_dict( data_directory=data_directory, dataset_type=dataset_type, downscale_factor=downscale_factor - ) + )["train"] image_dataset = ImageDataset( image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor ) @@ -72,14 +72,17 @@ def test_dataloader(): # visualize the batch image = visualize_batch(batch) - imageio.imwrite("temp0.png", image) + if visualize: + imageio.imwrite("temp0.png", image) # visualize the RayBundle small_ray_bundle = ray_bundle.sample(100) fig = visualize_dataset(camera_origins=dataset_inputs.camera_to_world[:, :3, 3], ray_bundle=small_ray_bundle) - fig.write_image("temp1.png") - fig.write_html("temp1.html") + if visualize: + fig.write_image("temp1.png") + fig.write_html("temp1.html") + assert True if __name__ == "__main__": - test_dataloader() + test_dataloader(visualize=True)