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})