docstring for ArtDataset
| 71 | |
| 72 | |
| 73 | class ImageFolder(Dataset): |
| 74 | """docstring for ArtDataset""" |
| 75 | def __init__(self, root, transform=None): |
| 76 | super( ImageFolder, self).__init__() |
| 77 | self.root = root |
| 78 | |
| 79 | self.frame = self._parse_frame() |
| 80 | self.transform = transform |
| 81 | |
| 82 | def _parse_frame(self): |
| 83 | frame = [] |
| 84 | img_names = os.listdir(self.root) |
| 85 | img_names.sort() |
| 86 | for i in range(len(img_names)): |
| 87 | image_path = os.path.join(self.root, img_names[i]) |
| 88 | if image_path[-4:] == '.jpg' or image_path[-4:] == '.png' or image_path[-5:] == '.jpeg': |
| 89 | frame.append(image_path) |
| 90 | return frame |
| 91 | |
| 92 | def __len__(self): |
| 93 | return len(self.frame) |
| 94 | |
| 95 | def __getitem__(self, idx): |
| 96 | file = self.frame[idx] |
| 97 | img = Image.open(file).convert('RGB') |
| 98 | |
| 99 | if self.transform: |
| 100 | img = self.transform(img) |
| 101 | |
| 102 | return img |
| 103 | |
| 104 | |
| 105 |
no outgoing calls
no test coverage detected