def get_data(file_name):
if file_name.endswith('.lmdb'):
ds = LMDBDataPoint(file_name, shuffle=True)
ds = ImageDecode(ds, index=0)
elif file_name.endswith('.zip'):
ds = ImageDataFromZIPFile(file_name, shuffle=True)
ds = ImageDecode(ds, index=0)
ds = RejectTooSmallImages(ds, index=0)
ds = CenterSquareResize(ds, index=0)
else:
raise ValueError("Unknown file format " + file_name)
augmentors = [imgaug.RandomCrop(128),
imgaug.Flip(horiz=True)]
ds = AugmentImageComponent(ds, augmentors, index=0, copy=True)
ds = MapData(ds, lambda x: [cv2.resize(x[0], (32, 32), interpolation=cv2.INTER_CUBIC), x[0]])
ds = PrefetchDataZMQ(ds, 3)
ds = BatchData(ds, BATCH_SIZE)
return ds
def get_data(lmdb):
ds = LMDBDataPoint(lmdb, shuffle=True)
ds = ImageDecode(ds, index=0)
augmentors = [imgaug.RandomCrop(128), imgaug.Flip(horiz=True)]
ds = AugmentImageComponent(ds, augmentors, index=0, copy=True)
ds = MapData(ds, lambda x: x - VGG_MEAN)
ds = MapData(
ds, lambda x:
[cv2.resize(x[0], (32, 32), interpolation=cv2.INTER_AREA), x[0]])
ds = PrefetchDataZMQ(ds, 8)
ds = BatchData(ds, BATCH_SIZE)
return ds
def get_data(lmdb_path):
ds_train = LMDBDataPoint(os.path.join(lmdb_path, 'train2017.lmdb'), shuffle=True)
ds_train = ImageDecode(ds_train, index=0)
augmentors = [
imgaug.RandomCrop((HR_SIZE, HR_SIZE)),
imgaug.Flip(horiz=True)
]
ds_train = AugmentImageComponent(ds_train, augmentors, 0)
ds_train = MapData(ds_train, lambda x: [x[0], x[0].copy()])
ds_train = AugmentImageComponent(ds_train, [imgaug.Resize(LR_SIZE, interp=cv2.INTER_CUBIC)], 0)
ds_train = BatchData(ds_train, BATCH)
ds_train = PrefetchDataZMQ(ds_train, 4)
ds_val = LMDBDataPoint(os.path.join(lmdb_path, 'val2017.lmdb'), shuffle=False)
ds_val = FixedSizeData(ds_val, 500)
ds_val = ImageDecode(ds_val, index=0)
augmentors = [
imgaug.CenterCrop((HR_SIZE, HR_SIZE))
]
ds_val = AugmentImageComponent(ds_val, augmentors, 0)
ds_val = MapData(ds_val, lambda x: [x[0], x[0].copy()])
ds_val = AugmentImageComponent(ds_val, [imgaug.Resize(LR_SIZE, interp=cv2.INTER_CUBIC)], 0)
ds_val = BatchData(ds_val, BATCH)
return ds_train, ds_val
