[CODEMOD][pytorch] replace uses of np.ndarray with npt.NDArray

Summary:
X-link: pytorch/opacus#681

X-link: pytorch/captum#1389

X-link: pytorch/botorch#2586

X-link: pytorch/audio#3846

This replaces uses of `numpy.ndarray` in type annotations with `numpy.typing.NDArray`. In Numpy-1.24.0+ `numpy.ndarray` is annotated as generic type. Without template parameters it triggers static analysis errors:
```counterexample
Generic type `ndarray` expects 2 type parameters.
```
`numpy.typing.NDArray` is an alias that provides default template parameters.

Reviewed By: ryanthomasjohnson

Differential Revision: D64619891

fbshipit-source-id: dffc096b1ce90d11e73d475f0bbcb8867ed9ef01

references/depth/stereo/train.py

@@ -5,6 +5,7 @@
 from typing import List, Union
 
 import numpy as np
+import numpy.typing as npt
 import torch
 import torch.distributed as dist
 import torchvision.models.optical_flow
@@ -33,7 +34,7 @@ def make_stereo_flow(flow: Union[torch.Tensor, List[torch.Tensor]], model_out_ch
     return flow
 
 
-def make_lr_schedule(args: argparse.Namespace, optimizer: torch.optim.Optimizer) -> np.ndarray:
+def make_lr_schedule(args: argparse.Namespace, optimizer: torch.optim.Optimizer) -> npt.NDArray:
     """Helper function to return a learning rate scheduler for CRE-stereo"""
     if args.decay_after_steps < args.warmup_steps:
         raise ValueError(f"decay_after_steps: {args.function} must be greater than warmup_steps: {args.warmup_steps}")

references/depth/stereo/visualization.py

@@ -2,6 +2,7 @@
 from typing import List
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from torch import Tensor
 from torchvision.utils import make_grid
@@ -64,7 +65,7 @@ def make_training_sample_grid(
     disparities: Tensor,
     masks: Tensor,
     predictions: List[Tensor],
-) -> np.ndarray:
+) -> npt.NDArray:
     # detach images and renormalize to [0, 1]
     images_left = left_images.detach().cpu() * 0.5 + 0.5
     images_right = right_images.detach().cpu() * 0.5 + 0.5
@@ -84,7 +85,7 @@ def make_training_sample_grid(
 
 
 @torch.no_grad()
-def make_disparity_sequence_grid(predictions: List[Tensor], disparities: Tensor) -> np.ndarray:
+def make_disparity_sequence_grid(predictions: List[Tensor], disparities: Tensor) -> npt.NDArray:
     # right most we will be adding the ground truth
     seq_len = len(predictions) + 1
     predictions = list(map(lambda x: x[:, :1, :, :].detach().cpu(), predictions + [disparities]))

test/datasets_utils.py

@@ -18,6 +18,7 @@
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 
 import PIL
 import PIL.Image
@@ -825,8 +826,8 @@ def size(idx: int) -> Tuple[int, int, int]:
 def shape_test_for_stereo(
     left: PIL.Image.Image,
     right: PIL.Image.Image,
-    disparity: Optional[np.ndarray] = None,
-    valid_mask: Optional[np.ndarray] = None,
+    disparity: Optional[npt.NDArray] = None,
+    valid_mask: Optional[npt.NDArray] = None,
 ):
     left_dims = get_dimensions(left)
     right_dims = get_dimensions(right)

torchvision/datasets/_optical_flow.py

@@ -6,6 +6,7 @@
 from typing import Callable, List, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from PIL import Image
 
@@ -164,7 +165,7 @@ def __getitem__(self, index: int) -> Union[T1, T2]:
         """
         return super().__getitem__(index)
 
-    def _read_flow(self, file_name: str) -> np.ndarray:
+    def _read_flow(self, file_name: str) -> npt.NDArray:
         return _read_flo(file_name)
 
 
@@ -225,7 +226,7 @@ def __getitem__(self, index: int) -> Union[T1, T2]:
         """
         return super().__getitem__(index)
 
-    def _read_flow(self, file_name: str) -> Tuple[np.ndarray, np.ndarray]:
+    def _read_flow(self, file_name: str) -> Tuple[npt.NDArray, npt.NDArray]:
         return _read_16bits_png_with_flow_and_valid_mask(file_name)
 
 
@@ -293,7 +294,7 @@ def __getitem__(self, index: int) -> Union[T1, T2]:
         """
         return super().__getitem__(index)
 
-    def _read_flow(self, file_name: str) -> np.ndarray:
+    def _read_flow(self, file_name: str) -> npt.NDArray:
         return _read_flo(file_name)
 
 
@@ -391,7 +392,7 @@ def __getitem__(self, index: int) -> Union[T1, T2]:
         """
         return super().__getitem__(index)
 
-    def _read_flow(self, file_name: str) -> np.ndarray:
+    def _read_flow(self, file_name: str) -> npt.NDArray:
         return _read_pfm(file_name)
 
 
@@ -443,7 +444,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt
                 "Could not find the HD1K images. Please make sure the directory structure is correct."
             )
 
