class HEDReshaper():
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='reshaping')
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_reshaping(session)
idx = 0
saver = tf.train.Saver()
saver.save(session,
os.path.join(self.cfgs['save_dir'],
'reshape_models/hed-model'),
global_step=idx)
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 HEDTrainer():
def __init__(self,dataDir=None,saveDir=None,initmodelfile=None,configfile=None):
self.io = IO()
self.init = True
self.dataDir = dataDir
self.saveDir = saveDir
self.initmodelfile = initmodelfile
try:
pfile = open(configfile)
self.cfgs = yaml.load(pfile)
pfile.close()
except Exception as err:
print('Error reading config file {}, {}'.format(configfile, err))
def setup(self):
try:
self.model = Vgg16(self.cfgs, self.saveDir, self.initmodelfile)
# self.model = HedNet(self.cfgs,self.saveDir,self.initmodelfile)
self.io.print_info('Done initializing VGG-16 model')
dirs = ['train', 'val', 'test', 'models']
save_dir = self.saveDir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
dirs = [os.path.join(save_dir + '/{}'.format(d)) for d in dirs]
_ = [os.makedirs(d) for d in dirs if not os.path.exists(d)]
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
train_data = DataParser(self.cfgs,self.dataDir)
self.model.setup_training(session)
opt = tf.train.AdamOptimizer(self.cfgs['optimizer_params']['learning_rate'])
train = opt.minimize(self.model.loss)
session.run(tf.global_variables_initializer())
for idx in range(self.cfgs['max_iterations']):
im, em, _ = train_data.get_training_batch()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
_, summary, loss,error = session.run([train, self.model.merged_summary, self.model.loss,self.model.error],
feed_dict={self.model.images: im, self.model.edgemaps: em},
options=run_options,
run_metadata=run_metadata)
self.model.train_writer.add_run_metadata(run_metadata, 'step{:06}'.format(idx))
self.io.print_info('[{}/{}] TRAINING loss : {}'.format(idx, self.cfgs['max_iterations'], loss))
self.io.print_info('[{}/{}] TRAINING error : {}'.format(idx, self.cfgs['max_iterations'], error))
if idx % 10 == 0:
self.model.train_writer.add_summary(summary, idx)
'''
if idx % self.cfgs['save_interval'] == 0:
saver = tf.train.Saver()
saver.save(session, os.path.join(self.saveDir, 'models/hed-model'), global_step=idx)
if idx % self.cfgs['val_interval'] == 0:
im, em, _ = train_data.get_validation_batch()
summary, error = session.run([self.model.merged_summary, self.model.error], feed_dict={self.model.images: im, self.model.edgemaps: em})
self.model.val_writer.add_summary(summary, idx)
self.io.print_info('[{}/{}] VALIDATION error : {}'.format(idx, self.cfgs['max_iterations'], error))
'''
if idx == self.cfgs['max_iterations'] - 1:
save_dir = self.saveDir
if not os.path.exists(save_dir):
os.makedirs(save_dir)
graph = tf.graph_util.convert_variables_to_constants(session, session.graph_def, ["fuse"])
tf.train.write_graph(graph, os.path.join(save_dir, 'models'), 'testgraph.pb', as_text=False)
self.model.train_writer.close()
class HEDTrainer():
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:
print(('Error reading config file {}, {}'.format(config_file,
err)))
raise err
def setup(self):
try:
self.model = Vgg16(self.cfgs)
self.io.print_info('Done initializing VGG-16 model')
dirs = ['train', 'val', 'test', 'models']
dirs = [
os.path.join(self.cfgs['save_dir'] + '/{}'.format(d))
for d in dirs
]
[os.makedirs(d) for d in dirs if not os.path.exists(d)]
except Exception as err:
self.io.print_error(
'Error setting up VGG-16 model, {}'.format(err))
self.init = False
raise err
def run(self, session):
if not self.init:
return
train_data = DataParser(self.cfgs)
self.model.setup_training(session)
opt = tf.train.AdamOptimizer(
self.cfgs['optimizer_params']['learning_rate'])
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train = opt.minimize(self.model.loss)
session.run(tf.global_variables_initializer())
for idx in range(self.cfgs['max_iterations']):
im, em, _ = train_data.get_training_batch()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
_, summary, loss = session.run(
[train, self.model.merged_summary, self.model.loss],
feed_dict={
self.model.images: im,
self.model.edgemaps: em,
self.model.training: True
},
options=run_options,
run_metadata=run_metadata)
self.model.train_writer.add_run_metadata(run_metadata,
'step{:06}'.format(idx))
self.model.train_writer.add_summary(summary, idx)
self.io.print_info('[{}/{}] TRAINING loss : {}'.format(
idx, self.cfgs['max_iterations'], loss))
if idx % self.cfgs['save_interval'] == 0:
saver = tf.train.Saver()
saver.save(session,
os.path.join(self.cfgs['save_dir'],
'models/hed-model'),
global_step=idx)
if idx % self.cfgs['val_interval'] == 0:
im, em, _ = train_data.get_validation_batch()
val_im_summary, summary, error = session.run(
[
self.model.val_im_summary, self.model.merged_summary,
self.model.error
],
feed_dict={
self.model.images: im,
self.model.edgemaps: em,
self.model.training: True
})
self.model.val_writer.add_summary(summary, idx)
self.model.val_writer.add_summary(val_im_summary, idx)
self.io.print_info('[{}/{}] VALIDATION error : {}'.format(
idx, self.cfgs['max_iterations'], error))
self.model.train_writer.close()
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)))
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, dataset):
if not self.init:
return
self.model.setup_testing(session)
filepath = os.path.join(self.cfgs['download_path'], self.cfgs['testing']['list'])
train_list = os.listdir(dataset)
self.io.print_info('Writing PNGs at {}'.format(self.cfgs['test_output']))
for idx, img in enumerate(train_list):
idx += 1
test_filename = dataset+'/'+img
im = self.fetch_image(test_filename)
edgemap = session.run(self.model.predictions, feed_dict={self.model.images: [im]})
self.save_egdemaps(edgemap, idx, img)
print 'processing img ', test_filename
def save_egdemaps(self, em_maps, index, name):
# 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):
if idx == 5:
sio.savemat('holy-edge/edgemaps/'+name[:-4] + '.mat', {'predict':em})
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))
if idx == 5:
em.save('holy-edge/edgemaps/'+name)
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 = cStringIO.StringIO(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 = cStringIO.StringIO(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 HEDTester():
def __init__(self, config_file, number_iterations, path_test_image):
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)
try:
self.test_snapshot = number_iterations
except Exception as err:
self.io.print_error('please choose an existing pretrained model')
try:
self.test_image = path_test_image
except Exception as err:
self.io.print_error('the chosen test image does not exist')
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.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)
test_filename = self.test_image
im = sio.loadmat(test_filename)['noisy']
new_im = im.astype(np.float32)
MM, NN = np.shape(new_im)
GR_n = np.zeros((MM, NN, 6), dtype=np.float32)
GR_p = np.zeros((MM, NN, 6), dtype=np.float32)
for alpha in range(1, 7):
GR_n[:, :, alpha - 1], GR_p[:, :,
alpha - 1] = compute_GR(new_im, alpha)
GR_n -= self.cfgs['mean_pixel_value']
im = GR_n[:, :, 1:5]
edgemap = session.run(self.model.predictions,
feed_dict={self.model.images: [im]})
em_maps = [e[0] for e in edgemap]
em_maps = em_maps + [np.mean(np.array(em_maps), axis=0)]
plt.figure()
for idx, em in enumerate(em_maps):
em = np.squeeze(em)
save_mat_dir = 'temp_field'
if not isdir(save_mat_dir):
os.makedirs(save_mat_dir)
save_mat_path = os.path.join('temp_field',
'testing-example-{}'.format(idx))
sio.savemat(save_mat_path, {'magnitude_field': em})