Tancik/lint and format scripts (#140)

* Lint and formatter to include scripts folder

.github/workflows/code_checks.yml

@@ -28,10 +28,10 @@ jobs:
         run: |
           ./scripts/licensing/license_headers.sh --check
       - name: Run Black
-        run: black pyrad/ notebooks/ --check
+        run: black pyrad/ scripts/ notebooks/ --check
       - name: Python Pylint
         run: |
-          pylint pyrad tests
+          pylint pyrad tests scripts
       - name: Test with pytest
         run: |
           pytest

pyproject.toml

@@ -16,6 +16,7 @@ ignore-paths = ["^pyrad/cuda/.*$", "^tests/cuda/.*$", "^pyrad/viewer/.*$", "^tes
 disable = [
   "too-many-arguments",
   "fixme",
+  "too-many-branches",
   "too-many-locals",
   "too-many-instance-attributes",
   "duplicate-code",

scripts/benchmarking/run_benchmark.py

@@ -90,5 +90,5 @@ if __name__ == "__main__":
     parser.add_argument("-g", "--graph", type=str, required=True, help="name of nerf graph to benchmark")
     parser.add_argument("-d", "--benchmark_date", type=str, required=True, help="timestamp of run to benchmark")
     parser.add_argument("-o", "--hydra_base_dir", type=str, default="outputs/", help="hydra base output path")
-    args = parser.parse_args()
-    main(args)
+    parsed_args = parser.parse_args()
+    main(parsed_args)

scripts/run_data_preprocessor.py

@@ -4,8 +4,8 @@ Script to process different dataset types into our format.
 
 # TODO(ethan): we should write a script here to convert from other datasets to our formats
 # e.g., Blender, LLFF, ShapeNet, DTU
+from typing import List
 from hydra import compose, initialize
-import json
 
 from pyrad.data.utils import get_dataset_inputs
 from pyrad.utils.io import get_absolute_path, make_dir, write_to_pkl
@@ -13,10 +13,20 @@ from pyrad.utils.misc import get_hash_str_from_dict
 
 
 def save_dataset_inputs_to_cache(
-    dataset_name="friends_TBBT-big_living_room", splits=["train", "val", "test"], downscale_factors=[1, 2, 4]
+    dataset_name: str = "friends_TBBT-big_living_room", splits: List = None, downscale_factors: List = None
 ):
     """Save the datasets input to cache.
-    TODO(ethan): use a serializable representation that's cleaner instead of pickle."""
+    TODO(ethan): use a serializable representation that's cleaner instead of pickle.
+
+    Args:
+        dataset_name (str): name of the dataset
+        splits (List[str]): list of splits to save. Defaults to ["train", "val", "test"]
+        downscale_factors (List[int]): list of downscale factors to save. Defaults to [1, 2, 4]
+    """
+    if splits is None:
+        splits = ["train", "val", "test"]
+    if downscale_factors is None:
+        downscale_factors = [1, 2, 4]
     with initialize(version_base="1.2", config_path="../configs"):
         config = compose(config_name="default_setup.yaml", overrides=[f"data/dataset={dataset_name}"])
     for downscale_factor in downscale_factors:

scripts/run_eval.py

@@ -5,7 +5,6 @@ import argparse
 import os
 from typing import Tuple
 
-import hydra
 import mediapy as media
 import torch
 from hydra import compose, initialize
@@ -13,7 +12,6 @@ from omegaconf import DictConfig
 from tqdm import tqdm
 
 from pyrad.cameras.camera_paths import get_spiral_path
-from pyrad.cameras.cameras import SimplePinholeCamera
 from pyrad.data.dataloader import setup_dataset_eval, setup_dataset_train
 from pyrad.graphs.base import Graph, setup_graph
 from pyrad.utils.writer import TimeWriter
@@ -90,6 +88,14 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1)
 def create_spiral_video(
     config: DictConfig, local_rank: int = 0, world_size: int = 1, output_filename: str = None
 ) -> None:
+    """Helper function to create a video of the spiral trajectory.
+
+    Args:
+        config (DictConfig): Configuration for loading the evaluation.
+        local_rank (int): Local rank of the process.
+        world_size (int): Total number of GPUs.
+        output_filename (str): Name of the output file.
+    """
     print("Creating spiral video")
     device = "cpu" if world_size == 0 else f"cuda:{local_rank}"
     # setup graph and dataset

scripts/run_train.py

@@ -93,13 +93,13 @@ def _distributed_worker(
         logger.error("Process group URL: %s", dist_url)
         raise e
 
-    assert comms._LOCAL_PROCESS_GROUP is None
+    assert comms._LOCAL_PROCESS_GROUP is None  # pylint: disable=protected-access
     num_machines = world_size // num_gpus_per_machine
     for i in range(num_machines):
         ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine))
         pg = dist.new_group(ranks_on_i)
         if i == machine_rank:
-            comms._LOCAL_PROCESS_GROUP = pg
+            comms._LOCAL_PROCESS_GROUP = pg  # pylint: disable=protected-access
 
     assert num_gpus_per_machine <= torch.cuda.device_count()
     torch.cuda.set_device(local_rank)
@@ -229,4 +229,4 @@ def main(config: DictConfig):
 
 
 if __name__ == "__main__":
-    main()
+    main()  # pylint: disable=no-value-for-parameter

scripts/run_viewer_bridge_server.py

@@ -1,9 +1,11 @@
+"""View Bridge Server"""
+
 import argparse
 from pyrad.viewer.server.server import ZMQWebSocketBridge
 
 
 def main():
-
+    """Run the viewer bridge server"""
     parser = argparse.ArgumentParser(description="Listen for ZeroMQ commands")
     parser.add_argument("--zmq-url", "-z", type=str, nargs="?", default=None)
     parser.add_argument("--websocket-port", "-wp", type=str, nargs="?", default=None)