class DataProvider:
def __init__(self, is_training, folds):
self.cursor = 0
records = load_data(folds)
self.is_training = is_training
self.records = records
random.shuffle(self.records)
self.data_count = len(self.records)
print ('#records:', self.data_count, 'preprocessing...')
self.preprocess()
self.batch_size = None
self.async_task = None
def preprocess(self):
images = []
nrgbs = []
illums = []
for i in range(len(self.records)):
images.append(self.records[i].img)
nrgbs.append(None)
illums.append(self.records[i].illum)
# No same size...
self.images, self.nrgbs, self.illums = images, nrgbs, np.vstack(illums)
def set_batch_size(self, batch_size):
assert self.batch_size is None
self.batch_size = batch_size
def shuffle(self):
ind = list(range(self.data_count))
random.shuffle(ind)
images = [self.images[i] for i in ind]
nrgbs = [self.nrgbs[i] for i in ind]
illums = [self.illums[i] for i in ind]
self.images = images
self.nrgbs = nrgbs
self.illums = illums
def get_batch_(self):
batch_size = self.batch_size
indices = []
while len(indices) < batch_size:
s = min(self.data_count - self.cursor, batch_size - len(indices))
indices += range(self.cursor, self.cursor + s)
if self.cursor + s >= self.data_count:
if self.is_training and DATA_SHUFFLE:
self.shuffle()
self.cursor = (self.cursor + s) % self.data_count
next_batch = [[], [], []]
for i in indices:
ldr, nrgb = self.images[i], self.nrgbs[i]
illum = self.illums[i]
if self.is_training and AUGMENTATION:
ldr, illum = augment(ldr, illum)
else:
ldr = ldr[:FCN_INPUT_SIZE, :FCN_INPUT_SIZE]
nrgb = None
next_batch[0].append(ldr)
next_batch[1].append(nrgb)
next_batch[2].append(illum)
next_batch = (np.stack(next_batch[0]), np.stack(next_batch[1]),
np.vstack(next_batch[2]))
return next_batch
def get_batch(self):
if self.async_task is None:
self.async_task = AsyncTaskManager(self.get_batch_)
return self.async_task.get_next()
def stop(self):
self.async_task.stop()
class DataProvider:
def __init__(self, is_training, folds):
self.cursor = 0
records = load_data(folds)
self.is_training = is_training
self.records = records
random.shuffle(self.records)
self.data_count = len(self.records)
print('#records:', self.data_count, 'preprocessing...')
self.preprocess()
self.batch_size = None
self.async_task = None
def preprocess(self):
images = []
nrgbs = []
illums = []
for i in range(len(self.records)):
images.append(self.records[i].img)
nrgbs.append(None)
illums.append(self.records[i].illum)
# No same size...
self.images, self.nrgbs, self.illums = images, nrgbs, np.vstack(illums)
def set_batch_size(self, batch_size):
assert self.batch_size is None
self.batch_size = batch_size
def shuffle(self):
ind = range(self.data_count)
random.shuffle(list(ind))
images = [self.images[i] for i in ind]
nrgbs = [self.nrgbs[i] for i in ind]
illums = [self.illums[i] for i in ind]
self.images = images
self.nrgbs = nrgbs
self.illums = illums
def get_batch_(self):
batch_size = self.batch_size
indices = []
while len(indices) < batch_size:
s = min(self.data_count - self.cursor, batch_size - len(indices))
indices += range(self.cursor, self.cursor + s)
if self.cursor + s >= self.data_count:
if self.is_training and FLAGS.DATA_SHUFFLE:
self.shuffle()
self.cursor = (self.cursor + s) % self.data_count
next_batch = [[], [], []]
for i in indices:
ldr, nrgb = self.images[i], self.nrgbs[i]
illum = self.illums[i]
# random choise another sample for splition
i_ = random.choice(indices)
ldr_, nrgb_ = self.images[i_], self.nrgbs[i_]
illum_ = self.illums[i_]
ldr = cv2.resize(ldr, (FLAGS.img_size, FLAGS.img_size))
ldr_ = cv2.resize(ldr_, (FLAGS.img_size, FLAGS.img_size))
if self.is_training:
if FLAGS.AUGMENTATION:
ldr, illum = augment(ldr, illum)
ldr_, illum_ = augment(ldr_, illum_)
if FLAGS.SPLITION:
ldr, illum = split(ldr, illum, ldr_, illum_)
nrgb = None
next_batch[0].append(ldr)
next_batch[1].append(nrgb)
next_batch[2].append(illum)
next_batch = (np.stack(next_batch[0]), np.stack(next_batch[1]),
np.vstack(next_batch[2]))
return next_batch
def get_batch(self):
if self.async_task is None:
self.async_task = AsyncTaskManager(self.get_batch_)
return self.async_task.get_next()
def stop(self):
self.async_task.stop()