More type annotations

Summary: More type annotations: device, shaders, pluggable I/O, stats in NeRF project, cameras, textures, etc...

Reviewed By: nikhilaravi

Differential Revision: D29327396

fbshipit-source-id: cdf0ceaaa010e22423088752688c8dd81f1acc3c

Author: patricklabatut

projects/nerf/nerf/dataset.py

@@ -136,7 +136,7 @@ def download_data(
     dataset_names: Optional[List[str]] = None,
     data_root: str = DEFAULT_DATA_ROOT,
     url_root: str = DEFAULT_URL_ROOT,
-):
+) -> None:
     """
     Downloads the relevant dataset files.
 

projects/nerf/nerf/stats.py

@@ -29,7 +29,7 @@ def __init__(self) -> None:
         self.history = []
         self.reset()
 
-    def reset(self):
+    def reset(self) -> None:
         """
         Reset the running average meter.
         """
@@ -38,7 +38,7 @@ def reset(self):
         self.sum = 0
         self.count = 0
 
-    def update(self, val: float, n: int = 1, epoch: int = 0):
+    def update(self, val: float, n: int = 1, epoch: int = 0) -> None:
         """
         Updates the average meter with a value `val`.
 
@@ -123,7 +123,7 @@ def __init__(
         self.plot_file = plot_file
         self.hard_reset(epoch=epoch)
 
-    def reset(self):
+    def reset(self) -> None:
         """
         Called before an epoch to clear current epoch buffers.
         """
@@ -138,7 +138,7 @@ def reset(self):
         # Set a new timestamp.
         self._epoch_start = time.time()
 
-    def hard_reset(self, epoch: int = -1):
+    def hard_reset(self, epoch: int = -1) -> None:
         """
         Erases all logged data.
         """
@@ -149,7 +149,7 @@ def hard_reset(self, epoch: int = -1):
         self.stats = {}
         self.reset()
 
-    def new_epoch(self):
+    def new_epoch(self) -> None:
         """
         Initializes a new epoch.
         """
@@ -166,7 +166,7 @@ def _gather_value(self, val):
             val = float(val.sum())
         return val
 
-    def update(self, preds: dict, stat_set: str = "train"):
+    def update(self, preds: dict, stat_set: str = "train") -> None:
         """
         Update the internal logs with metrics of a training step.
 
@@ -211,7 +211,7 @@ def update(self, preds: dict, stat_set: str = "train"):
             if val is not None:
                 self.stats[stat_set][stat].update(val, epoch=epoch, n=1)
 
-    def print(self, max_it: Optional[int] = None, stat_set: str = "train"):
+    def print(self, max_it: Optional[int] = None, stat_set: str = "train") -> None:
         """
         Print the current values of all stored stats.
 
@@ -247,7 +247,7 @@ def plot_stats(
         viz: Visdom = None,
         visdom_env: Optional[str] = None,
         plot_file: Optional[str] = None,
-    ):
+    ) -> None:
         """
         Plot the line charts of the history of the stats.
 

pytorch3d/io/experimental_gltf_io.py

@@ -42,14 +42,13 @@
 from collections import deque
 from enum import IntEnum
 from io import BytesIO
-from pathlib import Path
 from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union, cast
 
 import numpy as np
 import torch
 from iopath.common.file_io import PathManager
 from PIL import Image
-from pytorch3d.io.utils import _open_file
+from pytorch3d.io.utils import PathOrStr, _open_file
 from pytorch3d.renderer.mesh import TexturesBase, TexturesUV, TexturesVertex
 from pytorch3d.structures import Meshes, join_meshes_as_scene
 from pytorch3d.transforms import Transform3d, quaternion_to_matrix
@@ -498,7 +497,7 @@ def load(self, include_textures: bool) -> List[Tuple[Optional[str], Meshes]]:
 
 
 def load_meshes(
-    path: Union[str, Path],
+    path: PathOrStr,
     path_manager: PathManager,
     include_textures: bool = True,
 ) -> List[Tuple[Optional[str], Meshes]]:
@@ -544,7 +543,7 @@ def __init__(self) -> None:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         include_textures: bool,
         device,
         path_manager: PathManager,
@@ -566,7 +565,7 @@ def read(
     def save(
         self,
         data: Meshes,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         **kwargs,

pytorch3d/io/obj_io.py

@@ -18,7 +18,12 @@
 from PIL import Image
 from pytorch3d.common.types import Device
 from pytorch3d.io.mtl_io import load_mtl, make_mesh_texture_atlas
-from pytorch3d.io.utils import _check_faces_indices, _make_tensor, _open_file
+from pytorch3d.io.utils import (
+    PathOrStr,
+    _check_faces_indices,
+    _make_tensor,
+    _open_file,
+)
 from pytorch3d.renderer import TexturesAtlas, TexturesUV
 from pytorch3d.structures import Meshes, join_meshes_as_batch
 
@@ -213,7 +218,7 @@ def load_obj(
               None.
     """
     data_dir = "./"
-    if isinstance(f, (str, bytes, os.PathLike)):
+    if isinstance(f, (str, bytes, Path)):
         data_dir = os.path.dirname(f)
     if path_manager is None:
         path_manager = PathManager()
