| 15 | class DatasetGivenLabels(Dataset): |
| 16 | """ A self-defined dataset with updated input labels""" |
| 17 | def __init__(self, root, labels=None, transform=None, suffix='.jpg'): |
| 18 | dir_list = [] |
| 19 | for dir_item in os.listdir(root): |
| 20 | if osp.isdir(osp.join(root, dir_item)): |
| 21 | dir_list.append(dir_item) |
| 22 | dir_list.sort() |
| 23 | |
| 24 | imagedirs = [] |
| 25 | gt = [] |
| 26 | for label_id, dir_item in enumerate(dir_list): |
| 27 | sub_folder = osp.join(root, dir_item) |
| 28 | imagedirs_item = listdir(sub_folder, suffix=suffix) |
| 29 | imagedirs += imagedirs_item |
| 30 | gt += [label_id for _ in range(len(imagedirs_item))] |
| 31 | |
| 32 | self.transform = transform |
| 33 | self.imgs = imagedirs |
| 34 | self.gt = gt |
| 35 | self.labels = labels |
| 36 | |
| 37 | def __getitem__(self, index): |
| 38 | with open(self.imgs[index], 'rb') as f: |