| 1234567891011121314151617181920212223242526272829 |
- from imageio import imread
- import numpy as np
- class Dataset(object):
- def __init__(self, uuids, annotations):
- super(Dataset, self).__init__()
- self.uuids = uuids
- self._annot = annotations
- def __len__(self):
- return len(self.uuids)
- def _get(self, method, i):
- return getattr(self._annot, method)(self.uuids[i])
- def get_example(self, i, mode="RGB"):
- methods = ["image", "parts", "label"]
- im_path, parts, label = [self._get(m, i) for m in methods]
- im = imread(im_path, pilmode=mode)
- h,w,c = im.shape
- # fit to the dimensions of the image
- parts[:, 1] = np.minimum(parts[:, 1], w - 1)
- parts[:, 2] = np.minimum(parts[:, 2], h - 1)
- return im, parts, label
- __getitem__ = get_example
|
