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 __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 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,
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 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 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 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 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 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 __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 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 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 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 __init__(self, batch_size, checkpoint_dir):
self.model = InpaintCAModel()
self.images_ph = tf.placeholder(tf.float32,
shape=[batch_size, 256, 512, 3])
# with tf.device('/gpu:0'):
# with tf.device('/cpu:0'):
# for i, d in enumerate(['/gpu:0', '/gpu:1', '/gpu:2', '/gpu:3']):
# with tf.device(d):
output = self.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.allow_soft_placement = True
# sess_config.log_device_placement = True
self.sess = tf.Session(config=sess_config)
self.sess.run(self.assign_ops)
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 main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='gated_convolution.yml',
help='path to config file')
parser.add_argument('--model', type=str, default='v1',
help='model type: v1 or v2')
parser.add_argument('--snapshot', type=str, default='model.npz',
help='path to snapshot')
parser.add_argument('--name', type=str, default='result.png',
help='file name to save results')
args = parser.parse_args()
config = Config(args.config_path)
if args.model == "v1":
config.FREE_FORM = False
inpaint_model = InpaintCAModel(config)
elif args.model == "v2":
inpaint_model = InpaintGCModel(config)
else:
assert False, "Model name '{args.model}' is invalid."
if config.GPU_ID >= 0:
chainer.cuda.get_device(config.GPU_ID).use()
inpaint_model.to_gpu()
if os.path.exists(args.snapshot):
serializers.load_npz(args.snapshot, inpaint_model)
else:
assert False, "Flie '{args.snapshot}' does not exist."
xp = inpaint_model.xp
# training data
test_dataset = Dataset(config, test=True, return_mask=True)
test_iter = chainer.iterators.SerialIterator(test_dataset, 8)
batch_and_mask = test_iter.next()
batch_data, mask_data = zip(*batch_and_mask)
batch_data = xp.array(batch_data)
mask = xp.array(mask_data)
batch_pos = batch_data / 127.5 - 1.
# edges = None
batch_incomplete = batch_pos * (1. - mask[:, :1])
# inpaint
with chainer.using_config("train", False), chainer.using_config("enable_backprop", False):
x1, x2, offset_flow = inpaint_model.inpaintnet(batch_incomplete, mask, config)
batch_complete = x2 * mask[:, :1] + batch_incomplete * (1. - mask[:, :1])
# visualization
viz_img = [batch_pos, batch_incomplete - mask[:, 1:] + mask[:, :1], batch_complete.data]
batch_w = len(viz_img)
batch_h = viz_img[0].shape[0]
viz_img = xp.concatenate(viz_img, axis=0)
viz_img = batch_postprocess_images(viz_img, batch_w, batch_h)
viz_img = cuda.to_cpu(viz_img)
Image.fromarray(viz_img).save(args.name)
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
if __name__ == "__main__":
# training data
FLAGS = ng.Config('inpaint.yml')
img_shapes = FLAGS.img_shapes
with open(FLAGS.data_flist[FLAGS.dataset][0]) as f:
fnames = f.read().splitlines()
if FLAGS.guided:
fnames = [(fname, fname[:-4] + '_edge.jpg') for fname in fnames]
img_shapes = [img_shapes, img_shapes]
data = ng.data.DataFromFNames(fnames,
img_shapes,
random_crop=FLAGS.random_crop,
nthreads=FLAGS.num_cpus_per_job)
images = data.data_pipeline(FLAGS.batch_size)
# main model
model = InpaintCAModel()
g_vars, d_vars, losses = model.build_graph_with_losses(FLAGS, images)
# validation images
if FLAGS.val:
with open(FLAGS.data_flist[FLAGS.dataset][1]) as f:
val_fnames = f.read().splitlines()
if FLAGS.guided:
val_fnames = [(fname, fname[:-4] + '_edge.jpg')
for fname in val_fnames]
# progress monitor by visualizing static images
for i in range(FLAGS.static_view_size):
static_fnames = val_fnames[i:i + 1]
static_images = ng.data.DataFromFNames(
static_fnames,
img_shapes,
nthreads=1,
image = image.astype(dtype)
image = image[np.newaxis, ...]
fo = rasterio.open(filename, 'w', **kwargs)
fo.write(image)
fo.close()
if __name__ == "__main__":
FLAGS = ng.Config('inpaint.yml')
# ng.get_gpus(1)
args, unknown = parser.parse_known_args()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
model = InpaintCAModel()
image, meta = read_tif_file(args.image)
print('Read mask form {}'.format(args.mask))
mask = cv2.imread(args.mask, -1)
if len(mask.shape) == 3:
mask = mask[:, :, 0]
im_min = image.min()
im_max = image.max()
image = cv2.normalize(image, None, 255, 0, cv2.NORM_MINMAX, cv2.CV_32F)
image_tiles = tiling_image(image, args.num_tiles_x, args.num_tiles_y)
mask_tiles = tiling_image(mask, args.num_tiles_x, args.num_tiles_y)
result = np.zeros(image.shape)
ntiles = len(image_tiles)
for i in range(ntiles):
startY, endY = image_tiles[i][1]
config = ng.Config('inpaint.yml')
if config.GPU_ID != -1:
ng.set_gpus(config.GPU_ID)
else:
ng.get_gpus(config.NUM_GPUS)
# training data
with open(config.DATA_FLIST[config.DATASET][0]) as f:
fnames = f.read().splitlines()
data = ng.data.DataFromFNames(fnames,
config.IMG_SHAPES,
random_crop=config.RANDOM_CROP)
images = data.data_pipeline(config.BATCH_SIZE)
# main model
model = InpaintCAModel()
g_vars, d_vars, losses = model.build_graph_with_losses(images,
config=config)
# validation images
if config.VAL:
with open(config.DATA_FLIST[config.DATASET][1]) as f:
val_fnames = f.read().splitlines()
# progress monitor by visualizing static images
for i in range(config.STATIC_VIEW_SIZE):
static_fnames = val_fnames[i:i + 1]
static_images = ng.data.DataFromFNames(
static_fnames,
config.IMG_SHAPES,
nthreads=1,
random_crop=config.RANDOM_CROP).data_pipeline(1)
parser = argparse.ArgumentParser()
parser.add_argument('--image', default='', type=str,
help='The filename of image to be completed.')
parser.add_argument('--mask', default='', type=str,
help='The filename of mask, value 255 indicates mask.')
parser.add_argument('--output', default='output.png', type=str,
help='Where to write output.')
parser.add_argument('--checkpoint_dir', default='', type=str,
help='The directory of tensorflow checkpoint.')
if __name__ == "__main__":
ng.get_gpus(num_gpus=1, dedicated=False)
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)
parser.add_argument('--checkpoint_dir',
default='',
type=str,
help='The directory of tensorflow checkpoint.')
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))
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
count += 1
if count == 1:
fnames = f.read().splitlines()
elif count == 2:
fnames = [(l.split(' ')[0], l.split(' ')[1]) for l in f.read().splitlines()]
elif count == 3:
fnames = [(l.split(' ')[0], l.split(' ')[1], l.split(' ')[2]) for l in f.read().splitlines()]
else:
print('invalid data count')
exit()
data = ng.data.DataFromFNames(
fnames, config.IMG_SHAPES, random_crop=config.RANDOM_CROP, gamma=config.GAMMA, exposure=config.EXPOSURE, random_flip = config.RANDOM_FLIP)
images = data.data_pipeline(config.BATCH_SIZE)
# main model
model = InpaintCAModel()
g_vars, d_vars, losses = model.build_graph_with_losses(
images[0], config=config, exclusionmask=images[exclusionmask_index] if config.EXC_MASKS else None, mask=None if config.GEN_MASKS else images[mask_index])
# validation images
if config.VAL:
with open(config.DATA_FLIST[config.DATASET][1]) as f:
val_fnames = [(l.split(' ')[0], l.split(' ')[1]) for l in f.read().splitlines()]
# progress monitor by visualizing static images
for i in range(config.STATIC_VIEW_SIZE):
static_fnames = val_fnames[i:i+1]
static_images = ng.data.DataFromFNames(
static_fnames, config.IMG_SHAPES, nthreads=1,
random_crop=config.RANDOM_CROP, random_flip=config.RANDOM_FLIP).data_pipeline(1)
static_inpainted_images = model.build_static_infer_graph(
static_images[0], config, name='static_view/%d' % i, exclusionmask=images[exclusionmask_index] if config.EXC_MASKS else None)
# training settings
