Add annotations to vision/fair/pytorch3d

Reviewed By: shannonzhu

Differential Revision: D33970393

fbshipit-source-id: 9b4dfaccfc3793fd37705a923d689cb14c9d26ba

pytorch3d/ops/knn.py

@@ -106,7 +106,7 @@ def knn_points(
     version: int = -1,
     return_nn: bool = False,
     return_sorted: bool = True,
-):
+) -> _KNN:
     """
     K-Nearest neighbors on point clouds.
 

pytorch3d/ops/points_normals.py

@@ -166,7 +166,7 @@ def estimate_pointcloud_local_coord_frames(
     return curvatures, local_coord_frames
 
 
-def _disambiguate_vector_directions(pcl, knns, vecs):
+def _disambiguate_vector_directions(pcl, knns, vecs: float) -> float:
     """
     Disambiguates normal directions according to [1].
 
@@ -180,6 +180,7 @@ def _disambiguate_vector_directions(pcl, knns, vecs):
     # each element of the neighborhood
     df = knns - pcl[:, :, None]
     # projection of the difference on the principal direction
+    # pyre-fixme[16]: `float` has no attribute `__getitem__`.
     proj = (vecs[:, :, None] * df).sum(3)
     # check how many projections are positive
     n_pos = (proj > 0).type_as(knns).sum(2, keepdim=True)

pytorch3d/ops/points_to_volumes.py

@@ -479,7 +479,7 @@ def _check_points_to_volumes_inputs(
     volume_features: torch.Tensor,
     grid_sizes: torch.LongTensor,
     mask: Optional[torch.Tensor] = None,
-):
+) -> None:
 
     max_grid_size = grid_sizes.max(dim=0).values
     if torch.prod(max_grid_size) > volume_densities.shape[1]:

pytorch3d/ops/subdivide_meshes.py

@@ -400,7 +400,7 @@ def create_verts_index(verts_per_mesh, edges_per_mesh, device=None):
     return verts_idx
 
 
-def create_faces_index(faces_per_mesh, device=None):
+def create_faces_index(faces_per_mesh: int, device=None):
     """
     Helper function to group the faces indices for each mesh. New faces are
     stacked at the end of the original faces tensor, so in order to have
@@ -417,7 +417,9 @@ def create_faces_index(faces_per_mesh, device=None):
     """
     # e.g. faces_per_mesh = [2, 5, 3]
 
+    # pyre-fixme[16]: `int` has no attribute `sum`.
     F = faces_per_mesh.sum()  # e.g. 10
+    # pyre-fixme[16]: `int` has no attribute `cumsum`.
     faces_per_mesh_cumsum = faces_per_mesh.cumsum(dim=0)  # (N,) e.g. (2, 7, 10)
 
     switch1_idx = faces_per_mesh_cumsum.clone()

pytorch3d/ops/utils.py

@@ -150,7 +150,7 @@ def convert_pointclouds_to_tensor(pcl: Union[torch.Tensor, "Pointclouds"]):
     return X, num_points
 
 
-def is_pointclouds(pcl: Union[torch.Tensor, "Pointclouds"]):
+def is_pointclouds(pcl: Union[torch.Tensor, "Pointclouds"]) -> bool:
     """Checks whether the input `pcl` is an instance of `Pointclouds`
     by checking the existence of `points_padded` and `num_points_per_cloud`
     functions.