def generative_inpainting(image, mask):
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
model = InpaintCAModel()
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
checkpoint_dir = 'generative_inpainting/model_logs/release_imagenet_256/'
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
result = sess.run(output)
tf.reset_default_graph()
return result
def freeze(ckpt_path):
image = tf.placeholder(tf.float32, shape=(1, None, None, 3), name='image')
mask = tf.placeholder(tf.float32, shape=(1, None, None, 3), name='mask')
model = InpaintCAModel()
input = tf.concat([image, mask],axis=2)
output = model.build_server_graph(input)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
output = tf.add(output, 0, name='output')
init_op = tf.global_variables_initializer()
restore_saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init_op)
restore_saver.restore(sess, ckpt_path)
frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def,
output_node_names=['output'])
path = os.path.dirname(ckpt_path)
with open(path + '/deepfill.pb', 'wb') as f:
f.write(frozen_graph_def.SerializeToString())
print("frozen model path: {}".format( path + '/deepfill.pb'))
def setup(opts):
global g
global sess
global input_image
model = InpaintCAModel()
g = tf.get_default_graph()
sess = tf.Session(graph=g)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
input_image = tf.placeholder(tf.float32, shape=(1, 256, 256 * 2, 3))
output_image = model.build_server_graph(input_image)
output_image = (output_image + 1.) * 127.5
output_image = tf.reverse(output_image, [-1])
output_image = tf.saturate_cast(output_image, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(opts['checkpoint_dir'],
from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
return output_image
def __init__(self, checkpoint_dir):
model = InpaintCAModel()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
input_image_ph = tf.placeholder(tf.float32, shape=(1, 256, 512, 3))
output = model.build_server_graph(input_image_ph, reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# === END OF BUILD GRAPH ===
sess = tf.Session(config=sess_config)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
# sess.run(assign_ops)
self.sess = sess
self.graph = output
self.placeholder = input_image_ph
self.assign_ops = assign_ops
def Output(img_in, mask_in, img_out):
print("import from test:", img_in, mask_in, img_out)
#if __name__ == "__main__":
FLAGS = ng.Config('inpaint.yml')
# ng.get_gpus(1)
#args, unknown = parser.parse_known_args()
model = InpaintCAModel()
#image = cv2.imread(args.image)
image = cv2.imread("examples/places356/" + img_in)
#mask = cv2.imread(args.mask)
mask = cv2.imread("places356_mask/" + mask_in)
# mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
#var_value = tf.contrib.framework.load_variable(args.checkpoint_dir, from_name)
var_value = tf.contrib.framework.load_variable(
"model_logs/release_places2_256", from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
#cv2.imwrite(args.output, result[0][:, :, ::-1])
#cv2.imshow("result", result[0][:, :, ::-1])
cv2.imwrite("examples/places356/" + img_out, result[0][:, :, ::-1])
show1 = cv2.imread("examples/places356/" + img_in)
show2 = cv2.imread("examples/places356/" + img_out)
show = np.hstack([show1, show2])
cv2.imshow("result", show)
def inpaint(arg_image_dir, arg_mask_dir, arg_checkpoint_dir, arg_output_dir):
tf.reset_default_graph()
FLAGS = ng.Config('inpaint.yml')
# ng.get_gpus(1)
model = InpaintCAModel()
for arg_image in os.listdir(arg_image_dir):
arg_mask = arg_image # assume the mask has the same name as the image
if os.path.exists(arg_output_dir + arg_image):
print("note |", arg_image, "already inpainted.")
continue
if os.path.exists(arg_image_dir +
arg_image) and os.path.exists(arg_mask_dir +
arg_mask):
pass
else:
continue
image = cv2.imread(os.path.join(arg_image_dir, arg_image))
mask = cv2.imread(os.path.join(arg_mask_dir, arg_mask))
name = arg_image
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
tf.reset_default_graph()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
arg_checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv2.imwrite(arg_output_dir + name, result[0][:, :, ::-1])
def main():
ng.get_gpus(1)
args = parser.parse_args()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
model = InpaintCAModel()
input_image_ph = tf.placeholder(tf.float32,
shape=(1, IMG_HEIGHT, IMG_WIDTH * 2, 3))
output = model.build_server_graph(input_image_ph)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir,
from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
while (True):
try:
new_jobs = [
f for f in listdir(SYNC_FLODER_NAME) if f.startswith("go_")
]
if len(new_jobs) == 0:
sleep(0.1)
continue
job_id = new_jobs[0][3:]
mask_name = open(SYNC_FLODER_NAME + new_jobs[0]).readline().strip()
os.remove(SYNC_FLODER_NAME + new_jobs[0]) #delete go_ file
print("Processing image: ", job_id)
pic_folder = FILES_FLODER_NAME + job_id + "/"
process_img_new (img_name=pic_folder + "init", \
mask_name=mask_name, \
result_name=pic_folder + "out.png", \
sess=sess, \
output=output, \
input_image_ph=input_image_ph)
open(SYNC_FLODER_NAME + "done_" + job_id,
'a').close() #create done_ file
print("Job {} done.".format(job_id))
except KeyboardInterrupt:
sys.exit()
def inpaintP(name):
CUR_DIR = os.path.join(WORK_DIR, name)
TMASK_DIR = os.path.join(WORK_DIR, name + "//tmask")
INPAINT_MODEL_PATH = os.path.join(ROOT_DIR,
"model_logs/release_places2_256")
FLAGS = ng.Config('inpaint.yml')
model = InpaintCAModel()
image = cv.imread(os.path.join(CUR_DIR, f"{name}.png"))
mask = cv.imread(os.path.join(TMASK_DIR, "mask.png"))
filename = f'4#_{name}.png'
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.per_process_gpu_memory_fraction = 0.5
tf2 = tf.Graph()
with tf2.as_default():
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
INPAINT_MODEL_PATH, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv.imwrite(os.path.join(CUR_DIR, filename), result[0][:, :, ::-1])
print('Image has been made')
return 0
def test_single_image(image_path, mask_path, output_path, image_height,
image_width, ckpt_dir):
# generate input image
input_image = get_input_image(image_path, mask_path, image_height,
image_width)
# start sess configuration
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
# inpaint model
model = InpaintCAModel()
# input node placeholder
input_image_ph = tf.placeholder(tf.float32,
name="input",
shape=(1, image_height, image_width * 2,
3))
output = model.build_server_graph(input_image_ph)
output = (output + 1.0) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load variables from checkpoint
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(ckpt_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
# input, output
print("input = {}".format(input_image_ph))
print("output = {}".format(output))
# run model
print("Running model ...")
result = sess.run(output, feed_dict={input_image_ph: input_image})
print("Running model done.")
# save output image to file
print("Saving output to {}".format(output_path))
cv2.imwrite(output_path, result[0][:, :, ::-1])
print("Saving output done.")
save_tensorboard_log(sess, "trensorboard")
return sess
def generate_counterfactual(image_fp,
mask_fp,
output_fp,
checkpoint_dir,
model_id=None):
try:
FLAGS = ng.Config('config/inpaint.yml')
except AssertionError:
raise ValueError('check directory above')
# ng.get_gpus(1)
# args, unknown = parser.parse_known_args()
model = InpaintCAModel()
image = cv2.imread(image_fp)
mask = cv2.imread(mask_fp)
# mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv2.imwrite(output_fp, result[0][:, :, ::-1])
print(f'IMAGE WROTE TO {output_fp}\n\n\n')
tf.reset_default_graph()
def downloadFile():
start = time.time()
os.environ['CUDA_VISIABLE_DEVICE'] = '0'
args = parser.parse_args()
model = InpaintCAModel()
image = cv2.imread(args.image)
mask = cv2.imread(args.mask)
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
#sess_config = tf.ConfigProto()
#sess_config.gpu_options.allow_growth = True
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir,
from_name)
assign_ops.append(tf.assign(var, var_value))
global sess_config
with tf.Session(config=sess_config) as sess:
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv2.imwrite(args.output, result[0][:, :, ::-1])
print(time.time() - start)
return send_file('output.png')
def fillVideo(img, msk, i):
checkpoint_dir = 'model_logs/snap-0'
output_dir = 'output/'
ng.get_gpus(0)
model = InpaintCAModel()
image = cv2.imread(img)
# image = cv2.imread('./examples/places2/canyon_input.png')
print(image.shape)
# cv2.imshow(image)
# cv2.waitKey(0)
mask = cv2.imread(msk)
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
output_name = output_dir + 'output_' + str(i) + '.png'
cv2.imwrite(output_name, result[0][:, :, ::-1])
def pixel_fill(image, mask):
import tensorflow as tf2
import neuralgym as ng
from inpaint_model import InpaintCAModel
if image.ndim > mask.ndim:
mask = np.dstack([mask] * image.shape[2])
assert image.shape == mask.shape
model = InpaintCAModel()
FLAGS = ng.Config('inpaint.yml')
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf2.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf2.Session(config=sess_config) as sess:
input_image = tf2.constant(input_image, dtype=tf2.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf2.reverse(output, [-1])
output = tf2.saturate_cast(output, tf2.uint8)
# load pretrained model
vars_list = tf2.get_collection(tf2.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
print("checkpoint_dir = ", checkpoint_dir)
for var in vars_list:
vname = var.name
from_name = vname
if "inpaint_net" in var.name: # or else is going to mix with mrcnn
var_value = tf2.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf2.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.' * 10)
result = sess.run(output)
sess.close()
tf2.reset_default_graph()
return result[0][:, :, ::-1]
def generate_img_mask(self, image, mask):
mask, image = try_to_adjust_to_3c_mask(image, mask)
assert (image.shape == mask.shape
), "Mask shape {0} != img shape {1}".format(
mask.shape, image.shape)
# ng.get_gpus(1)
model = InpaintCAModel()
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
temp_graph = tf.Graph()
with temp_graph.as_default():
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
self.generative_model_path, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
result = sess.run(output)
out_img = result[0][:, :, ::-1]
return [out_img]
def test_model(image,
mask,
output_dir='output_images/output.png',
checkpoint_dir='model_logs/release_places2_256'):
ng.get_gpus(1)
model = InpaintCAModel()
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
# cv2.imwrite(output_dir, result[0][:, :, ::-1])
# plt.imsave('out.jpg', result[0][:, :, ::-1])
return result[0]
def complete(image_file):
ng.get_gpus(1,verbose=False)
tf.reset_default_graph()
model = InpaintCAModel()
image = cv2.imread(os.path.join(args.image_dir, image_file))
mask = cv2.imread(os.path.join(args.mask_dir, image_file))
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image_rs = image[:h // grid * grid, :w // grid * grid, :]
mask_rs = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image_rs.shape))
image_rs = np.expand_dims(image_rs, 0)
mask_rs = np.expand_dims(mask_rs, 0)
input_image = np.concatenate([image_rs, mask_rs], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
result = sess.run(output)
image[:h // grid * grid, :w // grid * grid, :] = result[0][:, :, ::-1]
save_value = cv2.imwrite(os.path.join(args.output_dir, image_file), image)
print("Image saved:", save_value)
sess.close()
def inpainting_api(image, mask):
FLAGS = ng.Config('inpaint.yml')
tf.reset_default_graph()
model = InpaintCAModel()
# image = cv2.imread(img_path)
# mask = cv2.imread(mask_path)
# cv2.imwrite('new.png', image - mask)
# mask = cv2.resize(mask, (0,0), fx=0.5, fy=0.5)
assert image.shape == mask.shape
image = crop(image)
mask = crop(mask)
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
'./model_logs/inpaint', from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
return result[0][:, :, ::-1]
def deepfill_model(checkpoint_dir):
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
model = InpaintCAModel()
input_image_ph = tf.placeholder(tf.float32, shape=(1, None, None, 3))
output = model.build_server_graph(input_image_ph, dynamic=True)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(checkpoint_dir,
from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
return input_image_ph, output, sess
class InpaintModel:
def __init__(self, model_path, max_size=768):
self.max_size = max_size
# build the graph
self.model = InpaintCAModel()
self.input_image = tf.placeholder(name='input_image',
shape=[1, max_size, max_size * 2, 3],
dtype=tf.float32)
net_output = self.model.build_server_graph(self.input_image)
net_output = (net_output + 1.) * 127.5
self.output_image = tf.saturate_cast(net_output, tf.uint8)
# restore pretrained model
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
self.sess = tf.Session(config=sess_config)
saver = tf.train.Saver([v for v in tf.global_variables()])
saver.restore(self.sess, model_path)
def inpaint(self, image, mask):
assert (image.shape == mask.shape)
h, w, _ = image.shape
assert (h <= self.max_size and w <= self.max_size)
grid = 8
ih, iw = h // grid * grid, w // grid * grid
input_im = np.zeros((self.max_size, self.max_size, 3), dtype=np.uint8)
input_mask = np.zeros((self.max_size, self.max_size, 3),
dtype=np.uint8)
input_im[:ih, :iw, :] = image[:ih, :iw, :]
input_mask[:ih, :iw, :] = mask[:ih, :iw, :]
input_im = np.expand_dims(input_im, 0)
input_mask = np.expand_dims(input_mask, 0)
input_image = np.concatenate([input_im, input_mask], axis=2)
output, = self.sess.run([self.output_image],
feed_dict={self.input_image: input_image})
image[:ih, :iw, :] = output[0][:ih, :iw, :]
return image
def __init__(self, batch_size, checkpoint_dir):
model = InpaintCAModel()
self.images_ph = tf.placeholder(tf.float32,
shape=[batch_size, 256, 512, 3])
# with tf.device('/device:GPU:0'):
# with tf.device('/cpu:0'):
output = model.build_server_graph(self.images_ph)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
self.output = output
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
self.assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
self.assign_ops.append(tf.assign(var, var_value))
print('Model loaded.')
self.pth_mean = np.ones((1, 3, 1, 1), dtype='float32')
self.pth_mean[0, :, 0, 0] = np.array([0.485, 0.456, 0.406])
self.pth_std = np.ones((1, 3, 1, 1), dtype='float32')
self.pth_std[0, :, 0, 0] = np.array([0.229, 0.224, 0.225])
self.upsample = torch.nn.Upsample(size=(256, 256), mode='bilinear')
self.downsample = torch.nn.Upsample(size=(224, 224), mode='bilinear')
# Create a session
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
# sess_config.log_device_placement = True
self.sess = tf.Session(config=sess_config)
self.sess.run(self.assign_ops)
h, w, _ = image.shape
grid = 8
image = image[:h//grid*grid, :w//grid*grid, :]
mask = mask[:h//grid*grid, :w//grid*grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv2.imwrite(args.output, result[0][:, :, ::-1])
if __name__ == "__main__":
FLAGS = ng.Config('inpaint.yml')
ng.get_gpus(1)
# os.environ['CUDA_VISIBLE_DEVICES'] =''
args = parser.parse_args()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
model = InpaintCAModel()
input_image_ph = tf.placeholder(tf.float32,
shape=(1, args.image_height,
args.image_width * 2, 3))
output = model.build_server_graph(FLAGS, input_image_ph)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir,
from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
with open(args.flist, 'r') as f:
class Synthesizer:
def __init__(self,
patch_size=512,
saved_model_path='./model_logs/release_places2_256'):
'''
Saved model weights url:
https://drive.google.com/drive/folders/1y7Irxm3HSHGvp546hZdAZwuNmhLUVcjO
'''
self.FLAGS = ng.Config('inpaint.yml')
self.model = InpaintCAModel()
self.checkpoint_dir = saved_model_path
self.patch_size = patch_size
self.sess_config = tf.ConfigProto()
self.sess_config.gpu_options.allow_growth = True
self._inpaint_input_placeholder = tf.placeholder(
shape=(1, self.patch_size, self.patch_size * 2, 3),
dtype=tf.float32)
self.sess = tf.Session(config=self.sess_config)
output = self.model.build_server_graph(self.FLAGS,
self._inpaint_input_placeholder,
reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
self._cached_inpaint_output = output
self.load_model()
def load_model(self):
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
self.checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
_ = self.sess.run(assign_ops)
def resize_image(self, image):
image = resize(image,
output_shape=(self.patch_size, self.patch_size, 3))
return image
def get_background(self, image, mask, reference_mask):
# image is ground_truth_image
merged_mask = mask + reference_mask
mask_3d = np.zeros((merged_mask.shape[0], merged_mask.shape[1], 3))
mask_3d[:, :, 0] = merged_mask
mask_3d[:, :, 1] = merged_mask
mask_3d[:, :, 2] = merged_mask
mask_3d = mask_3d * 255
mask_3d = mask_3d.astype(np.uint8)
image = img_as_ubyte(image)
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask_3d = mask_3d[:h // grid * grid, :w // grid * grid, :]
image = np.expand_dims(image, 0)
mask_3d = np.expand_dims(mask_3d, 0)
input_image = np.concatenate([image, mask_3d], axis=2)
input_image = input_image.astype(np.float32)
res = self.inpaint(input_image)
return res
def inpaint(self, input_image):
result = self.sess.run(
self._cached_inpaint_output,
feed_dict={self._inpaint_input_placeholder: input_image})
return result[0][:, :, ::-1]
def tweak_foreground(self, image):
"""
tweak foreground by applying random factor
"""
tweaked = image * np.random.uniform(0.1, 2)
tweaked = np.clip(tweaked, 0, 1)
# new_image = image + tweaked
# new_image *= (1.0/new_image.max())
return tweaked
def synthesize(self, image, mask, reference_mask):
inpainted_background_image = self.get_background(
image, mask, reference_mask)
inpainted_background_image = self.resize_image(
inpainted_background_image)
background_image = img_as_ubyte(inpainted_background_image.copy())
background_image[mask == True] = (0, 0, 0)
foreground_object = image.copy()
foreground_object[mask == False] = (0, 0, 0)
foreground_object = self.tweak_foreground(foreground_object)
background_image = background_image / 255.0
synthesized_image = background_image + foreground_object
return synthesized_image
with tf.Session(config=sess_config) as sess:
sess.run(tf.global_variables_initializer())
####
#### Might need to change the placeholder shape based on the
#### Model you are using
####
#### Can obtain this shape by using the `load_inputs` function
#### to load and preprocess the input image and mask, and print
#### its shape
####
#### input = load_inputs(input_image_path, input_mask_path)
#### print(input.shape)
#### >>> (1, 512, 1360, 3)
####
input_image_placeholder = tf.placeholder(tf.float32, shape=(1, 512, 1360, 3))
output = model.build_server_graph(input_image_placeholder, reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
if not vname.startswith("custom_input_image"):
var_value = tf.contrib.framework.load_variable(args.checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
if not os.path.exists(outdir):
os.makedirs(outdir)
if __name__ == "__main__":
FLAGS = ng.Config('inpaint.yml')
# ng.get_gpus(1)
# args, unknown = parser.parse_known_args()
model = InpaintCAModel()
images = os.listdir(maskdir)
images.sort()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
imin = tf.placeholder(tf.float32, (1, 256, 512, 3))
output = model.build_server_graph(FLAGS, imin)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
for x in images:
curdir = os.path.join(maskdir, x)
def iter_generate_image(image,
bbox,
id,
checkpoint_dir='model_logs/release_imagenet_256',
iterate=False):
# ng.get_gpus(1)
model = InpaintCAModel()
mask_obj = Mask_obj(top=bbox['y'],
left=bbox['x'],
width=bbox['w'],
height=bbox['h'])
mask = mask_obj.image
assert image.shape == mask.shape
print('Shape of image: {}'.format(image.shape))
iter1_input = preprocess_image(image, mask)
input_shape = iter1_input.shape
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
# input_image = tf.constant(input_image, dtype=tf.float32)
input_image = tf.placeholder(tf.float32, shape=input_shape)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
output_images = []
# iter1
result = sess.run(output, feed_dict={input_image: iter1_input})
iter1_image = result[0][:, :, ::-1]
output_images.append(iter1_image)
if iterate:
# iter2
iter2_submask_objs = generate_submask_objs([mask_obj])
temp_image = iter1_image
for s in iter2_submask_objs:
submask = s.image
temp_input = preprocess_image(temp_image, submask)
result = sess.run(output, feed_dict={input_image: temp_input})
temp_image = result[0][:, :, ::-1]
iter2_image = temp_image
output_images.append(iter2_image)
# iter3
# iter3_submask_objs = generate_submask_objs(iter2_submask_objs)
# temp_image = iter2_image
# for s in iter3_submask_objs:
# submask = s.image
# temp_input = preprocess_image(temp_image, submask)
# result = sess.run(output, feed_dict={input_image: temp_input})
# temp_image = result[0][:, :, ::-1]
# iter3_image = temp_image
# output_images.append(iter3_image)
print('hi')
return output_images[-1]
class GenerativeInpaintingTest(object):
def __init__(self, cfg):
self.cfg = cfg
self.log = {}
self.log_file = None
self.data_loader = None
def create_data_loaders(self):
self.data_loader = Dataset(self.cfg)
def create_model(self):
self.model = InpaintCAModel()
def plot_and_save(self, batch_idx, batch_data, subpath):
# load cfg
batch_size = self.cfg['batch_size']
batch_stride = self.cfg['batch_stride']
traverse_step = self.cfg['traverse_step']
def save(imgs, subpath, subsubpath):
res_dir = os.path.join(self.cfg['res_dir'], subpath, subsubpath)
for i, img in enumerate(imgs):
if img is None:
continue
sequence_path = os.path.join(
res_dir, 'sequence',
'{:03}.png'.format(batch_idx + i * batch_stride))
img_pil = Image.fromarray(img.astype(np.uint8))
if traverse_step <= batch_size // 2:
if batch_idx < batch_stride or i >= batch_size // 2:
img_pil.save(sequence_path)
else:
img_pil.save(sequence_path)
# load batch data
img_batch = batch_data['img_batch']
mask_batch = batch_data['mask_batch']
contour_batch = batch_data['contour_batch']
out_img_batch = batch_data['out_img_batch']
## use the original mask to fuse the result
out_img_batch_raw = img_batch * (
1 - mask_batch) + out_img_batch * mask_batch
out_img_batch = out_img_batch_raw.copy()
# draw mask boundary
for i in range(batch_size):
for con in contour_batch[i]:
for pt in con:
x, y = pt[0]
out_img_batch[i][y, x, :] = [255, 0, 0]
# store images
if not self.cfg['res_dir'] is None:
save(out_img_batch, subpath, 'full')
save(out_img_batch_raw, subpath, 'full_raw')
def build_dir(self, res_dir, subpath):
os.makedirs(os.path.join(res_dir, subpath))
res_type_list = ['full_raw', 'full']
for res_type in res_type_list:
sub_res_dir = os.path.join(res_dir, subpath, res_type)
os.makedirs(sub_res_dir)
os.makedirs(os.path.join(sub_res_dir, 'sequence'))
def train(self, sleep=False):
# randomly sleep seconds to avoid confilctc
if sleep:
t = np.random.randint(1, 20)
time.sleep(t)
# start time
start_time = time.time()
self.log['cfg'] = self.cfg.copy()
# build result folder
res_dir = self.cfg['res_dir']
if not res_dir is None:
if not os.path.exists(res_dir):
os.mkdir(res_dir)
# video folder
if not self.cfg['mask_name'] is None:
res_dir = os.path.join(
res_dir,
self.cfg['video_name'] + '_' + self.cfg['mask_name'])
else:
res_dir = os.path.join(res_dir, self.cfg['video_name'])
self.cfg['res_dir'] = res_dir
if not os.path.exists(res_dir):
os.mkdir(res_dir)
# pass folder
res_dir = os.path.join(self.cfg['res_dir'], '{:03}'.format(1))
os.mkdir(res_dir)
self.build_dir(res_dir, 'final')
self.log_file = open(os.path.join(self.cfg['res_dir'], 'log.txt'),
'w')
for key in sorted(self.cfg):
self.log_file.write(key + ' ' + str(self.cfg[key]) + '\n')
# print cfg to console
for key in sorted(self.cfg):
print(key, self.cfg[key])
self.create_data_loaders()
self.infer(0)
# save time
running_time = time.time() - start_time
print("Running time: ", running_time)
if not self.cfg['res_dir'] is None:
self.log_file.write("Running time: " + str(running_time) + '\n')
self.log_file.close()
def infer(self, pass_idx):
# load cfg
batch_size = self.cfg['batch_size']
print('Pass %d infer start...' % (pass_idx))
if not self.log_file is None:
self.log_file.write('Pass ' + str(pass_idx) + 'infer start...\n')
while True:
batch_data = self.data_loader.next_batch()
if batch_data is None:
break
batch_idx = batch_data['batch_idx']
self.infer_batch(batch_idx, batch_data)
def infer_batch(self, batch_idx, batch_data):
self.create_model()
tf.reset_default_graph()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = batch_data['input_batch']
input_image = tf.constant(input_image, dtype=tf.float32)
output = self.model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
self.cfg['checkpoint_dir'], from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
# print('Model loaded.')
result = sess.run(output)
out_img_batch = result
batch_data['out_img_batch'] = out_img_batch
self.plot_and_save(batch_idx, batch_data, '{:03}/final'.format(1))
# log
log_str = 'Batch {:05}'.format(batch_idx)
print(log_str)
sys.stdout.flush()
if not self.log_file is None:
self.log_file.write(log_str + '\n')
self.log_file.flush()
mask = np.expand_dims(mask, 0)
else:
mask = static_image[mask_index]
if args.image is not None:
image = image[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
if args.image is not None:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image, config=config)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
else:
logger.info('input image size: ' + str(static_image[0].shape))
#input_image = tf.concat([tf.expand_dims(static_image[0], 0), mask], axis=2) #TODO: handle this consistently
input_image = tf.concat([static_image[0], mask],
axis=2) #TODO: handle this consistently
logger.info('concat image size: ' + str(input_image.shape))
logger.info('mask size: ' + str(mask.shape))
output = model.build_server_graph(
input_image,
config=config,
exclusionmask=static_image[exclusionmask_index]
if config.EXC_MASKS else None)
def inpaintP(name):
import os
import time
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import cv2 as cv
import numpy as np
import tensorflow as tf
tf.compat.v1.logging.set_verbosity("ERROR")
import neuralgym as ng
from inpaint_model import InpaintCAModel
ROOT_DIR = os.path.abspath("./")
WAIT_DIR = os.path.abspath("./waiting")
WORK_DIR = os.path.abspath("./workspace")
CUR_DIR = os.path.join(WORK_DIR, name)
TMASK_DIR = os.path.join(WORK_DIR, name+"//tmask")
INPAINT_MODEL_PATH = os.path.join(ROOT_DIR, "model_logs/release_places2_256")
FLAGS = ng.Config('inpaint.yml')
model = InpaintCAModel()
image = cv.imread(os.path.join(CUR_DIR, f"{name}.png"))
mask = cv.imread(os.path.join(TMASK_DIR, "mask.png"))
filename = f'4#_{name}.png'
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 8
image = image[:h//grid*grid, :w//grid*grid, :]
mask = mask[:h//grid*grid, :w//grid*grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
sess_config = tf.ConfigProto()
sess_config.gpu_options.per_process_gpu_memory_fraction = 0.5
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(INPAINT_MODEL_PATH, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
cv.imwrite(os.path.join(CUR_DIR, filename), result[0][:, :, ::-1])
mask = cv2.imread(mask_folder + "/" + base_file_name)
assert image.shape == mask.shape
h, w, _ = image.shape
grid = 1
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image, reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
args.checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
