pyre strict in dataset and io

Differential Revision: D73199680

fbshipit-source-id: 39e745f4669fe1c84fda951ec879160f80eb771c

Author: Pau Gargallo

opensfm/dataset.py

@@ -1,4 +1,4 @@
-# pyre-unsafe
+# pyre-strict
 import gzip
 import json
 import logging
@@ -8,6 +8,7 @@
 from typing import Any, Dict, IO, List, Optional, Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 from opensfm import config, features, geo, io, masking, pygeometry, pymap, rig, types
 from opensfm.dataset_base import DataSetBase
 from PIL.PngImagePlugin import PngImageFile
@@ -30,14 +31,17 @@ class DataSet(DataSetBase):
     """
 
     io_handler: io.IoFilesystemBase = io.IoFilesystemDefault()
-    config = None
+    config: Dict[str, Any] = {}
     image_files: Dict[str, str] = {}
     mask_files: Dict[str, str] = {}
     image_list: List[str] = []
 
-    def __init__(self, data_path: str, io_handler=io.IoFilesystemDefault) -> None:
+    def __init__(
+        self, data_path: str, io_handler: Optional[io.IoFilesystemBase] = None
+    ) -> None:
         """Init dataset associated to a folder."""
-        self.io_handler = io_handler
+        if io_handler is not None:
+            self.io_handler = io_handler
         self.data_path = data_path
         self.load_config()
         self.load_image_list()
@@ -80,7 +84,7 @@ def _image_file(self, image: str) -> str:
         """Path to the image file."""
         return self.image_files[image]
 
-    def open_image_file(self, image: str) -> IO[Any]:
+    def open_image_file(self, image: str) -> IO[bytes]:
         """Open image file and return file object."""
         return self.io_handler.open_rb(self._image_file(image))
 
@@ -90,7 +94,7 @@ def load_image(
         unchanged: bool = False,
         anydepth: bool = False,
         grayscale: bool = False,
-    ) -> np.ndarray:
+    ) -> NDArray:
         """Load image pixels as numpy array.
 
         The array is 3D, indexed by y-coord, x-coord, channel.
@@ -117,7 +121,7 @@ def load_mask_list(self) -> None:
         else:
             self._set_mask_path(os.path.join(self.data_path, "masks"))
 
-    def load_mask(self, image: str) -> Optional[np.ndarray]:
+    def load_mask(self, image: str) -> Optional[NDArray]:
         """Load image mask if it exists, otherwise return None."""
         if image in self.mask_files:
             mask_path = self.mask_files[image]
@@ -138,7 +142,7 @@ def _instances_path(self) -> str:
     def _instances_file(self, image: str) -> str:
         return os.path.join(self._instances_path(), image + ".png")
 
-    def load_instances(self, image: str) -> Optional[np.ndarray]:
+    def load_instances(self, image: str) -> Optional[NDArray]:
         """Load image instances file if it exists, otherwise return None."""
         instances_file = self._instances_file(image)
         if self.io_handler.isfile(instances_file):
@@ -153,10 +157,10 @@ def _segmentation_path(self) -> str:
     def _segmentation_file(self, image: str) -> str:
         return os.path.join(self._segmentation_path(), image + ".png")
 
-    def segmentation_labels(self) -> List[Any]:
+    def segmentation_labels(self) -> List[Dict[str, Any]]:
         return []
 
-    def load_segmentation(self, image: str) -> Optional[np.ndarray]:
+    def load_segmentation(self, image: str) -> Optional[NDArray]:
         """Load image segmentation if it exists, otherwise return None."""
         segmentation_file = self._segmentation_file(image)
         if self.io_handler.isfile(segmentation_file):
@@ -318,12 +322,12 @@ def _words_file(self, image: str) -> str:
     def words_exist(self, image: str) -> bool:
         return self.io_handler.isfile(self._words_file(image))
 
-    def load_words(self, image: str) -> np.ndarray:
+    def load_words(self, image: str) -> NDArray:
         with self.io_handler.open_rb(self._words_file(image)) as f:
             s = np.load(f)
             return s["words"].astype(np.int32)
 
-    def save_words(self, image: str, words: np.ndarray) -> None:
+    def save_words(self, image: str, words: NDArray) -> None:
         with self.io_handler.open_wb(self._words_file(image)) as f:
             np.savez_compressed(f, words=words.astype(np.uint16))
 
@@ -338,7 +342,7 @@ def _matches_file(self, image: str) -> str:
     def matches_exists(self, image: str) -> bool:
         return self.io_handler.isfile(self._matches_file(image))
 
-    def load_matches(self, image: str) -> Dict[str, np.ndarray]:
+    def load_matches(self, image: str) -> Dict[str, NDArray]:
         # Prevent pickling of anything except what we strictly need
         # as 'pickle.load' is RCE-prone. Will raise on any class other
         # than the numpy ones we allow.
@@ -363,7 +367,7 @@ def find_class(self, module, name):
             matches = MatchingUnpickler(BytesIO(gzip.decompress(fin.read()))).load()
         return matches
 
