fix root type annotations for remaining datasets

Author: pmeier

torchvision/datasets/_optical_flow.py

@@ -13,7 +13,6 @@
 from .utils import _read_pfm, verify_str_arg
 from .vision import VisionDataset
 
-
 T1 = Tuple[Image.Image, Image.Image, Optional[np.ndarray], Optional[np.ndarray]]
 T2 = Tuple[Image.Image, Image.Image, Optional[np.ndarray]]
 
@@ -33,7 +32,7 @@ class FlowDataset(ABC, VisionDataset):
     # and it's up to whatever consumes the dataset to decide what valid_flow_mask should be.
     _has_builtin_flow_mask = False
 
-    def __init__(self, root: str, transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None:
 
         super().__init__(root=root)
         self.transforms = transforms
@@ -125,7 +124,7 @@ class Sintel(FlowDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         pass_name: str = "clean",
         transforms: Optional[Callable] = None,
@@ -191,7 +190,7 @@ class KittiFlow(FlowDataset):
 
     _has_builtin_flow_mask = True
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root=root, transforms=transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "test"))
@@ -256,7 +255,7 @@ class FlyingChairs(FlowDataset):
             return a built-in valid mask, such as :class:`~torchvision.datasets.KittiFlow`.
     """
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root=root, transforms=transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "val"))
@@ -329,7 +328,7 @@ class FlyingThings3D(FlowDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         pass_name: str = "clean",
         camera: str = "left",
@@ -419,7 +418,7 @@ class HD1K(FlowDataset):
 
     _has_builtin_flow_mask = True
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root=root, transforms=transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "test"))

torchvision/datasets/_stereo_matching.py

@@ -27,7 +27,7 @@ class StereoMatchingDataset(ABC, VisionDataset):
 
     _has_built_in_disparity_mask = False
 
-    def __init__(self, root: str, transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None:
         """
         Args:
             root(str): Root directory of the dataset.
@@ -163,7 +163,7 @@ class CarlaStereo(StereoMatchingDataset):
         transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
     """
 
-    def __init__(self, root: str, transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         root = Path(root) / "carla-highres"
@@ -240,7 +240,7 @@ class Kitti2012Stereo(StereoMatchingDataset):
 
     _has_built_in_disparity_mask = True
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "test"))
@@ -328,7 +328,7 @@ class Kitti2015Stereo(StereoMatchingDataset):
 
     _has_built_in_disparity_mask = True
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "test"))
@@ -480,7 +480,7 @@ class Middlebury2014Stereo(StereoMatchingDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         calibration: Optional[str] = "perfect",
         use_ambient_views: bool = False,
@@ -576,7 +576,7 @@ def _read_disparity(self, file_path: str) -> Union[Tuple[None, None], Tuple[np.n
         valid_mask = (disparity_map > 0).squeeze(0)  # mask out invalid disparities
         return disparity_map, valid_mask
 
-    def _download_dataset(self, root: str) -> None:
+    def _download_dataset(self, root: Union[str, Path]) -> None:
         base_url = "https://vision.middlebury.edu/stereo/data/scenes2014/zip"
         # train and additional splits have 2 different calibration settings
         root = Path(root) / "Middlebury2014"
@@ -675,7 +675,7 @@ class CREStereo(StereoMatchingDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         transforms: Optional[Callable] = None,
     ) -> None:
         super().__init__(root, transforms)
@@ -762,7 +762,7 @@ class FallingThingsStereo(StereoMatchingDataset):
         transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
     """
 
-    def __init__(self, root: str, variant: str = "single", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], variant: str = "single", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         root = Path(root) / "FallingThings"
@@ -877,7 +877,7 @@ class SceneFlowStereo(StereoMatchingDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         variant: str = "FlyingThings3D",
         pass_name: str = "clean",
         transforms: Optional[Callable] = None,
@@ -980,7 +980,7 @@ class SintelStereo(StereoMatchingDataset):
 
     _has_built_in_disparity_mask = True
 
-    def __init__(self, root: str, pass_name: str = "final", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], pass_name: str = "final", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         verify_str_arg(pass_name, "pass_name", valid_values=("final", "clean", "both"))
@@ -1087,7 +1087,7 @@ class InStereo2k(StereoMatchingDataset):
         transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
     """
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         root = Path(root) / "InStereo2k" / split
@@ -1176,7 +1176,7 @@ class ETH3DStereo(StereoMatchingDataset):
 
     _has_built_in_disparity_mask = True
 
-    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None) -> None:
+    def __init__(self, root: Union[str, Path], split: str = "train", transforms: Optional[Callable] = None) -> None:
         super().__init__(root, transforms)
 
         verify_str_arg(split, "split", valid_values=("train", "test"))

torchvision/datasets/caltech.py

@@ -1,5 +1,6 @@
 import os
 import os.path
+from pathlib import Path
 from typing import Any, Callable, List, Optional, Tuple, Union
 
 from PIL import Image
@@ -38,7 +39,7 @@ class Caltech101(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         target_type: Union[List[str], str] = "category",
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,

torchvision/datasets/celeba.py

@@ -1,6 +1,7 @@
 import csv
 import os
 from collections import namedtuple
+from pathlib import Path
 from typing import Any, Callable, List, Optional, Tuple, Union
 
 import PIL
@@ -63,7 +64,7 @@ class CelebA(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         target_type: Union[List[str], str] = "attr",
         transform: Optional[Callable] = None,

torchvision/datasets/cifar.py

@@ -1,6 +1,7 @@
 import os.path
 import pickle
-from typing import Any, Callable, Optional, Tuple
+from pathlib import Path
+from typing import Any, Callable, Optional, Tuple, Union
 
 import numpy as np
 from PIL import Image
@@ -50,7 +51,7 @@ class CIFAR10(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         train: bool = True,
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,

torchvision/datasets/cityscapes.py

@@ -1,6 +1,7 @@
 import json
 import os
 from collections import namedtuple
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from PIL import Image
@@ -103,7 +104,7 @@ class Cityscapes(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         mode: str = "fine",
         target_type: Union[List[str], str] = "instance",

torchvision/datasets/clevr.py

@@ -1,6 +1,6 @@
 import json
 import pathlib
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 from PIL import Image
@@ -30,7 +30,7 @@ class CLEVRClassification(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, pathlib.Path],
         split: str = "train",
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,

torchvision/datasets/coco.py

@@ -1,5 +1,6 @@
 import os.path
-from typing import Any, Callable, List, Optional, Tuple
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 from PIL import Image
 
@@ -24,7 +25,7 @@ class CocoDetection(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         annFile: str,
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,

torchvision/datasets/country211.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 
 from .folder import ImageFolder
 from .utils import download_and_extract_archive, verify_str_arg
@@ -28,7 +28,7 @@ class Country211(ImageFolder):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, Path],
         split: str = "train",
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,

torchvision/datasets/dtd.py

@@ -1,6 +1,6 @@
 import os
 import pathlib
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, Union
 
 import PIL.Image
 
@@ -34,7 +34,7 @@ class DTD(VisionDataset):
 
     def __init__(
         self,
-        root: str,
+        root: Union[str, pathlib.Path],
         split: str = "train",
         partition: int = 1,
         transform: Optional[Callable] = None,