| 172 | #---------------------------------------------------------------------------- |
| 173 | |
| 174 | class ImageFolderDataset(Dataset): |
| 175 | def __init__(self, |
| 176 | path, # Path to directory or zip. |
| 177 | resolution = None, # Ensure specific resolution, None = highest available. |
| 178 | **super_kwargs, # Additional arguments for the Dataset base class. |
| 179 | ): |
| 180 | self._path = path |
| 181 | self._zipfile = None |
| 182 | |
| 183 | if os.path.isdir(self._path): |
| 184 | self._type = 'dir' |
| 185 | self._all_fnames = {os.path.relpath(os.path.join(root, fname), start=self._path) for root, _dirs, files in os.walk(self._path) for fname in files} |
| 186 | elif self._file_ext(self._path) == '.zip': |
| 187 | self._type = 'zip' |
| 188 | self._all_fnames = set(self._get_zipfile().namelist()) |
| 189 | else: |
| 190 | raise IOError('Path must point to a directory or zip') |
| 191 | |
| 192 | PIL.Image.init() |
| 193 | self._image_fnames = sorted(fname for fname in self._all_fnames if self._file_ext(fname) in PIL.Image.EXTENSION) |
| 194 | if len(self._image_fnames) == 0: |
| 195 | raise IOError('No image files found in the specified path') |
| 196 | |
| 197 | name = os.path.splitext(os.path.basename(self._path))[0] |
| 198 | raw_shape = [len(self._image_fnames)] + list(self._load_raw_image(0).shape) |
| 199 | if resolution is not None and (raw_shape[2] != resolution or raw_shape[3] != resolution): |
| 200 | raise IOError(f'Image files do not match the specified resolution. Resolution is {resolution}, shape is {raw_shape}') |
| 201 | super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) |
| 202 | |
| 203 | def _get_zipfile(self): |
| 204 | assert self._type == 'zip' |
| 205 | if self._zipfile is None: |
| 206 | self._zipfile = zipfile.ZipFile(self._path) |
| 207 | return self._zipfile |
| 208 | |
| 209 | def _open_file(self, fname): |
| 210 | if self._type == 'dir': |
| 211 | return open(os.path.join(self._path, fname), 'rb') |
| 212 | if self._type == 'zip': |
| 213 | return self._get_zipfile().open(fname, 'r') |
| 214 | return None |
| 215 | |
| 216 | def close(self): |
| 217 | try: |
| 218 | if self._zipfile is not None: |
| 219 | self._zipfile.close() |
| 220 | finally: |
| 221 | self._zipfile = None |
| 222 | |
| 223 | def __getstate__(self): |
| 224 | return dict(super().__getstate__(), _zipfile=None) |
| 225 | |
| 226 | def _load_raw_image(self, raw_idx): |
| 227 | fname = self._image_fnames[raw_idx] |
| 228 | |
| 229 | with self._open_file(fname) as f: |
| 230 | use_pyspng = pyspng is not None and self._file_ext(fname) == '.png' |
| 231 | image = load_image_from_buffer(f, use_pyspng=use_pyspng) |
nothing calls this directly
no outgoing calls
no test coverage detected