-    def save_matches(self, image: str, matches: Dict[str, np.ndarray]) -> None:
+    def save_matches(self, image: str, matches: Dict[str, NDArray]) -> None:
         self.io_handler.mkdir_p(self._matches_path())
 
         with BytesIO() as buffer:
@@ -372,7 +376,7 @@ def save_matches(self, image: str, matches: Dict[str, np.ndarray]) -> None:
             with self.io_handler.open_wb(self._matches_file(image)) as fw:
                 fw.write(buffer.getvalue())
 
-    def find_matches(self, im1: str, im2: str) -> np.ndarray:
+    def find_matches(self, im1: str, im2: str) -> NDArray:
         if self.matches_exists(im1):
             im1_matches = self.load_matches(im1)
             if im2 in im1_matches:
@@ -422,7 +426,7 @@ def save_reconstruction(
         self,
         reconstruction: List[types.Reconstruction],
         filename: Optional[str] = None,
-        minify=False,
+        minify: bool = False,
     ) -> None:
         with self.io_handler.open_wt(self._reconstruction_file(filename)) as fout:
             io.json_dump(io.reconstructions_to_json(reconstruction), fout, minify)
@@ -628,11 +632,11 @@ def save_ground_control_points(
         with self.io_handler.open_wt(self._ground_control_points_file()) as fout:
             io.write_ground_control_points(points, fout)
 
-    def image_as_array(self, image: str) -> np.ndarray:
+    def image_as_array(self, image: str) -> NDArray:
         logger.warning("image_as_array() is deprecated. Use load_image() instead.")
         return self.load_image(image)
 
-    def mask_as_array(self, image: str) -> Optional[np.ndarray]:
+    def mask_as_array(self, image: str) -> Optional[NDArray]:
         logger.warning("mask_as_array() is deprecated. Use load_mask() instead.")
         return self.load_mask(image)
 
@@ -707,18 +711,20 @@ class UndistortedDataSet:
     base: DataSetBase
     config: Dict[str, Any] = {}
     data_path: str
+    io_handler: io.IoFilesystemBase = io.IoFilesystemDefault()
 
     def __init__(
         self,
         base_dataset: DataSetBase,
         undistorted_data_path: str,
-        io_handler=io.IoFilesystemDefault,
+        io_handler: Optional[io.IoFilesystemBase] = None,
     ) -> None:
         """Init dataset associated to a folder."""
         self.base = base_dataset
         self.config = self.base.config
         self.data_path = undistorted_data_path
-        self.io_handler = io_handler
+        if io_handler is not None:
+            self.io_handler = io_handler
 
     def load_undistorted_shot_ids(self) -> Dict[str, List[str]]:
         filename = os.path.join(self.data_path, "undistorted_shot_ids.json")
@@ -738,11 +744,11 @@ def _undistorted_image_file(self, image: str) -> str:
         """Path of undistorted version of an image."""
         return os.path.join(self._undistorted_image_path(), image)
 
-    def load_undistorted_image(self, image: str) -> np.ndarray:
+    def load_undistorted_image(self, image: str) -> NDArray:
         """Load undistorted image pixels as a numpy array."""
         return self.io_handler.imread(self._undistorted_image_file(image))
 
-    def save_undistorted_image(self, image: str, array: np.ndarray) -> None:
+    def save_undistorted_image(self, image: str, array: NDArray) -> None:
         """Save undistorted image pixels."""
         self.io_handler.mkdir_p(self._undistorted_image_path())
         self.io_handler.imwrite(self._undistorted_image_file(image), array)
@@ -762,13 +768,13 @@ def undistorted_mask_exists(self, image: str) -> bool:
         """Check if the undistorted mask file exists."""
         return self.io_handler.isfile(self._undistorted_mask_file(image))
 
-    def load_undistorted_mask(self, image: str) -> np.ndarray:
+    def load_undistorted_mask(self, image: str) -> NDArray:
         """Load undistorted mask pixels as a numpy array."""
         return self.io_handler.imread(
             self._undistorted_mask_file(image), grayscale=True
         )
 
-    def save_undistorted_mask(self, image: str, array: np.ndarray) -> None:
+    def save_undistorted_mask(self, image: str, array: NDArray) -> None:
         """Save the undistorted image mask."""
         self.io_handler.mkdir_p(self._undistorted_mask_path())
         self.io_handler.imwrite(self._undistorted_mask_file(image), array)
@@ -784,7 +790,7 @@ def undistorted_segmentation_exists(self, image: str) -> bool:
         """Check if the undistorted segmentation file exists."""
         return self.io_handler.isfile(self._undistorted_segmentation_file(image))
 
-    def load_undistorted_segmentation(self, image: str) -> np.ndarray:
+    def load_undistorted_segmentation(self, image: str) -> NDArray:
         """Load an undistorted image segmentation."""
         segmentation_file = self._undistorted_segmentation_file(image)
         with self.io_handler.open_rb(segmentation_file) as fp:
@@ -804,12 +810,12 @@ def load_undistorted_segmentation(self, image: str) -> np.ndarray:
                 else:
                     raise IndexError
 
-    def save_undistorted_segmentation(self, image: str, array: np.ndarray) -> None:
+    def save_undistorted_segmentation(self, image: str, array: NDArray) -> None:
         """Save the undistorted image segmentation."""
         self.io_handler.mkdir_p(self._undistorted_segmentation_path())
         self.io_handler.imwrite(self._undistorted_segmentation_file(image), array)
 
-    def load_undistorted_segmentation_mask(self, image: str) -> Optional[np.ndarray]:
+    def load_undistorted_segmentation_mask(self, image: str) -> Optional[NDArray]:
         """Build a mask from the undistorted segmentation.
 
         The mask is non-zero only for pixels with segmentation
@@ -828,7 +834,7 @@ def load_undistorted_segmentation_mask(self, image: str) -> Optional[np.ndarray]
 
         return masking.mask_from_segmentation(segmentation, ignore_values)
 
-    def load_undistorted_combined_mask(self, image: str) -> Optional[np.ndarray]:
+    def load_undistorted_combined_mask(self, image: str) -> Optional[NDArray]:
         """Combine undistorted binary mask with segmentation mask.
 
         Return a mask that is non-zero only where the binary
@@ -854,16 +860,16 @@ def point_cloud_file(self, filename: str = "merged.ply") -> str:
 
     def load_point_cloud(
         self, filename: str = "merged.ply"
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    ) -> Tuple[NDArray, NDArray, NDArray, NDArray]:
         with self.io_handler.open_rt(self.point_cloud_file(filename)) as fp:
             return io.point_cloud_from_ply(fp)
 
     def save_point_cloud(
         self,
-        points: np.ndarray,
-        normals: np.ndarray,
-        colors: np.ndarray,
-        labels: np.ndarray,
+        points: NDArray,
+        normals: NDArray,
+        colors: NDArray,
+        labels: NDArray,
         filename: str = "merged.ply",
     ) -> None:
         self.io_handler.mkdir_p(self._depthmap_path())
@@ -876,11 +882,11 @@ def raw_depthmap_exists(self, image: str) -> bool:
     def save_raw_depthmap(
         self,
         image: str,
-        depth: np.ndarray,
-        plane: np.ndarray,
-        score: np.ndarray,
-        nghbr: np.ndarray,
-        nghbrs: np.ndarray,
+        depth: NDArray,
+        plane: NDArray,
+        score: NDArray,
+        nghbr: NDArray,
+        nghbrs: NDArray,
     ) -> None:
         self.io_handler.mkdir_p(self._depthmap_path())
         filepath = self.depthmap_file(image, "raw.npz")
@@ -891,7 +897,7 @@ def save_raw_depthmap(
 
     def load_raw_depthmap(
         self, image: str
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    ) -> Tuple[NDArray, NDArray, NDArray, NDArray, NDArray]:
         with self.io_handler.open_rb(self.depthmap_file(image, "raw.npz")) as f:
             o = np.load(f)
             return o["depth"], o["plane"], o["score"], o["nghbr"], o["nghbrs"]
@@ -900,16 +906,14 @@ def clean_depthmap_exists(self, image: str) -> bool:
         return self.io_handler.isfile(self.depthmap_file(image, "clean.npz"))
 
     def save_clean_depthmap(
-        self, image: str, depth: np.ndarray, plane: np.ndarray, score: np.ndarray
+        self, image: str, depth: NDArray, plane: NDArray, score: NDArray
     ) -> None:
         self.io_handler.mkdir_p(self._depthmap_path())
         filepath = self.depthmap_file(image, "clean.npz")
         with self.io_handler.open_wb(filepath) as f:
             np.savez_compressed(f, depth=depth, plane=plane, score=score)
 
-    def load_clean_depthmap(
-        self, image: str
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def load_clean_depthmap(self, image: str) -> Tuple[NDArray, NDArray, NDArray]:
         with self.io_handler.open_rb(self.depthmap_file(image, "clean.npz")) as f:
             o = np.load(f)
             return o["depth"], o["plane"], o["score"]
@@ -920,10 +924,10 @@ def pruned_depthmap_exists(self, image: str) -> bool:
     def save_pruned_depthmap(
         self,
         image: str,
-        points: np.ndarray,
-        normals: np.ndarray,
-        colors: np.ndarray,
-        labels: np.ndarray,
+        points: NDArray,
+        normals: NDArray,
+        colors: NDArray,
+        labels: NDArray,
     ) -> None:
         self.io_handler.mkdir_p(self._depthmap_path())
         filepath = self.depthmap_file(image, "pruned.npz")
@@ -938,7 +942,7 @@ def save_pruned_depthmap(
 
     def load_pruned_depthmap(
         self, image: str
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    ) -> Tuple[NDArray, NDArray, NDArray, NDArray]:
         with self.io_handler.open_rb(self.depthmap_file(image, "pruned.npz")) as f:
             o = np.load(f)
             return (