| 8 | |
| 9 | |
| 10 | class Dataset(object): |
| 11 | def __init__(self, images, imsize, embeddings=None, |
| 12 | filenames=None, workdir=None, |
| 13 | labels=None, aug_flag=True, |
| 14 | class_id=None, class_range=None): |
| 15 | self._images = images |
| 16 | self._embeddings = embeddings |
| 17 | self._filenames = filenames |
| 18 | self.workdir = workdir |
| 19 | self._labels = labels |
| 20 | self._epochs_completed = -1 |
| 21 | self._num_examples = len(images) |
| 22 | self._saveIDs = self.saveIDs() |
| 23 | |
| 24 | # shuffle on first run |
| 25 | self._index_in_epoch = self._num_examples |
| 26 | self._aug_flag = aug_flag |
| 27 | self._class_id = np.array(class_id) |
| 28 | self._class_range = class_range |
| 29 | self._imsize = imsize |
| 30 | self._perm = None |
| 31 | |
| 32 | @property |
| 33 | def images(self): |
| 34 | return self._images |
| 35 | |
| 36 | @property |
| 37 | def embeddings(self): |
| 38 | return self._embeddings |
| 39 | |
| 40 | @property |
| 41 | def filenames(self): |
| 42 | return self._filenames |
| 43 | |
| 44 | @property |
| 45 | def num_examples(self): |
| 46 | return self._num_examples |
| 47 | |
| 48 | @property |
| 49 | def epochs_completed(self): |
| 50 | return self._epochs_completed |
| 51 | |
| 52 | def saveIDs(self): |
| 53 | self._saveIDs = np.arange(self._num_examples) |
| 54 | np.random.shuffle(self._saveIDs) |
| 55 | return self._saveIDs |
| 56 | |
| 57 | def readCaptions(self, filenames, class_id): |
| 58 | name = filenames |
| 59 | if name.find('jpg/') != -1: # flowers dataset |
| 60 | class_name = 'class_%05d/' % class_id |
| 61 | name = name.replace('jpg/', class_name) |
| 62 | cap_path = '%s/text_c10/%s.txt' %\ |
| 63 | (self.workdir, name) |
| 64 | with open(cap_path, "r") as f: |
| 65 | captions = f.read().split('\n') |
| 66 | captions = [cap for cap in captions if len(cap) > 0] |
| 67 | return captions |