| 482 | pbar.desc = 'Caching images (%.1fGB)' % (gb / 1E9) |
| 483 | |
| 484 | def cache_labels(self, path='labels.cache'): |
| 485 | # Cache dataset labels, check images and read shapes |
| 486 | x = {} # dict |
| 487 | pbar = tqdm(zip(self.img_files, self.label_files), desc='Scanning images', total=len(self.img_files)) |
| 488 | for (img, label) in pbar: |
| 489 | try: |
| 490 | l = [] |
| 491 | im = Image.open(img) |
| 492 | im.verify() # PIL verify |
| 493 | shape = exif_size(im) # image size |
| 494 | assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels' |
| 495 | if os.path.isfile(label): |
| 496 | with open(label, 'r') as f: |
| 497 | l = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) # labels |
| 498 | if len(l) == 0: |
| 499 | l = np.zeros((0, 5), dtype=np.float32) |
| 500 | x[img] = [l, shape] |
| 501 | except Exception as e: |
| 502 | print('WARNING: Ignoring corrupted image and/or label %s: %s' % (img, e)) |
| 503 | |
| 504 | x['hash'] = get_hash(self.label_files + self.img_files) |
| 505 | torch.save(x, path) # save for next time |
| 506 | return x |
| 507 | |
| 508 | def __len__(self): |
| 509 | return len(self.img_files) |