pylint and test cases for dataloader

mattport/nerf/dataset/collate.py

@@ -2,14 +2,15 @@
 Classes to help with caching data while training.
 """
 
-from typing import Callable
-import torch
 import random
+from typing import Callable
 
+import torch
 from torch.utils.data import default_collate
 
 
 def collate_batch_size_one(batch_list):
+    """Use the default collate function but then squeze to avoid the batch dimension being added."""
     assert len(batch_list) == 1
     collated_batch = default_collate(batch_list)
     for key in collated_batch:
@@ -41,7 +42,11 @@ class CollateIterDataset(torch.utils.data.IterableDataset):
         self.num_repeated = self.num_times_to_repeat  # starting value
         self.cached_batch_list = None
 
+    def __getitem__(self, idx):
+        return self.dataset.__getitem__(idx)
+
     def get_batch_list(self):
+        """Returns a list of batches from the dataset attribute."""
         # sampling without replacement. TODO(ethan): don't use 'random'
         indices = random.sample(range(len(self.dataset)), k=self.num_samples_to_collate)
         batch_list = [self.dataset.__getitem__(idx) for idx in indices]

mattport/nerf/dataset/image_dataset.py

@@ -32,9 +32,11 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False)
     assert pixels.shape == (num_rays_per_batch, 3), pixels.shape
 
     # Needed to correct the random indices to their actual camera idx locations.
+    local_indices = indices.clone()
     indices[:, 0] = batch["image_idx"][c]
     collated_batch = {
-        "indices": indices,
+        "local_indices": local_indices, # local to the batch returned
+        "indices": indices, # with the abs camera indices
         "pixels": pixels,
         "mask": mask,
     }

mattport/nerf/field_modules/mlp.py

@@ -1,8 +1,7 @@
 """
 Multi Layer Perceptron
 """
-from tkinter import X
-from typing import List, Optional
+from typing import Tuple, Optional
 
 import torch
 from torch import nn
@@ -20,7 +19,7 @@ class MLP(FieldModule):
         out_dim: int,
         num_layers: int,
         layer_width: int = None,
-        skip_connections: List[int] = [],
+        skip_connections: Tuple[int] = (),
         activation: Optional[nn.Module] = nn.ReLU(),
         out_activation: Optional[nn.Module] = None,
     ) -> None:

mattport/nerf/graph/base.py

@@ -34,7 +34,7 @@ class Node:
 class Graph(nn.Module):
     """_summary_"""
 
-    def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None:
+    def __init__(self, intrinsics=None, camera_to_world=None) -> None:
         super().__init__()
         self.intrinsics = intrinsics
         self.camera_to_world = camera_to_world

mattport/nerf/graph/vanilla_nerf.py

@@ -3,7 +3,7 @@ Implementation of vanilla nerf.
 """
 
 
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import torch
 from torch import nn
@@ -17,14 +17,14 @@ from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import RGBRenderer
-from mattport.nerf.sampler import PDFSampler, UniformSampler
+from mattport.nerf.sampler import PDFSampler, UniformSampler  # pylint: disable=unused-import
 from mattport.structures.rays import RaySamples
 
 
 class NeRFField(nn.Module):
     """NeRF module"""
 
-    def __init__(self, num_layers=8, layer_width=256, skip_connections=[4]) -> None:
+    def __init__(self, num_layers=8, layer_width=256, skip_connections: Tuple = (4)) -> None:
         super().__init__()
         self.num_layers = num_layers
         self.layer_width = layer_width
@@ -99,6 +99,7 @@ class NeRFGraph(Graph):
         self.sampler_uniform = UniformSampler(
             near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples
         )
+        # TODO: this needs to be corrected
         # self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples)
         self.sampler_pdf = UniformSampler(
             near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_importance_samples

mattport/nerf/optimizers.py

@@ -56,11 +56,10 @@ class Optimizers:
 
     def optimizer_step_all(self):
         """Run step for all optimizers."""
-        # logging.info("optimizer_step_all")
-        for param_name, optimizer in self.optimizers.items():
+        for _, optimizer in self.optimizers.items():
+            # note that they key is the parameter name
             optimizer.step()
 
     def scheduler_step_all(self, step):
         """Run step for all schedulers."""
-        # logging.info("optimizer_step_all")
-        pass
+        # TODO: this needs to be implemented

mattport/nerf/trainer.py

@@ -5,6 +5,7 @@ import datetime
 import logging
 import os
 from time import time
+from typing import Dict
 
 import torch
 import torch.distributed as dist
@@ -32,6 +33,7 @@ class Trainer:
         self.local_rank = local_rank
         self.world_size = world_size
         # dataset variables
+        self.train_image_dataset = None
         self.train_dataset = None
         self.train_dataloader = None
         self.test_dataset = None
@@ -66,7 +68,7 @@ class Trainer:
             self.train_dataset,
             batch_size=1,
             num_workers=self.config.dataloader.num_workers,
-            collate_fn=lambda batch_list: collate_batch_size_one(batch_list),
+            collate_fn=collate_batch_size_one,
             pin_memory=True,
         )
         # TODO(ethan): implement the test dataset
