(dir, classes_idx=None)
| 21 | |
| 22 | |
| 23 | def find_classes(dir, classes_idx=None): |
| 24 | classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] |
| 25 | classes.sort() |
| 26 | if classes_idx is not None: |
| 27 | assert type(classes_idx) == tuple |
| 28 | start, end = classes_idx |
| 29 | classes = classes[start:end] |
| 30 | class_to_idx = {classes[i]: i for i in range(len(classes))} |
| 31 | return classes, class_to_idx |
| 32 | |
| 33 | |
| 34 | def make_dataset(dir, class_to_idx): |