return graph
graph = load_graph(sys.argv[1])
image_dir = sys.argv[2]
# DeepLabv3+ input and output tensors
image_input = graph.get_tensor_by_name('prefix/ImageTensor:0')
softmax = graph.get_tensor_by_name('prefix/SemanticPredictions:0')
# Create output directories in the image folder
if not os.path.exists(image_dir + 'segmented_images/'):
os.mkdir(image_dir + 'segmented_images/')
if not os.path.exists(image_dir + 'segmented_images_colored/'):
os.mkdir(image_dir + 'segmented_images_colored/')
image_dir_segmented = image_dir + 'segmented_images/'
image_dir_segmented_colored = image_dir + 'segmented_images_colored/'
with tf.Session(graph=graph) as sess:
for fname in sorted(os.listdir(image_dir)):
if fname.endswith(".png"):
img = imageio.imread(os.path.join(image_dir, fname))
img = np.expand_dims(img, axis=0)
probs = sess.run(softmax, {image_input: img})
img = tf.squeeze(probs).eval()
img_colored = logits2image(img)
cv2.imwrite(os.path.join(image_dir_segmented + fname), img)
cv2.imwrite(os.path.join(image_dir_segmented_colored + fname),
cv2.cvtColor(img_colored, cv2.COLOR_BGR2RGB))
print(fname)
if os.path.exists(lab_fn):
continue
duration = time.time() - global_start_time
print('%s - %s %d:%02d' %
(seq, fname, duration / 60, duration % 60))
img = cv2.imread(os.path.join(image_dir, fname))
img = cv2.resize(img, (width, height),
interpolation=cv2.INTER_AREA)
img = img[:, :, ::-1]
img = np.expand_dims(img, axis=0)
lab, probs = sess.run([softmax, probs_op], {image_input: img})
lab = np.squeeze(lab)
probs = np.squeeze(probs)
col = logits2image(lab)
probs = (probs * 255).astype(np.uint8)
# save labels
if WRITE:
col_fn = os.path.join(col_out_dir, fname)
cv2.imwrite(col_fn, col)
for k in range(probs.shape[2]):
prob_k_fn = os.path.join(prob_out_dir, 'class_%d' % k,
fname)
#print(prob_k_fn)
cv2.imwrite(prob_k_fn, probs[:, :, k])
#cv2.imwrite(prob_k_fn, (probs[:,:,k]*255).astype(np.uint8))
for k in range(probs.shape[2]):