Args: mapper: a symbolic function that takes a tf.string (the raw bytes read from file) and produces a BGR image. Defaults to `fbresnet_mapper(isTrain)`. Returns: A `tf.data.Dataset`. If training, the dataset is infinite. The dataset contains BGR images
(datadir, name, batch_size, mapper=None, parallel=None)
| 113 | |
| 114 | |
| 115 | def get_imagenet_tfdata(datadir, name, batch_size, mapper=None, parallel=None): |
| 116 | """ |
| 117 | Args: |
| 118 | mapper: a symbolic function that takes a tf.string (the raw bytes read from file) and produces a BGR image. |
| 119 | Defaults to `fbresnet_mapper(isTrain)`. |
| 120 | |
| 121 | Returns: |
| 122 | A `tf.data.Dataset`. If training, the dataset is infinite. |
| 123 | The dataset contains BGR images and labels. |
| 124 | """ |
| 125 | |
| 126 | def get_imglist(dir, name): |
| 127 | """ |
| 128 | Returns: |
| 129 | [(full filename, label)] |
| 130 | """ |
| 131 | dir = os.path.join(dir, name) |
| 132 | meta = dataset.ILSVRCMeta() |
| 133 | imglist = meta.get_image_list( |
| 134 | name, |
| 135 | dataset.ILSVRCMeta.guess_dir_structure(dir)) |
| 136 | |
| 137 | def _filter(fname): |
| 138 | # png |
| 139 | return 'n02105855_2933.JPEG' in fname |
| 140 | |
| 141 | ret = [] |
| 142 | for fname, label in imglist: |
| 143 | if _filter(fname): |
| 144 | logger.info("Image {} was filtered out.".format(fname)) |
| 145 | continue |
| 146 | fname = os.path.join(dir, fname) |
| 147 | ret.append((fname, label)) |
| 148 | return ret |
| 149 | |
| 150 | assert name in ['train', 'val', 'test'] |
| 151 | assert datadir is not None |
| 152 | isTrain = name == 'train' |
| 153 | if mapper is None: |
| 154 | mapper = fbresnet_mapper(isTrain) |
| 155 | if parallel is None: |
| 156 | parallel = min(40, multiprocessing.cpu_count() // 2) # assuming hyperthreading |
| 157 | imglist = get_imglist(datadir, name) |
| 158 | |
| 159 | N = len(imglist) |
| 160 | filenames = tf.constant([k[0] for k in imglist], name='filenames') |
| 161 | labels = tf.constant([k[1] for k in imglist], dtype=tf.int32, name='labels') |
| 162 | |
| 163 | ds = tf.data.Dataset.from_tensor_slices((filenames, labels)) |
| 164 | |
| 165 | if isTrain: |
| 166 | ds = ds.shuffle(N, reshuffle_each_iteration=True).repeat() |
| 167 | |
| 168 | ds = ds.apply( |
| 169 | tf.data.experimental.map_and_batch( |
| 170 | lambda fname, label: (mapper(tf.read_file(fname)), label), |
| 171 | batch_size=batch_size, |
| 172 | num_parallel_batches=parallel)) |
no test coverage detected
searching dependent graphs…