diff --git a/.gitignore b/.gitignore index 68bc17f9f..de34d8e54 100644 --- a/.gitignore +++ b/.gitignore @@ -158,3 +158,10 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +# Experiments and outputs +runs/ +outputs/ + +# Data +data/ \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 000000000..13aa61c4c --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,16 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Python: run_train_nerf", + "type": "python", + "request": "launch", + "program": "scripts/run_train_nerf.py", + "console": "integratedTerminal", + "args": [] + } + ] +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index 5c0e9d025..2f20bab22 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -13,6 +13,11 @@ "python.linting.pylintEnabled": true, "python.linting.flake8Enabled": false, "python.linting.enabled": true, + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, "python.linting.pylintArgs": [ "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*", "--max-line-length=120" diff --git a/README.md b/README.md index f0be6b824..092eea60a 100644 --- a/README.md +++ b/README.md @@ -15,10 +15,20 @@ pip install -r environment/requirements.txt # Install helper library cd external/goat python setup.py develop +# then go back, i.e., `cd /path/to/mattport` # For using with Jupyter python -m ipykernel install --user --name mattport --display-name "mattport" # Install mattport as a library python setup.py develop + +# Running the test cases +pytest tests +``` + +# Training a model + +``` +python scripts/run_train_nerf.py ``` \ No newline at end of file diff --git a/configs/config.yaml b/configs/config.yaml new file mode 100644 index 000000000..f95fc0ce2 --- /dev/null +++ b/configs/config.yaml @@ -0,0 +1,47 @@ +# OPTIONS FOR DIFFERENT NERF VARIANTS +# ---- +# sampler options +sampler: + asdf: null + num_coarse_samples: 64 + num_importance_samples: 128 + perturb: 1.0 + noise_std: 1.0 +# encoding options +encodings: + asdf: null +# network options +network: + asdf: null +# renderer options +renderer: + asdf: null +# loss options +loss: + asdf: null + type: mse +# ---- + +# dataset options +dataset: + data_directory: /path/to/data + dataset_type: blender # type of dataset + scene: lego + num_workers: 4 + +# camera options +camera: + model: pinhole + +# optimizer options +optimizer: + lr: 1e-4 + lr_scheduler: null + momentum: null + weight_decay: null + +# trainer options +trainer: + batch_size: 2 # rays per batch + chunk_size: 1024 # chunks per batch + device: cpu \ No newline at end of file diff --git a/environment/requirements.txt b/environment/requirements.txt index c1b62abdf..ef64b7b0c 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,5 +1,8 @@ +hydra-core pylint==2.13.4 +pytest torch==1.11.0 +torchtyping torchvision==0.12.0 torchaudio==0.11.0 black diff --git a/mattport/nerf/dataset/__init__.py b/mattport/nerf/dataset/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py new file mode 100644 index 000000000..2c48232fa --- /dev/null +++ b/mattport/nerf/dataset/base.py @@ -0,0 +1,68 @@ +""" +The dataset baseclass. +""" + +from abc import abstractmethod +import torch + +from mattport.nerf.cameras import Camera + + +class Dataset(torch.utils.data.Dataset): + """_summary_ + + Args: + torch (_type_): _description_ + """ + + def __init__( + self, + data_directory: str, + dataset_type: str, + scene: str + ): + """_summary_ + + Args: + data_directory (str): _description_ + type (str): _description_ + scene (str): _description_ + """ + super().__init__() + self.data_directory = data_directory + self.dataset_type = dataset_type + self.scene = scene + + + def __len__(self): + pass + + @abstractmethod + def get_image(self, idx): + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError + + @abstractmethod + 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 new file mode 100644 index 000000000..6966e87e6 --- /dev/null +++ b/mattport/nerf/dataset/build.py @@ -0,0 +1,26 @@ +""" +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 + + Args: + config (dict): _description_ + """ + + # TODO(ethan): better handle this with configs + module = importlib.import_module("mattport.nerf.dataset.base") + dataset = module.Dataset( + data_directory=config.data_directory, + dataset_type=config.dataset_type, + scene=config.scene, + ) + logging.info(dataset) + return dataset diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py new file mode 100644 index 000000000..01d1a6db1 --- /dev/null +++ b/mattport/nerf/graph.py @@ -0,0 +1,28 @@ +""" +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_ + + Args: + nn (_type_): _description_ + """ + + def __init__(self) -> None: + super().__init__() + raise NotImplementedError + + def forward(self): + """_summary_""" + raise NotImplementedError diff --git a/mattport/nerf/optimizer.py b/mattport/nerf/optimizer.py new file mode 100644 index 000000000..e6c6581cf --- /dev/null +++ b/mattport/nerf/optimizer.py @@ -0,0 +1,26 @@ +""" +Optimizer class. +""" + +from torch import optim + + +class Optimizer: + """_summary_""" + + 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): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e69de29bb..23c479f97 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -0,0 +1,61 @@ +""" +Code to train model. +""" +from mattport.nerf.dataset.build import build_dataset +from mattport.nerf.optimizer import Optimizer +from mattport.nerf.graph import Graph + + +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_ + """ + self.train_dataset = build_dataset(self.config.dataset) + self.test_dataset = build_dataset(self.config.dataset) + + def setup_graph(self): + """_summary_ + """ + self.graph = Graph() + + def setup_optimizer(self): + """_summary_ + """ + self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer) + + def load_checkpoint(self): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError + + def save_checkpoint(self): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError + + def train(self): + """_summary_""" + raise NotImplementedError + + def train_epoch(self): + """_summary_""" + raise NotImplementedError + + def train_iteration(self): + """_summary_""" + raise NotImplementedError diff --git a/mattport/viewer/README.md b/mattport/viewer/README.md index aca76ca0b..3ba37fcd0 100644 --- a/mattport/viewer/README.md +++ b/mattport/viewer/README.md @@ -6,4 +6,5 @@ The backend viewer code (python, websockets, ...). ``` https://github.com/lp-research/openvslam/tree/main/viewer +https://github.com/rdeits/meshcat ``` \ No newline at end of file diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index aaa3ffc7b..2ad9a5f97 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,5 +1,24 @@ -import os -import mattport.nerf as nerf +""" +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) + trainer.setup_dataset() + trainer.setup_graph() + trainer.setup_optimizer() + if __name__ == "__main__": - print("hello mattport") + main() diff --git a/tests/nerf/test_encodings.py b/tests/nerf/test_encodings.py index 49ad82215..99f26096c 100644 --- a/tests/nerf/test_encodings.py +++ b/tests/nerf/test_encodings.py @@ -2,4 +2,8 @@ import pytest def test_positional_encoding(): - raise NotImplementedError() + assert True + + +def test_lower_case(): + assert "MATTPORT".lower() == "mattport"