diff --git a/.vscode/settings.json b/.vscode/settings.json index 3f6d1df2c..f476de807 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -11,6 +11,7 @@ "--line-length=120" ], "python.linting.pylintEnabled": true, + "python.linting.pylintArgs": [], "python.linting.flake8Enabled": false, "python.linting.enabled": true, "python.testing.pytestArgs": [ diff --git a/configs/config.yaml b/configs/config.yaml index 2175cf59b..f95fc0ce2 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,7 +1,3 @@ -# defaults: -# - experiment: lego_bulldozer -# - dataset: blender - # OPTIONS FOR DIFFERENT NERF VARIANTS # ---- # sampler options @@ -28,7 +24,8 @@ loss: # dataset options dataset: - dataset: blender # type of dataset + data_directory: /path/to/data + dataset_type: blender # type of dataset scene: lego num_workers: 4 diff --git a/configs/dataset/blender.yaml b/configs/dataset/blender.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/dataset/dtu.yaml b/configs/dataset/dtu.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/dataset/shapenet.yaml b/configs/dataset/shapenet.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/experiment/lego_bulldozer.yaml b/configs/experiment/lego_bulldozer.yaml deleted file mode 100644 index 02e9326c7..000000000 --- a/configs/experiment/lego_bulldozer.yaml +++ /dev/null @@ -1,2 +0,0 @@ -name: lego_bulldozer -feature_size: 5 \ No newline at end of file diff --git a/configs/optimizer/adam.yaml b/configs/optimizer/adam.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/optimizer/sgd.yaml b/configs/optimizer/sgd.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py index 67cbd6422..2c48232fa 100644 --- a/mattport/nerf/dataset/base.py +++ b/mattport/nerf/dataset/base.py @@ -1,6 +1,12 @@ +""" +The dataset baseclass. +""" + from abc import abstractmethod import torch +from mattport.nerf.cameras import Camera + class Dataset(torch.utils.data.Dataset): """_summary_ @@ -12,20 +18,51 @@ class Dataset(torch.utils.data.Dataset): def __init__( self, data_directory: str, - ) -> None: + dataset_type: str, + scene: str + ): + """_summary_ + + Args: + data_directory (str): _description_ + type (str): _description_ + scene (str): _description_ + """ super().__init__() - pass + self.data_directory = data_directory + self.dataset_type = dataset_type + self.scene = scene + def __len__(self): pass @abstractmethod def get_image(self, idx): - pass + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError @abstractmethod - def get_pose(self, idx): - pass + def get_camera(self, idx) -> Camera: + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + + Returns: + Camera: _description_ + """ + raise NotImplementedError def __getitem__(self, idx): raise NotImplementedError diff --git a/mattport/nerf/dataset/build.py b/mattport/nerf/dataset/build.py index e9eceb00f..6966e87e6 100644 --- a/mattport/nerf/dataset/build.py +++ b/mattport/nerf/dataset/build.py @@ -1,8 +1,12 @@ -from argon2 import DEFAULT_RANDOM_SALT_LENGTH -from mattport.nerf.dataset.base import Dataset +""" +Builds the dataset +""" import importlib import logging +from mattport.nerf.dataset.base import Dataset + + def build_dataset(config: dict) -> Dataset: """Returns the dataset according to the specified config. # TODO(ethan): change the type from dict to be a dataclass @@ -11,7 +15,12 @@ def build_dataset(config: dict) -> Dataset: config (dict): _description_ """ + # TODO(ethan): better handle this with configs module = importlib.import_module("mattport.nerf.dataset.base") - dataset = module.Dataset(**config) + dataset = module.Dataset( + data_directory=config.data_directory, + dataset_type=config.dataset_type, + scene=config.scene, + ) logging.info(dataset) - return dataset \ No newline at end of file + return dataset diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index ec52c9850..01d1a6db1 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -1,6 +1,17 @@ +""" +The Graph module contains all trainable parameters. +""" +from dataclasses import dataclass from torch import nn +@dataclass +class GraphInputs: + """Datastucture to encode the inputs to the graph.""" + + points: float + + class Graph(nn.Module): """_summary_ @@ -8,8 +19,10 @@ class Graph(nn.Module): nn (_type_): _description_ """ - def __init__(self, config) -> None: + def __init__(self) -> None: super().__init__() - self.config = config + raise NotImplementedError - \ No newline at end of file + def forward(self): + """_summary_""" + raise NotImplementedError diff --git a/mattport/nerf/optimizer.py b/mattport/nerf/optimizer.py index 3484bada3..e6c6581cf 100644 --- a/mattport/nerf/optimizer.py +++ b/mattport/nerf/optimizer.py @@ -1,21 +1,26 @@ +""" +Optimizer class. +""" + from torch import optim class Optimizer: """_summary_""" - def __init__( - self, - params: dict, - lr: float) -> None: + def __init__(self, params: dict, lr: float): """ Args: params (dict): _description_ lr (float): Learning rate. """ self.optimizer = optim.Adam(params) + self.lr = lr def test(self): - raise NotImplementedError + """_summary_ - \ No newline at end of file + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 4722d16a8..23c479f97 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,15 +1,20 @@ -from abc import abstractmethod -import logging -import torch - +""" +Code to train model. +""" +from mattport.nerf.dataset.build import build_dataset from mattport.nerf.optimizer import Optimizer from mattport.nerf.graph import Graph -from mattport.nerf.dataset.build import build_dataset class Trainer(object): + """Training class""" + def __init__(self, config: dict): self.config = config + self.train_dataset = None + self.test_dataset = None + self.graph = None + self.optimizer = None def setup_dataset(self): """_summary_ @@ -20,13 +25,11 @@ class Trainer(object): def setup_graph(self): """_summary_ """ - self.graph = Graph( - encodings_config=self.config.encodings, - renderer_config=self.config.renderer, - loss_config=self.config.loss - ) + self.graph = Graph() def setup_optimizer(self): + """_summary_ + """ self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer) def load_checkpoint(self): @@ -47,12 +50,12 @@ class Trainer(object): def train(self): """_summary_""" - pass + raise NotImplementedError def train_epoch(self): """_summary_""" - pass + raise NotImplementedError def train_iteration(self): """_summary_""" - pass + raise NotImplementedError diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index 65a097aef..2ad9a5f97 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,21 +1,23 @@ -import os -from typing import Dict -from mattport.nerf.trainer import Trainer +""" +run_train_nerf.py +""" import logging import hydra from omegaconf import DictConfig +from mattport.nerf.trainer import Trainer + + logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) @hydra.main(config_path="../configs", config_name="config.yaml") def main(config: DictConfig): + """Main function.""" trainer = Trainer(config) - print("trainer!") - trainer.setup_dataset() - # trainer.setup_graph() - # trainer.setup_optimizer() + trainer.setup_graph() + trainer.setup_optimizer() if __name__ == "__main__":