|  | 
| 9 | 9 | from typing import Callable, cast, List, Optional, Tuple, Union | 
| 10 | 10 | 
 | 
| 11 | 11 | import numpy as np | 
| 12 |  | -import numpy.typing as npt | 
| 13 | 12 | from PIL import Image | 
| 14 | 13 | 
 | 
| 15 | 14 | from .utils import _read_pfm, download_and_extract_archive, verify_str_arg | 
| @@ -93,7 +92,7 @@ def _scan_pairs( | 
| 93 | 92 |         return paths | 
| 94 | 93 | 
 | 
| 95 | 94 |     @abstractmethod | 
| 96 |  | -    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], Optional[npt.NDArray]]: | 
|  | 95 | +    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: | 
| 97 | 96 |         # function that returns a disparity map and an occlusion map | 
| 98 | 97 |         pass | 
| 99 | 98 | 
 | 
| @@ -179,7 +178,7 @@ def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None | 
| 179 | 178 |         disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) | 
| 180 | 179 |         self._disparities = disparities | 
| 181 | 180 | 
 | 
| 182 |  | -    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]: | 
|  | 181 | +    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: | 
| 183 | 182 |         disparity_map = _read_pfm_file(file_path) | 
| 184 | 183 |         disparity_map = np.abs(disparity_map)  # ensure that the disparity is positive | 
| 185 | 184 |         valid_mask = None | 
| @@ -258,7 +257,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt | 
| 258 | 257 |         else: | 
| 259 | 258 |             self._disparities = list((None, None) for _ in self._images) | 
| 260 | 259 | 
 | 
| 261 |  | -    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], None]: | 
|  | 260 | +    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]: | 
| 262 | 261 |         # test split has no disparity maps | 
| 263 | 262 |         if file_path is None: | 
| 264 | 263 |             return None, None | 
| @@ -346,7 +345,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt | 
| 346 | 345 |         else: | 
| 347 | 346 |             self._disparities = list((None, None) for _ in self._images) | 
| 348 | 347 | 
 | 
| 349 |  | -    def _read_disparity(self, file_path: str) -> Tuple[Optional[npt.NDArray], None]: | 
|  | 348 | +    def _read_disparity(self, file_path: str) -> Tuple[Optional[np.ndarray], None]: | 
| 350 | 349 |         # test split has no disparity maps | 
| 351 | 350 |         if file_path is None: | 
| 352 | 351 |             return None, None | 
| @@ -566,7 +565,7 @@ def _read_img(self, file_path: Union[str, Path]) -> Image.Image: | 
| 566 | 565 |             file_path = random.choice(ambient_file_paths)  # type: ignore | 
| 567 | 566 |         return super()._read_img(file_path) | 
| 568 | 567 | 
 | 
| 569 |  | -    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]: | 
|  | 568 | +    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: | 
| 570 | 569 |         # test split has not disparity maps | 
| 571 | 570 |         if file_path is None: | 
| 572 | 571 |             return None, None | 
| @@ -696,7 +695,7 @@ def __init__( | 
| 696 | 695 |             disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) | 
| 697 | 696 |             self._disparities += disparities | 
| 698 | 697 | 
 | 
| 699 |  | -    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]: | 
|  | 698 | +    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: | 
| 700 | 699 |         disparity_map = np.asarray(Image.open(file_path), dtype=np.float32) | 
| 701 | 700 |         # unsqueeze the disparity map into (C, H, W) format | 
| 702 | 701 |         disparity_map = disparity_map[None, :, :] / 32.0 | 
| @@ -790,7 +789,7 @@ def __init__(self, root: Union[str, Path], variant: str = "single", transforms: | 
| 790 | 789 |             right_disparity_pattern = str(root / s / split_prefix[s] / "*.right.depth.png") | 
| 791 | 790 |             self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern) | 
| 792 | 791 | 
 | 
| 793 |  | -    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]: | 
|  | 792 | +    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: | 
| 794 | 793 |         # (H, W) image | 
| 795 | 794 |         depth = np.asarray(Image.open(file_path)) | 
| 796 | 795 |         # as per https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt | 
| @@ -913,7 +912,7 @@ def __init__( | 
| 913 | 912 |             right_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "right" / "*.pfm") | 
| 914 | 913 |             self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern) | 
| 915 | 914 | 
 | 
| 916 |  | -    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]: | 
|  | 915 | +    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: | 
| 917 | 916 |         disparity_map = _read_pfm_file(file_path) | 
| 918 | 917 |         disparity_map = np.abs(disparity_map)  # ensure that the disparity is positive | 
| 919 | 918 |         valid_mask = None | 
| @@ -1022,7 +1021,7 @@ def _get_occlussion_mask_paths(self, file_path: str) -> Tuple[str, str]: | 
| 1022 | 1021 | 
 | 
| 1023 | 1022 |         return occlusion_path, outofframe_path | 
| 1024 | 1023 | 
 | 
| 1025 |  | -    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]: | 
|  | 1024 | +    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: | 
| 1026 | 1025 |         if file_path is None: | 
| 1027 | 1026 |             return None, None | 
| 1028 | 1027 | 
 | 
| @@ -1103,7 +1102,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt | 
| 1103 | 1102 |         right_disparity_pattern = str(root / "*" / "right_disp.png") | 
| 1104 | 1103 |         self._disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern) | 
| 1105 | 1104 | 
 | 
| 1106 |  | -    def _read_disparity(self, file_path: str) -> Tuple[npt.NDArray, None]: | 
|  | 1105 | +    def _read_disparity(self, file_path: str) -> Tuple[np.ndarray, None]: | 
| 1107 | 1106 |         disparity_map = np.asarray(Image.open(file_path), dtype=np.float32) | 
| 1108 | 1107 |         # unsqueeze disparity to (C, H, W) | 
| 1109 | 1108 |         disparity_map = disparity_map[None, :, :] / 1024.0 | 
| @@ -1197,7 +1196,7 @@ def __init__(self, root: Union[str, Path], split: str = "train", transforms: Opt | 
| 1197 | 1196 |             disparity_pattern = str(root / anot_dir / "*" / "disp0GT.pfm") | 
| 1198 | 1197 |             self._disparities = self._scan_pairs(disparity_pattern, None) | 
| 1199 | 1198 | 
 | 
| 1200 |  | -    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[npt.NDArray, npt.NDArray]]: | 
|  | 1199 | +    def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.ndarray, np.ndarray]]: | 
| 1201 | 1200 |         # test split has no disparity maps | 
| 1202 | 1201 |         if file_path is None: | 
| 1203 | 1202 |             return None, None | 
|  | 
0 commit comments