@@ -297,7 +302,7 @@ def __init__(self) -> None:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         include_textures: bool,
         device: Device,
         path_manager: PathManager,
@@ -322,7 +327,7 @@ def read(
     def save(
         self,
         data: Meshes,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         decimal_places: Optional[int] = None,
@@ -650,7 +655,7 @@ def _load_obj(
 
 
 def save_obj(
-    f: Union[str, os.PathLike],
+    f: PathOrStr,
     verts,
     faces,
     decimal_places: Optional[int] = None,

pytorch3d/io/off_io.py

@@ -13,13 +13,12 @@
 http://www.geomview.org/docs/html/OFF.html .
 """
 import warnings
-from pathlib import Path
 from typing import Optional, Tuple, Union, cast
 
 import numpy as np
 import torch
 from iopath.common.file_io import PathManager
-from pytorch3d.io.utils import _check_faces_indices, _open_file
+from pytorch3d.io.utils import PathOrStr, _check_faces_indices, _open_file
 from pytorch3d.renderer import TexturesAtlas, TexturesVertex
 from pytorch3d.structures import Meshes
 
@@ -424,7 +423,7 @@ def __init__(self) -> None:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         include_textures: bool,
         device,
         path_manager: PathManager,
@@ -460,7 +459,7 @@ def read(
     def save(
         self,
         data: Meshes,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         decimal_places: Optional[int] = None,

pytorch3d/io/pluggable_formats.py

@@ -5,10 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from pathlib import Path
+import pathlib
 from typing import Optional, Tuple, Union
 
 from iopath.common.file_io import PathManager
+from pytorch3d.common.types import Device
+from pytorch3d.io.utils import PathOrStr
 from pytorch3d.structures import Meshes, Pointclouds
 
 
@@ -20,14 +22,14 @@
 """
 
 
-def endswith(path, suffixes: Tuple[str, ...]) -> bool:
+def endswith(path: PathOrStr, suffixes: Tuple[str, ...]) -> bool:
     """
     Returns whether the path ends with one of the given suffixes.
     If `path` is not actually a path, returns True. This is useful
     for allowing interpreters to bypass inappropriate paths, but
     always accepting streams.
     """
-    if isinstance(path, Path):
+    if isinstance(path, pathlib.Path):
         return path.suffix.lower() in suffixes
     if isinstance(path, str):
         return path.lower().endswith(suffixes)
@@ -42,9 +44,9 @@ class MeshFormatInterpreter:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         include_textures: bool,
-        device,
+        device: Device,
         path_manager: PathManager,
         **kwargs,
     ) -> Optional[Meshes]:
@@ -68,7 +70,7 @@ def read(
     def save(
         self,
         data: Meshes,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         **kwargs,
@@ -96,7 +98,7 @@ class PointcloudFormatInterpreter:
     """
 
     def read(
-        self, path: Union[str, Path], device, path_manager: PathManager, **kwargs
+        self, path: PathOrStr, device: Device, path_manager: PathManager, **kwargs
     ) -> Optional[Pointclouds]:
         """
         Read the data from the specified file and return it as
@@ -117,7 +119,7 @@ def read(
     def save(
         self,
         data: Pointclouds,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         **kwargs,

pytorch3d/io/ply_io.py

@@ -15,13 +15,17 @@
 import warnings
 from collections import namedtuple
 from io import BytesIO, TextIOBase
-from pathlib import Path
 from typing import List, Optional, Tuple, Union, cast
 
 import numpy as np
 import torch
 from iopath.common.file_io import PathManager
-from pytorch3d.io.utils import _check_faces_indices, _make_tensor, _open_file
+from pytorch3d.io.utils import (
+    PathOrStr,
+    _check_faces_indices,
+    _make_tensor,
+    _open_file,
+)
 from pytorch3d.renderer import TexturesVertex
 from pytorch3d.structures import Meshes, Pointclouds
 
@@ -1237,7 +1241,7 @@ def __init__(self) -> None:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         include_textures: bool,
         device,
         path_manager: PathManager,
@@ -1269,7 +1273,7 @@ def read(
     def save(
         self,
         data: Meshes,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         decimal_places: Optional[int] = None,
@@ -1318,7 +1322,7 @@ def __init__(self) -> None:
 
     def read(
         self,
-        path: Union[str, Path],
+        path: PathOrStr,
         device,
         path_manager: PathManager,
         **kwargs,
@@ -1339,7 +1343,7 @@ def read(
     def save(
         self,
         data: Pointclouds,
-        path: Union[str, Path],
+        path: PathOrStr,
         path_manager: PathManager,
         binary: Optional[bool],
         decimal_places: Optional[int] = None,

pytorch3d/io/utils.py

@@ -7,7 +7,7 @@
 import contextlib
 import pathlib
 import warnings
-from typing import IO, ContextManager, Optional
+from typing import IO, ContextManager, Optional, Union
 
 import numpy as np
 import torch
@@ -25,6 +25,9 @@ def nullcontext(x):
     yield x
 
 
+PathOrStr = Union[pathlib.Path, str]
+
+
 def _open_file(f, path_manager: PathManager, mode="r") -> ContextManager[IO]:
     if isinstance(f, str):
         f = path_manager.open(f, mode)