__init__.py 548 B

123456789101112131415161718192021222324
  1. from imageio import imread
  2. import numpy as np
  3. class Dataset(object):
  4. def __init__(self, uuids, annotations):
  5. super(Dataset, self).__init__()
  6. self.uuids = uuids
  7. self._annot = annotations
  8. def __len__(self):
  9. return len(self.uuids)
  10. def _get(self, method, i):
  11. return getattr(self._annot, method)(self.uuids[i])
  12. def get_example(self, i, mode="RGB"):
  13. methods = ["image", "parts", "label"]
  14. im_path, parts, label = [self._get(m, i) for m in methods]
  15. return imread(im_path, pilmode=mode), parts, label
  16. __getitem__ = get_example