-    def _read_flow(self, file_name: str) -> Tuple[np.ndarray, np.ndarray]:
+    def _read_flow(self, file_name: str) -> Tuple[npt.NDArray, npt.NDArray]:
         return _read_16bits_png_with_flow_and_valid_mask(file_name)
 
     def __getitem__(self, index: int) -> Union[T1, T2]:
@@ -462,7 +463,7 @@ def __getitem__(self, index: int) -> Union[T1, T2]:
         return super().__getitem__(index)
 
 
-def _read_flo(file_name: str) -> np.ndarray:
+def _read_flo(file_name: str) -> npt.NDArray:
     """Read .flo file in Middlebury format"""
     # Code adapted from:
     # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
@@ -479,7 +480,7 @@ def _read_flo(file_name: str) -> np.ndarray:
         return data.reshape(h, w, 2).transpose(2, 0, 1)
 
 
-def _read_16bits_png_with_flow_and_valid_mask(file_name: str) -> Tuple[np.ndarray, np.ndarray]:
+def _read_16bits_png_with_flow_and_valid_mask(file_name: str) -> Tuple[npt.NDArray, npt.NDArray]:
 
     flow_and_valid = decode_png(read_file(file_name)).to(torch.float32)
     flow, valid_flow_mask = flow_and_valid[:2, :, :], flow_and_valid[2, :, :]

torchvision/datasets/_stereo_matching.py

@@ -9,6 +9,7 @@
 from typing import Callable, cast, List, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 from PIL import Image
 
 from .utils import _read_pfm, download_and_extract_archive, verify_str_arg
@@ -92,7 +93,7 @@ def _scan_pairs(
         return paths
 
     @abstractmethod
-    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
+    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], Optional[npt.NDArray]]:
         # function that returns a disparity map and an occlusion map
         pass
 
@@ -178,7 +179,7 @@ def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None
         disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
         self._disparities = disparities
 
-    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]:
+    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]:
         disparity_map = _read_pfm_file(file_path)
         disparity_map = np.abs(disparity_map)  # ensure that the disparity is positive
         valid_mask = None
@@ -257,7 +258,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt
         else:
             self._disparities = list((None, None) for _ in self._images)
 
-    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]:
+    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], None]:
         # test split has no disparity maps
         if file_path is None:
             return None, None
@@ -345,7 +346,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt
         else:
             self._disparities = list((None, None) for _ in self._images)
 
-    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]:
+    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], None]:
         # test split has no disparity maps
         if file_path is None:
             return None, None
@@ -565,7 +566,7 @@ def _read_img(self, file_path: Union[str, Path]) -> Image.Image:
             file_path = random.choice(ambient_file_paths)  # type: ignore
         return super()._read_img(file_path)
 
-    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]:
+    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]:
         # test split has not disparity maps
         if file_path is None:
             return None, None
@@ -695,7 +696,7 @@ def __init__(
             disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
             self._disparities += disparities
 
-    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]:
+    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]:
         disparity_map = np.asarray(Image.open(file_path), dtype=np.float32)
         # unsqueeze the disparity map into (C, H, W) format
         disparity_map = disparity_map[None, :, :] / 32.0
