output_image_path = path.join( '/home/bill/Desktop/thesis/code/dilation-tensorflow/data', dataset + 'a_out.png') # Build pretrained model and save it as TF checkpoint with tf.Session() as sess: # Choose input shape according to dataset characteristics input_h, input_w, input_c = CONFIG[dataset]['input_shape'] input_tensor = tf.placeholder(tf.float32, shape=(None, input_h, input_w, input_c), name='input_placeholder') # Create pretrained model model = dilation_model_pretrained(dataset, input_tensor, w_pretrained, trainable=False) sess.run(tf.global_variables_initializer()) # Save both graph and weights saver = tf.train.Saver(tf.get_collection( tf.GraphKeys.GLOBAL_VARIABLES)) saver.save(sess, path.join(checkpoint_dir, 'dilation')) # Restore both graph and weights from TF checkpoint with tf.Session() as sess: saver = tf.train.import_meta_graph( path.join(checkpoint_dir, 'dilation.meta')) saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir))
def main(): # Choose between 'cityscapes' and 'camvid' dataset = 'cityscapes' # Load dict of pretrained weights print('Loading pre-trained weights...') with open(CONFIG[dataset]['weights_file'], 'rb') as f: w_pretrained = pickle.load(f) print('Done.') # Create checkpoint directory checkpoint_dir = path.join('data/checkpoint', 'dilation_' + dataset) if not path.exists(checkpoint_dir): os.makedirs(checkpoint_dir) Path = '/Users/alasdair/Desktop/Thesis_git/Data/test_ims' Path_out = '/Users/alasdair/Desktop/Thesis_git/Data/outputs' im_names = [] for filename in os.listdir(Path): im_names.append(filename) im_names.sort() print('Number of images: ', len(im_names)) masks = np.zeros((256, 512, len(im_names))) # Build pretrained model and save it as TF checkpoint with tf.Session() as sess: # Choose input shape according to dataset characteristics input_h, input_w, input_c = CONFIG[dataset]['input_shape'] input_tensor = tf.placeholder(tf.float32, shape=(None, input_h, input_w, input_c), name='input_placeholder') # Create pretrained model model = dilation_model_pretrained(dataset, input_tensor, w_pretrained, trainable=False) sess.run(tf.global_variables_initializer()) # Save both graph and weights saver = tf.train.Saver(tf.get_collection( tf.GraphKeys.GLOBAL_VARIABLES)) saver.save(sess, path.join(checkpoint_dir, 'dilation')) # Restore both graph and weights from TF checkpoint with tf.Session() as sess: saver = tf.train.import_meta_graph( path.join(checkpoint_dir, 'dilation.meta')) saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir)) graph = tf.get_default_graph() model = graph.get_tensor_by_name('softmax:0') model = tf.reshape(model, shape=(1, ) + CONFIG[dataset]['output_shape']) for i in range(len(im_names)): file_n = im_names[i] input_image_path = os.path.join(Path, file_n) input_image = cv2.resize(cv2.imread(input_image_path), (2048, 1024)) print('Reading: ', file_n) output_image_path = path.join(Path_out, file_n[:-4] + '_out.png') output_mask_path = path.join(Path_out, file_n[:-4] + '_mask.png') # Read and predict on a test image # input_image = cv2.imread(input_image_path) input_tensor = graph.get_tensor_by_name('input_placeholder:0') predicted_image = predict(input_image, input_tensor, model, dataset, sess) # Convert colorspace (palette is in RGB) and save prediction result predicted_image = np.array(cv2.resize(predicted_image, (512, 256))) mask = np.sum((predicted_image >= 250).astype(int), axis=2) masks[:, :, i] = mask predicted_image = cv2.cvtColor(predicted_image, cv2.COLOR_BGR2RGB) print(output_image_path, i) cv2.imwrite(output_image_path, predicted_image) np.save(os.path.join(Path + 'masks.npy'), mask)