| import os | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| class PrecomputedDataset(Dataset): | |
| def __init__(self, | |
| root, | |
| transforms, | |
| student_augs, | |
| ): | |
| super(PrecomputedDataset, self).__init__() | |
| self.root = root | |
| self.transforms = transforms | |
| self.student_augs = student_augs | |
| self.image_files = [] | |
| self.label_files = [] | |
| self.pseudo_files = [] | |
| for file in os.listdir(os.path.join(self.root, 'imgs')): | |
| self.image_files.append(os.path.join(self.root, 'imgs', file)) | |
| self.label_files.append(os.path.join(self.root, 'gts', file)) | |
| self.pseudo_files.append(os.path.join(self.root, 'pseudos', file)) | |
| def __getitem__(self, index): | |
| image_path = self.image_files[index] | |
| label_path = self.label_files[index] | |
| pseudo_path = self.pseudo_files[index] | |
| img = Image.open(image_path).convert("RGB") | |
| label = Image.open(label_path) | |
| pseudo = Image.open(pseudo_path) | |
| if self.student_augs: | |
| img, label, aimg, pseudo = self.transforms(img, label, pseudo) | |
| return img, label.long(), aimg, pseudo.long() | |
| else: | |
| img, label, pseudo = self.transforms(img, label, pseudo) | |
| return img, label.long(), pseudo.long() | |
| def __len__(self): | |
| return len(self.image_files) |