@@ -160,14 +162,10 @@ class Trainer:
             mssg += f"{k}: {v} | "
         print(mssg, end="\r")
 
-    def get_aggregated_loss(self, losses):
-        """_summary_
-
-        Args:
-            losses (_type_): _description_
-
-        Returns:
-            _type_: _description_
+    @classmethod
+    def get_aggregated_loss(cls, losses: Dict[str, torch.tensor]):
+        """Returns the aggregated losses and the scalar for calling .backwards() on.
+        # TODO: move this out to another file/class/etc.
         """
         loss_sum = 0.0
         loss_dict = {}

mattport/utils/misc.py

@@ -1,3 +1,8 @@
+"""
+Miscellaneous helper code.
+"""
+
+
 class DotDict(dict):
     """
     dot.notation access to dictionary attributes

notebooks/visualize_dataset.ipynb

@@ -20,18 +20,19 @@
     "import torch\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
-    "from mattport.nerf.dataset.utils import get_dataset_inputs\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.field_modules.ray_generator import RayGenerator\n",
     "from mattport.structures.rays import RayBundle\n",
     "from mattport.utils.io import get_absolute_path\n",
+    "from mattport.viewer.plotly import get_line_segments_from_lines\n",
     "\n",
     "from hydra import compose, initialize\n",
     "from omegaconf import open_dict\n",
     "import pprint\n",
     "from tqdm import tqdm\n",
     "\n",
-    "import goat\n",
     "import random\n",
     "from tqdm import tqdm"
    ]
@@ -49,16 +50,23 @@
     "    \n",
     "cfg.dataset.data_directory = get_absolute_path(cfg.dataset.data_directory)\n",
     "dataset_inputs = get_dataset_inputs(**cfg.dataset)\n",
-    "dataset = ImageDataset(\n",
-    "    image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n",
+    "image_dataset = ImageDataset(\n",
+    "            image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n",
+    "        )\n",
+    "iter_dataset = CollateIterDataset(\n",
+    "    image_dataset,\n",
+    "    collate_fn=lambda batch_list: collate_batch(\n",
+    "        batch_list, cfg.dataloader.num_rays_per_batch, keep_full_image=True\n",
+    "    ),\n",
+    "    num_samples_to_collate=cfg.dataloader.num_images_to_sample_from,\n",
+    "    num_times_to_repeat=cfg.dataloader.num_times_to_repeat_images,\n",
     ")\n",
-    "\n",
     "dataloader = DataLoader(\n",
-    "    dataset,\n",
-    "    batch_size=cfg.dataloader.num_images_to_sample_from,\n",
-    "    collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True),\n",
+    "    iter_dataset,\n",
+    "    batch_size=1,\n",
     "    num_workers=cfg.dataloader.num_workers,\n",
-    "    shuffle=True,\n",
+    "    collate_fn=collate_batch_size_one,\n",
+    "    pin_memory=True,\n",
     ")\n",
     "dataloader_iter = iter(dataloader)\n",
     "\n",
@@ -67,7 +75,7 @@
     "num_batches = 10\n",
     "for _ in tqdm(range(num_batches)):\n",
     "    batch = next(dataloader_iter)\n",
-    "    ray_bundle = ray_generator.forward(batch.indices)\n",
+    "    ray_bundle = ray_generator.forward(batch[\"indices\"])\n",
     "    break"
    ]
   },