@@ -789,7 +790,7 @@ def __init__(self, root: Union[str, Path], variant: str = "single", transforms:
             right_disparity_pattern = str(root / s / split_prefix[s] / "*.right.depth.png")
             self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
 
-    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]:
+    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]:
         # (H, W) image
         depth = np.asarray(Image.open(file_path))
         # as per https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt
@@ -912,7 +913,7 @@ def __init__(
             right_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "right" / "*.pfm")
             self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
 
-    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]:
+    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]:
         disparity_map = _read_pfm_file(file_path)
         disparity_map = np.abs(disparity_map)  # ensure that the disparity is positive
         valid_mask = None
@@ -1021,7 +1022,7 @@ def _get_occlussion_mask_paths(self, file_path: str) -> Tuple[str, str]:
 
         return occlusion_path, outofframe_path
 
-    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]:
+    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]:
         if file_path is None:
             return None, None
 
@@ -1102,7 +1103,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt
         right_disparity_pattern = str(root / "*" / "right_disp.png")
         self._disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
 
-    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]:
+    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]:
         disparity_map = np.asarray(Image.open(file_path), dtype=np.float32)
         # unsqueeze disparity to (C, H, W)
         disparity_map = disparity_map[None, :, :] / 1024.0
@@ -1196,7 +1197,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt
             disparity_pattern = str(root / anot_dir / "*" / "disp0GT.pfm")
             self._disparities = self._scan_pairs(disparity_pattern, None)
 
-    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]:
+    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]:
         # test split has no disparity maps
         if file_path is None:
             return None, None

torchvision/datasets/phototour.py

@@ -3,6 +3,7 @@
 from typing import Any, Callable, List, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from PIL import Image
 
@@ -187,7 +188,7 @@ def extra_repr(self) -> str:
 def read_image_file(data_dir: str, image_ext: str, n: int) -> torch.Tensor:
     """Return a Tensor containing the patches"""
 
-    def PIL2array(_img: Image.Image) -> np.ndarray:
+    def PIL2array(_img: Image.Image) -> npt.NDArray:
         """Convert PIL image type to numpy 2D array"""
         return np.array(_img.getdata(), dtype=np.uint8).reshape(64, 64)
 

torchvision/datasets/sbd.py

@@ -4,6 +4,7 @@
 from typing import Any, Callable, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 from PIL import Image
 
 from .utils import download_and_extract_archive, download_url, verify_str_arg
@@ -102,7 +103,7 @@ def _get_segmentation_target(self, filepath: str) -> Image.Image:
         mat = self._loadmat(filepath)
         return Image.fromarray(mat["GTcls"][0]["Segmentation"][0])
 
