'make_dataset' as staticmethod of 'DatasetFolder' (#3215)

Summary: * 'make_dataset' as staticmethod of 'DatasetFolder' * a better fix Reviewed By: datumbox Differential Revision: D25954567 fbshipit-source-id: 514fde3bad4e27518a198276228a36c3217c2163 Co-authored-by: Francisco Massa <[email protected]>
pytorch · Jan 21, 2021 · d10758d · d10758d
1 parent 6a02f85
commit d10758d
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py
@@ -124,7 +124,7 @@ def __init__(
         super(DatasetFolder, self).__init__(root, transform=transform,
                                             target_transform=target_transform)
         classes, class_to_idx = self._find_classes(self.root)
-        samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file)
+        samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)
         if len(samples) == 0:
             msg = "Found 0 files in subfolders of: {}\n".format(self.root)
             if extensions is not None:
@@ -139,6 +139,15 @@ def __init__(
         self.samples = samples
         self.targets = [s[1] for s in samples]
 
+    @staticmethod
+    def make_dataset(
+        directory: str,
+        class_to_idx: Dict[str, int],
+        extensions: Optional[Tuple[str, ...]] = None,
+        is_valid_file: Optional[Callable[[str], bool]] = None,
+    ) -> List[Tuple[str, int]]:
+        return make_dataset(directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file)
+
     def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
         """
         Finds the class folders in a dataset.