@@ -78,7 +86,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "media.show_image(dataset.get_image(10))"
+    "media.show_image(image_dataset.get_image(10))"
    ]
   },
   {
@@ -90,11 +98,11 @@
    "source": [
     "def show_batch(batch):\n",
     "    # set the color of the sampled rays\n",
-    "    c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)]\n",
-    "    batch.image[c,y,x] = 1.0\n",
+    "    c, y, x = [i.flatten() for i in torch.split(batch[\"local_indices\"], 1, dim=-1)]\n",
+    "    batch[\"image\"][c,y,x] = 1.0\n",
     "\n",
-    "    # batch.image is num_images, h, w, 3\n",
-    "    images = torch.split(batch.image, 1, dim=0)\n",
+    "    # batch[\"image\"] is num_images, h, w, 3\n",
+    "    images = torch.split(batch[\"image\"], 1, dim=0)\n",
     "    image_list = [image[0] for image in images]\n",
     "    image = torch.cat(image_list, dim=1) # cat along the width dimension\n",
     "\n",
@@ -117,16 +125,6 @@
     "batch = sample_and_show_batch()"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "f5f2e2e8-8d12-4e12-9b9f-e70a4f25bcee",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# batch.keys()"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -134,7 +132,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ray_indices = batch.indices\n",
+    "ray_indices = batch[\"indices\"]\n",
     "ray_bundle = ray_generator(ray_indices)"
    ]
   },
@@ -180,7 +178,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data += goat.plotly_utils.get_line_segments_from_lines(lines)"
+    "data += get_line_segments_from_lines(lines)"
    ]
   },
   {
@@ -240,6 +238,14 @@
    "metadata": {},
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "722f5370-3b0f-4ad1-87a1-df479ae85b60",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

tests/nerf/dataloader/test_ray_generator.py

@@ -4,10 +4,10 @@ Code to test ray generation.
 
 import imageio
 import torch
-from hydra import compose, initialize
 from torch.utils.data import DataLoader
 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.field_modules.ray_generator import RayGenerator
@@ -19,11 +19,11 @@ def visualize_batch(batch):
     """Returns an image from a batch."""
     assert "image" in batch, "You must call collate_batch with keep_full_image=True for this to work."
     # set the color of the sampled rays
-    c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)]
-    batch.image[c, y, x] = 1.0
+    c, y, x = [i.flatten() for i in torch.split(batch["local_indices"], 1, dim=-1)]
+    batch["image"][c, y, x] = 1.0
 
-    # batch.image is num_images, h, w, 3
-    images = torch.split(batch.image, 1, dim=0)
+    # batch["image"] is num_images, h, w, 3
+    images = torch.split(batch["image"], 1, dim=0)
     image_list = [image[0] for image in images]
     image = torch.cat(image_list, dim=1)  # cat along the width dimension
     image = (image * 255.0).to(torch.uint8)
@@ -33,23 +33,33 @@ def visualize_batch(batch):
 def test_dataloader():
     """Testing for the dataloader from input dataset parameters to rays."""
 
-    with initialize(config_path="configs"):
-        cfg = compose(config_name="blender_lego.yaml")
+    data_directory = "data/blender/lego"
+    dataset_type = "blender"
+    downscale_factor = 1
+    num_images_to_sample_from = 8
+    num_times_to_repeat_images = 40
+    num_rays_per_batch = 1024
+    num_workers = 0
 
-    cfg.dataset.data_directory = get_absolute_path(
-        cfg.dataset.data_directory
-    )  # TODO(ethan): create a helper function for this
-    dataset_inputs = get_dataset_inputs(**cfg.dataset)
-    dataset = ImageDataset(
+    data_directory = get_absolute_path(data_directory)
+    dataset_inputs = get_dataset_inputs(
+        data_directory=data_directory, dataset_type=dataset_type, downscale_factor=downscale_factor
+    )
+    image_dataset = ImageDataset(
         image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor
     )
-
+    iter_dataset = CollateIterDataset(
+        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,
+    )
     dataloader = DataLoader(
-        dataset,
-        batch_size=cfg.dataloader.num_images_to_sample_from,
-        collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True),
-        num_workers=cfg.dataloader.num_workers,
-        shuffle=True,
+        iter_dataset,
+        batch_size=1,
+        num_workers=num_workers,
+        collate_fn=collate_batch_size_one,
+        pin_memory=True,
     )
     dataloader_iter = iter(dataloader)
 
@@ -58,7 +68,7 @@ def test_dataloader():
     num_batches = 10
     for _ in tqdm(range(num_batches)):
         batch = next(dataloader_iter)
-        ray_bundle = ray_generator.forward(batch.indices)
+        ray_bundle = ray_generator.forward(batch["indices"])
 
     # visualize the batch
     image = visualize_batch(batch)