-    def _get_boundaries_target(self, filepath: str) -> np.ndarray:
+    def _get_boundaries_target(self, filepath: str) -> npt.NDArray:
         mat = self._loadmat(filepath)
         return np.concatenate(
             [np.expand_dims(mat["GTcls"][0]["Boundaries"][0][i][0].toarray(), axis=0) for i in range(self.num_classes)],

torchvision/datasets/stl10.py

@@ -3,6 +3,7 @@
 from typing import Any, Callable, cast, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 from PIL import Image
 
 from .utils import check_integrity, download_and_extract_archive, verify_str_arg
@@ -63,7 +64,7 @@ def __init__(
             raise RuntimeError("Dataset not found or corrupted. You can use download=True to download it")
 
         # now load the picked numpy arrays
-        self.labels: Optional[np.ndarray]
+        self.labels: Optional[npt.NDArray]
         if self.split == "train":
             self.data, self.labels = self.__loadfile(self.train_list[0][0], self.train_list[1][0])
             self.labels = cast(np.ndarray, self.labels)
@@ -129,7 +130,7 @@ def __getitem__(self, index: int) -> Tuple[Any, Any]:
     def __len__(self) -> int:
         return self.data.shape[0]
 
-    def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    def __loadfile(self, data_file: str, labels_file: Optional[str] = None) -> Tuple[npt.NDArray, Optional[npt.NDArray]]:
         labels = None
         if labels_file:
             path_to_labels = os.path.join(self.root, self.base_folder, labels_file)

torchvision/datasets/utils.py

@@ -16,6 +16,7 @@
 from urllib.parse import urlparse
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from torch.utils.model_zoo import tqdm
 
@@ -434,7 +435,7 @@ def verify_str_arg(
     return value
 
 
-def _read_pfm(file_name: Union[str, pathlib.Path], slice_channels: int = 2) -> np.ndarray:
+def _read_pfm(file_name: Union[str, pathlib.Path], slice_channels: int = 2) -> npt.NDArray:
     """Read file in .pfm format. Might contain either 1 or 3 channels of data.
 
     Args:

torchvision/prototype/datasets/_builtin/cifar.py

@@ -5,6 +5,7 @@
 from typing import Any, BinaryIO, cast, Dict, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 from torchdata.datapipes.iter import Filter, IterDataPipe, Mapper
 from torchvision.prototype.datasets.utils import Dataset, HttpResource, OnlineResource
 from torchvision.prototype.datasets.utils._internal import (
@@ -24,7 +25,7 @@ def __init__(self, datapipe: IterDataPipe[Dict[str, Any]], *, labels_key: str) -
         self.datapipe = datapipe
         self.labels_key = labels_key
 
-    def __iter__(self) -> Iterator[Tuple[np.ndarray, int]]:
+    def __iter__(self) -> Iterator[Tuple[npt.NDArray, int]]:
         for mapping in self.datapipe:
             image_arrays = mapping["data"].reshape((-1, 3, 32, 32))
             category_idcs = mapping[self.labels_key]
@@ -67,7 +68,7 @@ def _unpickle(self, data: Tuple[str, io.BytesIO]) -> Dict[str, Any]:
         file.close()
         return content
 
-    def _prepare_sample(self, data: Tuple[np.ndarray, int]) -> Dict[str, Any]:
+    def _prepare_sample(self, data: Tuple[npt.NDArray, int]) -> Dict[str, Any]:
         image_array, category_idx = data
         return dict(
             image=Image(image_array),

torchvision/prototype/datasets/_builtin/svhn.py

@@ -2,6 +2,7 @@
 from typing import Any, BinaryIO, Dict, List, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 from torchdata.datapipes.iter import IterDataPipe, Mapper, UnBatcher
 from torchvision.prototype.datasets.utils import Dataset, HttpResource, OnlineResource
 from torchvision.prototype.datasets.utils._internal import hint_sharding, hint_shuffling, read_mat
@@ -50,7 +51,7 @@ def _resources(self) -> List[OnlineResource]:
 
         return [data]
 
-    def _read_images_and_labels(self, data: Tuple[str, BinaryIO]) -> List[Tuple[np.ndarray, np.ndarray]]:
+    def _read_images_and_labels(self, data: Tuple[str, BinaryIO]) -> List[Tuple[npt.NDArray, npt.NDArray]]:
         _, buffer = data
         content = read_mat(buffer)
         return list(
@@ -60,7 +61,7 @@ def _read_images_and_labels(self, data: Tuple[str, BinaryIO]) -> List[Tuple[np.n
             )
         )
 
-    def _prepare_sample(self, data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, Any]:
+    def _prepare_sample(self, data: Tuple[npt.NDArray, npt.NDArray]) -> Dict[str, Any]:
         image_array, label_array = data
 
         return dict(

torchvision/transforms/functional.py

@@ -6,6 +6,7 @@
 from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
+import numpy.typing as npt
 import torch
 from PIL import Image
 from PIL.Image import Image as PILImage
@@ -124,7 +125,7 @@ def _is_numpy_image(img: Any) -> bool:
     return img.ndim in {2, 3}
 
 
-def to_tensor(pic: Union[PILImage, np.ndarray]) -> Tensor:
+def to_tensor(pic: Union[PILImage, npt.NDArray]) -> Tensor:
     """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
     This function does not support torchscript.