class HEDTester():
def __init__(self, config_file):
self.io = IO()
self.init = True
try:
pfile = open(config_file)
self.cfgs = yaml.load(pfile)
pfile.close()
except Exception as err:
self.io.print_error('Error reading config file {}, {}'.format(config_file), err)
def setup(self, session):
try:
self.model = Vgg16(self.cfgs, run='testing')
meta_model_file = os.path.join(self.cfgs['save_dir'], 'models/hed-model-{}'.format(self.cfgs['test_snapshot']))
saver = tf.train.Saver()
saver.restore(session, meta_model_file)
self.io.print_info('Done restoring VGG-16 model from {}'.format(meta_model_file))
except Exception as err:
self.io.print_error('Error setting up VGG-16 model, {}'.format(err))
self.init = False
def run(self, session):
if not self.init:
return
self.model.setup_testing(session)
filepath = os.path.join(self.cfgs['download_path'], self.cfgs['testing']['list'])
train_list = self.io.read_file_list(filepath)
self.io.print_info('Writing PNGs at {}'.format(self.cfgs['test_output']))
for idx, img in enumerate(train_list):
test_filename = os.path.join(self.cfgs['download_path'], self.cfgs['testing']['dir'], img)
im = self.fetch_image(test_filename)
edgemap = session.run(self.model.predictions, feed_dict={self.model.images: [im]})
# self.save_egdemaps(edgemap, idx)
self.io.print_info('Done testing {}, {}'.format(test_filename, im.shape))
def save_egdemaps(self, em_maps, index):
# Take the edge map from the network from side layers and fuse layer
em_maps = [e[0] for e in em_maps]
em_maps = em_maps + [np.mean(np.array(em_maps), axis=0)]
for idx, em in enumerate(em_maps):
em[em < self.cfgs['testing_threshold']] = 0.0
em = 255.0 * (1.0 - em)
em = np.tile(em, [1, 1, 3])
em = Image.fromarray(np.uint8(em))
em.save(os.path.join(self.cfgs['test_output'], 'testing-{}-{:03}.png'.format(index, idx)))
im = 255 - cv2.imread(os.path.join(self.cfgs['test_output'], 'testing-{}-{:03}.png'.format(index, idx)), cv2.CV_LOAD_IMAGE_GRAYSCALE)
_, bw2 = cv2.threshold(im, 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY)
bw2 = cv2.erode(bw2, np.ones((3,3), np.uint8), iterations=1)
bw2 = thinning(bw2)
cv2.imwrite(os.path.join(self.cfgs['test_output'], 'testing-{}-{:03}-th.png'.format(index, idx)), 255-bw2)
def fetch_image(self, test_image):
# is url
image = None
if not urlparse.urlparse(test_image).scheme == "":
url_response = urllib.urlopen(test_image)
if url_response.code == 404:
print (self.io.print_error('[Testing] URL error code : {1} for {0}'.format(test_image, url_response.code)))
return None
try:
image_buffer = io.BytesIO(url_response.read())
image = self.capture_pixels(image_buffer)
except Exception as err:
print (self.io.print_error('[Testing] Error with URL {0} {1}'.format(test_image, err)))
return None
# read from disk
elif os.path.exists(test_image):
try:
fid = open(test_image, 'r')
stream = fid.read()
fid.close()
image_buffer = io.BytesIO(stream)
image = self.capture_pixels(image_buffer)
except Exception as err:
print (self.io.print_error('[Testing] Error with image file {0} {1}'.format(test_image, err)))
return None
return image
def capture_pixels(self, image_buffer):
image = Image.open(image_buffer)
image = image.resize((self.cfgs['testing']['image_width'], self.cfgs['testing']['image_height']))
image = np.array(image, np.float32)
image = self.colorize(image)
image = image[:, :, self.cfgs['channel_swap']]
image -= self.cfgs['mean_pixel_value']
return image
def colorize(self, image):
# BW to 3 channel RGB image
if image.ndim == 2:
image = image[:, :, np.newaxis]
image = np.tile(image, (1, 1, 3))
elif image.shape[2] == 4:
image = image[:, :, :3]
return image
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
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 HEDTester():
def __init__(self, config_file):
self.io = IO()
self.init = True
try:
pfile = open(config_file)
self.cfgs = yaml.load(pfile)
pfile.close()
except Exception as err:
self.io.print_error(
'Error reading config file {}, {}'.format(config_file), err)
def setup(self, session):
try:
self.model = Vgg16(self.cfgs, run='testing')
meta_model_file = os.path.join(
self.cfgs['save_dir'],
'models/hed-model-{}'.format(self.cfgs['test_snapshot']))
saver = tf.train.Saver()
saver.restore(session, meta_model_file)
self.io.print_info(
'Done restoring VGG-16 model from {}'.format(meta_model_file))
except Exception as err:
self.io.print_error(
'Error setting up VGG-16 model, {}'.format(err))
self.init = False
def run(self, session):
if not self.init:
return
self.model.setup_testing(session)
filepath = os.path.join(self.cfgs['download_path'],
self.cfgs['testing']['list'])
train_list = self.io.read_file_list(filepath)
self.io.print_info('Writing PNGs at {}'.format(
self.cfgs['test_output']))
for idx, img in enumerate(train_list):
test_filename = os.path.join(self.cfgs['download_path'],
self.cfgs['testing']['dir'], img)
im = sio.loadmat(test_filename)['noisy']
#im=cv2.resize(im,(self.cfgs['testing']['image_height'], self.cfgs['testing']['image_width']))
new_im = im.astype(np.float32)
new_im -= self.cfgs['mean_pixel_value']
im = new_im[:, :, 1:5]
#im=np.expand_dims(im,axis=2)
edgemap = session.run(self.model.predictions,
feed_dict={self.model.images: [im]})
self.save_egdemaps(edgemap, idx)
self.io.print_info('Done testing {}, {}'.format(
test_filename, im.shape))
def save_egdemaps(self, em_maps, index):
# Take the edge map from the network from side layers and fuse layer
em_maps = [e[0] for e in em_maps]
em_maps = em_maps + [np.mean(np.array(em_maps), axis=0)]
for idx, em in enumerate(em_maps):
em[em < self.cfgs['testing_threshold']] = 0.0
save_mat_path = os.path.join(self.cfgs['test_output'],
'testing-{}-{:03}'.format(index, idx))
sio.savemat(save_mat_path, {'magnitude_field': em})
em = 255.0 * (1.0 - em)
em = np.tile(em, [1, 1, 3])
em = Image.fromarray(np.uint8(em))
em.save(
os.path.join(self.cfgs['test_output'],
'testing-{}-{:03}.png'.format(index, idx)))