class DataParser():
def __init__(self, cfgs, dataDir=None, initmodelfile=None):
self.io = IO()
self.cfgs = cfgs
download_path = dataDir
self.train_file = os.path.join(download_path, cfgs['training']['list'])
self.train_data_dir = os.path.join(download_path,
cfgs['training']['dir'])
self.training_pairs = self.io.read_file_list(self.train_file)
self.samples = self.io.split_pair_names(self.training_pairs,
self.train_data_dir)
self.io.print_info('Training data set-up from {}'.format(
os.path.join(self.train_file)))
self.n_samples = len(self.training_pairs)
self.all_ids = list(range(self.n_samples))
np.random.shuffle(self.all_ids)
self.training_ids = self.all_ids[:int(self.cfgs['train_split'] *
len(self.training_pairs))]
self.validation_ids = self.all_ids[int(self.cfgs['train_split'] *
len(self.training_pairs)):]
self.io.print_info('Training samples {}'.format(len(
self.training_ids)))
self.io.print_info('Validation samples {}'.format(
len(self.validation_ids)))
def get_training_batch(self):
batch_ids = np.random.choice(self.training_ids,
self.cfgs['batch_size_train'])
return self.get_batch(batch_ids)
def get_validation_batch(self):
batch_ids = np.random.choice(self.validation_ids,
self.cfgs['batch_size_val'])
return self.get_batch(batch_ids)
def get_batch(self, batch):
tstart = time.time()
filenames = []
images = []
edgemaps = []
for idx, b in enumerate(batch):
im = Image.open(self.samples[b][0])
em = Image.open(self.samples[b][1])
im = im.resize((self.cfgs['training']['image_width'],
self.cfgs['training']['image_height']))
em = em.resize((self.cfgs['training']['image_width'],
self.cfgs['training']['image_height']))
im = np.array(im, dtype=np.float32)
im = im[:, :, self.cfgs['channel_swap']]
im -= self.cfgs['mean_pixel_value']
# Labels needs to be 1 or 0 (edge pixel or not)
# or can use regression targets as done by the author
# https://github.com/s9xie/hed/blob/9e74dd710773d8d8a469ad905c76f4a7fa08f945/src/caffe/layers/image_labelmap_data_layer.cpp#L213
em = np.array(em.convert('L'), dtype=np.float32)
if self.cfgs['target_regression']:
bin_em = em / 255.0
else:
bin_em = np.zeros_like(em)
bin_em[np.where(em)] = 1
# Some edge maps have 3 channels some dont
bin_em = bin_em if bin_em.ndim == 2 else bin_em[:, :, 0]
# To fit [batch_size, H, W, 1] output of the network
bin_em = np.expand_dims(bin_em, 2)
images.append(im)
edgemaps.append(bin_em)
filenames.append(self.samples[b])
return images, edgemaps, filenames
class DataParser():
def __init__(self, cfgs):
self.io = IO()
self.cfgs = cfgs
self.train_file = os.path.join(cfgs['download_path'],
cfgs['training']['list'])
self.train_data_dir = os.path.join(cfgs['download_path'],
cfgs['training']['dir'])
self.training_pairs = self.io.read_file_list(self.train_file)
self.samples = self.io.split_pair_names(self.training_pairs,
self.train_data_dir)
self.io.print_info('Training data set-up from {}'.format(
os.path.join(self.train_file)))
self.n_samples = len(self.training_pairs)
self.all_ids = range(self.n_samples)
np.random.shuffle(self.all_ids)
self.training_ids = self.all_ids[:int(self.cfgs['train_split'] *
len(self.training_pairs))]
self.validation_ids = self.all_ids[int(self.cfgs['train_split'] *
len(self.training_pairs)):]
self.io.print_info('Training samples {}'.format(len(
self.training_ids)))
self.io.print_info('Validation samples {}'.format(
len(self.validation_ids)))
def get_training_batch(self):
batch_ids = np.random.choice(self.training_ids,
self.cfgs['batch_size_train'])
return self.get_batch(batch_ids)
def get_validation_batch(self):
batch_ids = np.random.choice(self.validation_ids,
self.cfgs['batch_size_val'])
return self.get_batch(batch_ids)
def get_batch(self, batch):
tstart = time.time()
filenames = []
images = []
edgemaps = []
for idx, b in enumerate(batch):
im = Image.open(self.samples[b][0])
em = Image.open(self.samples[b][1])
im = im.resize((self.cfgs['training']['image_width'],
self.cfgs['training']['image_height']))
em = em.resize((self.cfgs['training']['image_width'],
self.cfgs['training']['image_height']))
im = np.array(im, dtype=np.float32)
im = im[:, :, self.cfgs['channel_swap']]
im -= self.cfgs['mean_pixel_value']
em = np.array(em.convert('L'), dtype=np.float32)
if self.cfgs['target_regression']:
bin_em = em / 255.0
else:
bin_em = np.zeros_like(em)
bin_em[np.where(em)] = 1
# Some edge maps have 3 channels some dont
bin_em = bin_em if bin_em.ndim == 2 else bin_em[:, :, 0]
# To fit [batch_size, H, W, 1] output of the network
bin_em = np.expand_dims(bin_em, 2)
images.append(im)
edgemaps.append(bin_em)
filenames.append(self.samples[b])
return images, edgemaps, filenames