| 297 | |
| 298 | |
| 299 | class ImageCephDataset(data.Dataset): |
| 300 | |
| 301 | def __init__(self, |
| 302 | root, |
| 303 | split, |
| 304 | parser=None, |
| 305 | transform=None, |
| 306 | target_transform=None, |
| 307 | on_memory=False): |
| 308 | if '22k' in root: |
| 309 | # Imagenet 22k |
| 310 | annotation_root = 'meta_data/' |
| 311 | else: |
| 312 | # Imagenet |
| 313 | annotation_root = 'meta_data/' |
| 314 | if parser is None or isinstance(parser, str): |
| 315 | parser = ParserCephImage(root=root, |
| 316 | split=split, |
| 317 | annotation_root=annotation_root, |
| 318 | on_memory=on_memory) |
| 319 | self.parser = parser |
| 320 | self.transform = transform |
| 321 | self.target_transform = target_transform |
| 322 | self._consecutive_errors = 0 |
| 323 | |
| 324 | def __getitem__(self, index): |
| 325 | img, target = self.parser[index] |
| 326 | self._consecutive_errors = 0 |
| 327 | if self.transform is not None: |
| 328 | img = self.transform(img) |
| 329 | if target is None: |
| 330 | target = -1 |
| 331 | elif self.target_transform is not None: |
| 332 | target = self.target_transform(target) |
| 333 | return img, target |
| 334 | |
| 335 | def __len__(self): |
| 336 | return len(self.parser) |
| 337 | |
| 338 | def filename(self, index, basename=False, absolute=False): |
| 339 | return self.parser.filename(index, basename, absolute) |
| 340 | |
| 341 | def filenames(self, basename=False, absolute=False): |
| 342 | return self.parser.filenames(basename, absolute) |
| 343 | |
| 344 | |
| 345 | class Parser: |