Misc cleanup (#494)

* misc cleanup

* Fix dependencies

.github/workflows/core_code_checks.yml

@@ -25,7 +25,7 @@ jobs:
           key: ${{ env.pythonLocation }}-${{ hashFiles('pyproject.toml') }}
       - name: Install dependencies
         run: |
-          pip install --upgrade --upgrade-strategy eager -e .
+          pip install --upgrade --upgrade-strategy eager -e .[dev]
       - name: Run license checks
         run: |
           ./scripts/licensing/license_headers.sh --check
@@ -38,7 +38,7 @@ jobs:
         run: black docs/ nerfactory/ scripts/ tests/ --check
       - name: Python Pylint
         run: |
-          pylint --disable=cyclic-import nerfactory tests scripts
+          pylint nerfactory tests scripts
       - name: Test with pytest
         run: |
           pytest

docs/reference/contributing.md

@@ -11,6 +11,21 @@ The project is set up for development in VSCode, we recommend using it if you pl
 | Docs            | Sphinx |
 | Docstring style | Google |
 
+## Requirements
+
+To install the required packages:
+
+```bash
+pip install -e .[dev]
+pip install -e .[docs]
+```
+
+You may also need to install [pandoc](https://pandoc.org/). If you are using `conda` you can run the following:
+
+```bash
+conda install -c conda-forge pandoc
+```
+
 ## Committing
 
 1. Make your modifications
@@ -22,12 +37,12 @@ The project is set up for development in VSCode, we recommend using it if you pl
    python scripts/debugging/run_actions.py
    ```
 
-   This will perform the following checks and actions:
+This will perform the following checks and actions:
 
-   - Black/ Linting style check: Ensure that the code is consistently and properly formatted.
-   - Pytests: Runs pytests locally to make sure the logic of any added code does not break existing logic.
-   - Documentation build: Builds the documentation locally to make sure none of the changes result in any warnings or errors in the docs.
-   - Licensing: automatically adds licensing headers to the correct files.
+- Black/ Linting style check: Ensure that the code is consistently and properly formatted.
+- Pytests: Runs pytests locally to make sure the logic of any added code does not break existing logic.
+- Documentation build: Builds the documentation locally to make sure none of the changes result in any warnings or errors in the docs.
+- Licensing: automatically adds licensing headers to the correct files.
 
 In order to merge changes to the code base, all of these checks must be passing. If you pass these tests locally, you will likely pass on github servers as well (results in a green checkmark next to your commit).
 
@@ -35,20 +50,6 @@ In order to merge changes to the code base, all of these checks must be passing.
 
 ## Documentation
 
-### Requirements
-
-To install the required packages:
-
-```bash
-pip install -e .[docs]
-```
-
-You may also need to install [pandoc](https://pandoc.org/). If you are using `conda` you can run the following:
-
-```bash
-conda install -c conda-forge pandoc
-```
-
 ### Building
 
 Run the following to build the documentation:

nerfactory/engine/base.py

@@ -1,13 +0,0 @@
-# Copyright 2022 The Plenoptix Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

nerfactory/engine/trainer.py

@@ -23,7 +23,7 @@ import logging
 import os
 import time
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
 
 import torch
 from rich import console
@@ -68,8 +68,21 @@ class Trainer:
         config: The configuration object.
         local_rank: Local rank of the process.
         world_size: World size of the process.
+
+    Attributes:
+        config: The configuration object.
+        local_rank: Local rank of the process.
+        world_size: World size of the process.
+        device: The device to run the training on.
+        pipeline: The pipeline object.
+        optimizers: The optimizers object.
+        callbacks: The callbacks object.
     """
 
+    pipeline: Pipeline
+    optimizers: Optimizers
+    callbacks: List[TrainingCallback]
+
     def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1):
         self.config = config
         self.local_rank = local_rank
@@ -79,12 +92,7 @@ class Trainer:
         if self.device == "cpu":
             self.mixed_precision = False
             logging.warning("Mixed precision is disabled for CPU training.")
-        # model variables
-        self.pipeline: Pipeline
-        self.optimizers: Optimizers
-        self.start_step = 0
-        # training callbacks
-        self.callbacks: List[TrainingCallback]
+        self._start_step = 0
         # optimizers
         self.grad_scaler = GradScaler(enabled=self.mixed_precision)
         # logging/viewer variables
@@ -112,7 +120,11 @@ class Trainer:
 
         # TODO(ethan): do this for pipeline, not pipeline.model
         self.callbacks = self.pipeline.get_training_callbacks(
-            TrainingCallbackAttributes(optimizers=self.optimizers, grad_scaler=self.grad_scaler, pipeline=self.pipeline)
+            TrainingCallbackAttributes(
+                optimizers=self.optimizers,  # type: ignore
+                grad_scaler=self.grad_scaler,  # type: ignore
+                pipeline=self.pipeline,  # type: ignore
+            )
         )
 
     def train(self) -> None:
@@ -122,7 +134,7 @@ class Trainer:
         self._init_viewer_scene()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
             num_iterations = self.config.trainer.max_num_iterations
-            for step in range(self.start_step, self.start_step + num_iterations):
+            for step in range(self._start_step, self._start_step + num_iterations):
                 # if the viewer used, the rendering of the viewer will be included in the iteration train time
                 with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:
 
@@ -275,7 +287,7 @@ class Trainer:
             load_path = load_dir / f"step-{load_step:09d}.ckpt"
             assert load_path.exists(), f"Checkpoint {load_path} does not exist"
             loaded_state = torch.load(load_path, map_location="cpu")
-            self.start_step = loaded_state["step"] + 1
+            self._start_step = loaded_state["step"] + 1
             # load the checkpoints for pipeline, optimizers, and gradient scalar
             self.pipeline.load_pipeline(loaded_state["pipeline"])
             self.optimizers.load_optimizers(loaded_state["optimizers"])
@@ -312,7 +324,7 @@ class Trainer:
         )
 
     @profiler.time_function
-    def train_iteration(self, step: int) -> Dict[str, torch.Tensor]:
+    def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
         """Run one iteration with a batch of inputs. Returns dictionary of model losses.
 
         Args:

nerfactory/viewer/server/visualizer.py

@@ -18,7 +18,7 @@
 import logging
 import signal
 import sys
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 import msgpack
 import msgpack_numpy
@@ -116,7 +116,7 @@ class Viewer:
     def __repr__(self):
         return f"<Viewer using: {self.window} at path: {self.path}>"
 
-    def write(self, data: Optional[Dict] = None):
+    def write(self, data: Union[Dict, str, None] = None):
         """Write data."""
         path = self.path.lower()
         return self.window.send({"type": "write", "path": path, "data": data})

pyproject.toml

@@ -17,7 +17,6 @@ dependencies = [
     "aiohttp==3.8.1",
     "aiortc==1.3.2",
     "av==9.2.0",
-    "black[jupyter]==22.3.0",
     "dcargs>=0.3.7",
     "gdown==4.5.1",
     "ninja==1.10.2.3",
@@ -28,36 +27,38 @@ dependencies = [
     "imageio==2.21.1",
     "ipywidgets>=7.6",
     "jupyterlab==3.3.4",
-    "kaleido==0.2.1",
     "matplotlib==3.5.3",
     "mediapy==1.1.0",
     "msgpack==1.0.4",
     "msgpack_numpy==0.4.8",
     "nerfacc==0.0.9",
-    "numpy-quaternion",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",
     "protobuf==3.20.0",
-    "pylint==2.13.4",
-    "pytest==7.1.2",
-    "pytest-xdist==2.5.0",
     "python-socketio==5.7.1",
     "rich==12.5.1",
     "tensorboard==2.9.0",
     "torch==1.12.1",
-    "torchmetrics[image]==0.8.2",
+    "torchmetrics[image]>=0.9.3",
     "torchtyping>=0.1.4",
     "torchvision>=0.13.0",
-    "tqdm==4.64.0",
     "u-msgpack-python>=2.4.1",
-    "wandb==0.13.3"
+    "wandb>=0.13.3"
 ]
 
 [project.urls]
 "Documentation" = "https://plenoptix-nerfactory.readthedocs-hosted.com/en/latest/?badge=latest"
 
 [project.optional-dependencies]
-# for building the docs
+# Development packages
+dev = [
+    "black[jupyter]==22.3.0",
+    "pylint==2.13.4",
+    "pytest==7.1.2",
+    "pytest-xdist==2.5.0",
+]
+
+# Documentation related packages
 docs = [
     "furo==2022.4.7",
     "readthedocs-sphinx-search==0.1.2",
@@ -91,6 +92,7 @@ max-line-length = 120
 generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"]
 good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$"
 ignore-paths = ["scripts/colmap2nerf.py", "^nerfactory/cuda/.*$", "^tests/cuda/.*$"]
+jobs = 0
 
 disable = [
   "duplicate-code",