Merge pull request #43 from ethanweber/ethan/apr29_using_real_test_set

Validation images integration.

.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

.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",

.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"

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
+                ]
+            ]
+        }
+    ]
+}

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
+                ]
+            ]
+        }
+    ]
+}

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
 

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

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))

mattport/nerf/graph/instant_ngp.py

@@ -0,0 +1,4 @@
+"""
+Instant NGP implementation.
+TODO:
+"""

mattport/nerf/graph/tensorf.py

@@ -0,0 +1,4 @@
+"""
+TensorRF implementation.
+TODO:
+"""

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")

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])

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",

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
 

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
 

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)