Spaces:
Sleeping
Sleeping
File size: 4,094 Bytes
b18cfd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
"""Compute depth maps for images in the input folder.
"""
import os
import glob
import utils
import cv2
import argparse
import tensorflow as tf
from transforms import Resize, NormalizeImage, PrepareForNet
def run(input_path, output_path, model_path, model_type="large"):
"""Run MonoDepthNN to compute depth maps.
Args:
input_path (str): path to input folder
output_path (str): path to output folder
model_path (str): path to saved model
"""
print("initialize")
# the runtime initialization will not allocate all memory on the device to avoid out of GPU memory
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
for gpu in gpus:
#tf.config.experimental.set_memory_growth(gpu, True)
tf.config.experimental.set_virtual_device_configuration(gpu,
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4000)])
except RuntimeError as e:
print(e)
# network resolution
if model_type == "large":
net_w, net_h = 384, 384
elif model_type == "small":
net_w, net_h = 256, 256
else:
print(f"model_type '{model_type}' not implemented, use: --model_type large")
assert False
# load network
graph_def = tf.compat.v1.GraphDef()
with tf.io.gfile.GFile(model_path, 'rb') as f:
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
model_operations = tf.compat.v1.get_default_graph().get_operations()
input_node = '0:0'
output_layer = model_operations[len(model_operations) - 1].name + ':0'
print("Last layer name: ", output_layer)
resize_image = Resize(
net_w,
net_h,
resize_target=None,
keep_aspect_ratio=False,
ensure_multiple_of=32,
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
)
def compose2(f1, f2):
return lambda x: f2(f1(x))
transform = compose2(resize_image, PrepareForNet())
# get input
img_names = glob.glob(os.path.join(input_path, "*"))
num_images = len(img_names)
# create output folder
os.makedirs(output_path, exist_ok=True)
print("start processing")
with tf.compat.v1.Session() as sess:
try:
# load images
for ind, img_name in enumerate(img_names):
print(" processing {} ({}/{})".format(img_name, ind + 1, num_images))
# input
img = utils.read_image(img_name)
img_input = transform({"image": img})["image"]
# compute
prob_tensor = sess.graph.get_tensor_by_name(output_layer)
prediction, = sess.run(prob_tensor, {input_node: [img_input] })
prediction = prediction.reshape(net_h, net_w)
prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
# output
filename = os.path.join(
output_path, os.path.splitext(os.path.basename(img_name))[0]
)
utils.write_depth(filename, prediction, bits=2)
except KeyError:
print ("Couldn't find input node: ' + input_node + ' or output layer: " + output_layer + ".")
exit(-1)
print("finished")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_path',
default='input',
help='folder with input images'
)
parser.add_argument('-o', '--output_path',
default='output',
help='folder for output images'
)
parser.add_argument('-m', '--model_weights',
default='model-f6b98070.pb',
help='path to the trained weights of model'
)
parser.add_argument('-t', '--model_type',
default='large',
help='model type: large or small'
)
args = parser.parse_args()
# compute depth maps
run(args.input_path, args.output_path, args.model_weights, args.model_type)
|