Merge pull request #21 from ethanweber/tancik/type_check

Add type checking

.github/workflows/code_checks.yml

@@ -34,4 +34,4 @@ jobs:
         pylint mattport --max-line-length=120 --generated-members=numpy.*,torch.*,cv2.*,cv.* --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W --disable=too-many-arguments --disable=fixme
     - name: Test with pytest
       run: |
-        pytest tests --ignore tests/nerf/test_ray_generator.py
+        pytest tests --ignore tests/nerf/test_ray_generator.py --typeguard-packages=mattport --torchtyping-patch-typeguard

.vscode/settings.json

@@ -18,7 +18,9 @@
         "--src=${workspaceFolder}"
     ],
     "python.testing.pytestArgs": [
-        "tests"
+        "--ignore=tests/nerf/test_ray_generator.py",
+        "--typeguard-packages=mattport",
+        "--torchtyping-patch-typeguard"
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,

mattport/nerf/renderers.py

@@ -22,13 +22,13 @@ class Renderer(nn.Module):
 
     def forward(
         self,
-        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict["str", TensorType[..., "num_samples", -1]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         """Composite samples along ray and render image
 
         Args:
-            in_tensor (TensorType[..., "num_samples", "features"]): feature for each sample
+            in_tensor (TensorType[..., "num_samples", -1]): feature for each sample
             deltas (TensorType[..., "num_samples"]): depth of each sample
 
         Returns:
@@ -50,13 +50,13 @@ class RGB(Renderer):
 
     def forward(
         self,
-        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict[str, TensorType[..., "num_samples", -1]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         """Composite samples along ray and render color image
 
         Args:
-            in_tensor (TensorType[..., "num_samples", 3]): RGB for each sample
+            in_tensor (TensorType[..., "num_samples", -1]): RGB for each sample
             deltas (TensorType[..., "num_samples"]): depth influence for each sample
 
         Returns:
@@ -88,7 +88,7 @@ class Depth(Renderer):
 
     def forward(
         self,
-        field_output_dict: Dict["str", TensorType[..., "num_samples", "features"]],
+        field_output_dict: Dict["str", TensorType[..., "num_samples", -1]],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         raise NotImplementedError