Unverified Commit 3b19d6fc authored by Ren Pang's avatar Ren Pang Committed by GitHub
Browse files

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



* 'make_dataset' as staticmethod of 'DatasetFolder'

* a better fix
Co-authored-by: default avatarFrancisco Massa <fvsmassa@gmail.com>
parent ae87c1e4
...@@ -124,7 +124,7 @@ class DatasetFolder(VisionDataset): ...@@ -124,7 +124,7 @@ class DatasetFolder(VisionDataset):
super(DatasetFolder, self).__init__(root, transform=transform, super(DatasetFolder, self).__init__(root, transform=transform,
target_transform=target_transform) target_transform=target_transform)
classes, class_to_idx = self._find_classes(self.root) 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: if len(samples) == 0:
msg = "Found 0 files in subfolders of: {}\n".format(self.root) msg = "Found 0 files in subfolders of: {}\n".format(self.root)
if extensions is not None: if extensions is not None:
...@@ -139,6 +139,15 @@ class DatasetFolder(VisionDataset): ...@@ -139,6 +139,15 @@ class DatasetFolder(VisionDataset):
self.samples = samples self.samples = samples
self.targets = [s[1] for s in 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]]: def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
""" """
Finds the class folders in a dataset. Finds the class folders in a dataset.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment