def main(checkpoint_dir, style_name):
ckpt_dir = '../checkpoint/' + 'generator_' + style_name + '_weight'
check_folder(ckpt_dir)
placeholder = tf.placeholder(tf.float32, [1, None, None, 3],
name='generator_input')
with tf.variable_scope("generator", reuse=False):
_ = generator.G_net(placeholder).fake
generator_var = [
var for var in tf.trainable_variables()
if var.name.startswith('generator')
]
saver = tf.train.Saver(generator_var)
gpu_options = tf.GPUOptions(allow_growth=True)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)) as sess:
sess.run(tf.global_variables_initializer())
# load model
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
if ckpt and ckpt.model_checkpoint_path:
print(ckpt.model_checkpoint_path)
ckpt_name = os.path.basename(
ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
counter = ckpt_name.split('-')[-1]
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
return
info = save(saver, sess, ckpt_dir, style_name + '-' + counter)
print(f'save over : {info} ')
def test(checkpoint_dir,
style_name,
test_dir,
if_adjust_brightness,
img_size=[256, 256]):
# tf.reset_default_graph()
test_files = glob('{}/*.*'.format(test_dir))
result_dir = 'results/' + style_name
check_folder(result_dir)
# test_real = tf.placeholder(tf.float32, [1, 256, 256, 3], name='test')
test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
with tf.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
generator_var = [
var for var in tf.trainable_variables()
if var.name.startswith('generator')
]
saver = tf.train.Saver(generator_var)
gpu_options = tf.GPUOptions(allow_growth=True)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(
ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
return
# FLOPs
# stats_graph(tf.get_default_graph())
begin = time.time()
for sample_file in tqdm(test_files):
# print('Processing image: ' + sample_file)
sample_image = np.asarray(load_test_data(sample_file, img_size))
image_path = os.path.join(
result_dir, '{0}'.format(os.path.basename(sample_file)))
fake_img = sess.run(test_generated,
feed_dict={test_real: sample_image})
if if_adjust_brightness:
save_images(fake_img, image_path, sample_file)
else:
save_images(fake_img, image_path, None)
end = time.time()
print(f'test-time: {end-begin} s')
print(f'one image test time : {(end-begin)/len(test_files)} s')
print(f'result path: {result_dir}')
def generator(self, x_init, reuse=False, scope="generator"):
if self.light:
with tf.variable_scope(scope, reuse=reuse):
G = generator_lite.G_net(x_init)
return G.fake
else:
with tf.variable_scope(scope, reuse=reuse):
G = generator.G_net(x_init)
return G.fake
def cvt2anime_video(video, output, checkpoint_dir, output_format='MP4V'):
'''
output_format: 4-letter code that specify codec to use for specific video type. e.g. for mp4 support use "H264", "MP4V", or "X264"
'''
gpu_stat = bool(len(tf.config.experimental.list_physical_devices('GPU')))
if gpu_stat:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
gpu_options = tf.GPUOptions(allow_growth=gpu_stat)
test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
with tf.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
saver = tf.train.Saver()
# load video
vid = cv2.VideoCapture(video)
vid_name = os.path.basename(video)
total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
codec = cv2.VideoWriter_fourcc(*output_format)
tfconfig = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
with tf.Session(config=tfconfig) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(checkpoint_dir) # checkpoint file information
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(os.path.join(checkpoint_dir, ckpt_name)))
else:
print(" [*] Failed to find a checkpoint")
return
video_out = cv2.VideoWriter(os.path.join(output, vid_name.rsplit('.', 1)[0] + "_AnimeGANv2.mp4"), codec, fps, (width, height))
pbar = tqdm(total=total, ncols=80)
pbar.set_description(f"Making: {os.path.basename(video).rsplit('.', 1)[0] + '_AnimeGANv2.mp4'}")
while True:
ret, frame = vid.read()
if not ret:
break
frame = np.asarray(np.expand_dims(process_image(frame),0))
fake_img = sess.run(test_generated, feed_dict={test_real: frame})
fake_img = post_precess(fake_img, (width, height))
video_out.write(cv2.cvtColor(fake_img, cv2.COLOR_BGR2RGB))
pbar.update(1)
pbar.close()
vid.release()
video_out.release()
return os.path.join(output, vid_name.rsplit('.', 1)[0] + "_AnimeGANv2.mp4")
def cvt2anime_dir(inputdir,
outputdir,
checkpoint_dir,
show_stats=False,
img_size=(256, 256)):
gpu_stat = bool(len(tf.config.experimental.list_physical_devices('GPU')))
if gpu_stat:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
gpu_options = tf.compat.v1.GPUOptions(allow_growth=gpu_stat)
test_real = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3],
name='test')
with tf.compat.v1.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
# get image list
img_files = glob('{}/*.*'.format(inputdir))
tfconfig = tf.compat.v1.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)
with tf.compat.v1.Session(config=tfconfig) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
saver = tf.compat.v1.train.Saver()
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(
ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
return
# show FLOPs
if show_stats:
stats_graph(tf.compat.v1.get_default_graph())
# convert
for imgfile in tqdm(img_files):
img = cv2.imread(imgfile).astype(np.float32)
img_cvt = convert_image(img, img_size)
fake_img = sess.run(test_generated, feed_dict={test_real: img_cvt})
fake_img_inv = inverse_image(fake_img)
output_path = os.path.join(outputdir,
'{0}'.format(os.path.basename(imgfile)))
cv2.imwrite(output_path, fake_img_inv)
def test(img_size=[256, 256]):
checkpoint_dir = 'checkpoint/generator_Hayao_weight'
# tf.reset_default_graph()
result_dir = 'static/client/img'
test_dir = 'static/client/img'
check_folder(result_dir)
test_files = glob('{}/*.*'.format(test_dir))
test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
with tf.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
saver = tf.train.Saver()
gpu_options = tf.GPUOptions(allow_growth=True)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(
ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
begin = time.time()
for sample_file in tqdm(test_files):
sample_image = np.asarray(load_test_data(sample_file, img_size))
image_path = os.path.join(
result_dir, '{0}'.format(os.path.basename(sample_file)))
fake_img = sess.run(test_generated,
feed_dict={test_real: sample_image})
#adjustBrightness
save_images(fake_img, image_path, sample_file)
#save_images(fake_img, image_path, None)
end = time.time()
print(f'test-time: {end-begin} s')
print(f'one image test time : {(end-begin)/len(test_files)} s')
def cvt2anime_video(video, output, checkpoint_dir, output_format='MP4V', if_adjust_brightness=False, img_size=(256,256)):
'''
output_format: 4-letter code that specify codec to use for specific video type. e.g. for mp4 support use "H264", "MP4V", or "X264"
'''
# tf.reset_default_graph()
# check_folder(result_dir)
gpu_stat = bool(len(tf.config.experimental.list_physical_devices('GPU')))
if gpu_stat:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
gpu_options = tf.GPUOptions(allow_growth=gpu_stat)
test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
with tf.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
# load video
vid = cv2.VideoCapture(video)
vid_name = os.path.basename(video)
total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
fps = vid.get(cv2.CAP_PROP_FPS)
# codec = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
codec = cv2.VideoWriter_fourcc(*output_format)
tfconfig = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
with tf.Session(config=tfconfig) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(checkpoint_dir) # checkpoint file information
saver = tf.train.Saver()
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
return
# determine output width and height
ret, img = vid.read()
if img is None:
print('Error! Failed to determine frame size: frame empty.')
return
img = preprocessing(img, img_size)
height, width = img.shape[:2]
out = cv2.VideoWriter(os.path.join(output, vid_name), codec, fps, (width, height))
pbar = tqdm(total=total)
vid.set(cv2.CAP_PROP_POS_FRAMES, 0)
while ret:
ret, frame = vid.read()
if frame is None:
print('Warning: got empty frame.')
continue
img = convert_image(frame, img_size)
fake_img = sess.run(test_generated, feed_dict={test_real: img})
fake_img = inverse_image(fake_img)
if if_adjust_brightness:
fake_img = cv2.cvtColor(adjust_brightness_from_src_to_dst(fake_img, frame), cv2.COLOR_BGR2RGB)
else:
fake_img = cv2.cvtColor(fake_img, cv2.COLOR_BGR2RGB)
fake_img = cv2.resize(fake_img, (width, height))
out.write(fake_img)
pbar.update(1)
pbar.close()
vid.release()
# cv2.destroyAllWindows()
return os.path.join(output, vid_name)
def generator(self, x_init, reuse=False, scope="generator"):
with tf.variable_scope(scope, reuse=reuse):
G = generator.G_net(x_init)
return G.fake
def cvt2anime_video(video,
output,
checkpoint_dir,
output_format='MP4V',
show_stats=False,
img_size=(256, 256)):
'''
output_format: 4-letter code that specify codec to use for specific video type. e.g. for mp4 support use "H264", "MP4V", or "XVID"
'''
# tf.reset_default_graph()
# check_folder(result_dir)
gpu_stat = bool(len(tf.config.experimental.list_physical_devices('GPU')))
if gpu_stat:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
gpu_options = tf.compat.v1.GPUOptions(allow_growth=gpu_stat)
test_real = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3],
name='test')
with tf.compat.v1.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
# load video
try:
vid = cv2.VideoCapture(int(video))
except:
vid = cv2.VideoCapture(video)
#width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
#height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
# codec = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
codec = cv2.VideoWriter_fourcc(*output_format)
tfconfig = tf.compat.v1.ConfigProto(allow_soft_placement=True,
gpu_options=gpu_options)
with tf.compat.v1.Session(config=tfconfig) as sess:
# tf.global_variables_initializer().run()
# load model
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
saver = tf.compat.v1.train.Saver()
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(
ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
return
# FLOPs
if show_stats:
stats_graph(tf.compat.v1.get_default_graph())
ret, img = vid.read()
if img is None:
print('Error! Failed to determine frame size: frame empty.')
return
img = convert_image(img, img_size)
fake_img = sess.run(test_generated, feed_dict={test_real: img})
fake_img = inverse_image(fake_img)
height, width = fake_img.shape[:2]
out = cv2.VideoWriter(output, codec, fps, (width, height))
bar_total = 100
bar_step = np.round(bar_total / total, 2)
pbar = tqdm(total=bar_total)
vid.set(cv2.CAP_PROP_POS_FRAMES, 0)
while ret:
ret, img = vid.read()
if img is None:
print('Warning: got empty frame.')
continue
img = convert_image(img, img_size)
fake_img = sess.run(test_generated, feed_dict={test_real: img})
fake_img = inverse_image(fake_img)
# cv2.imwrite(f'results/cut_{i:03}.jpg', fake_img)
out.write(fake_img)
pbar.update(bar_step)
# cv2.imshow('output', fake_img)
if cv2.waitKey(1) == ord('q'):
break
pbar.close()
vid.release()
cv2.destroyAllWindows()
# Flask utils
from flask import Flask, send_from_directory, request, render_template
from werkzeug.utils import secure_filename
from gevent.pywsgi import WSGIServer
# Define a flask app
app = Flask(__name__)
"""=============================读取模型==================================="""
# 存放模型的文件夹
checkpoint_dir = 'checkpoint/' + 'AnimeGAN_Hayao_lsgan_300_300_1_3_10'
# 输入的真实图片留存区，执行的时候再赋具体size的大小
test_real = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3],
name='test')
# 先初始化generator网络
with tf.compat.v1.variable_scope("generator", reuse=False):
test_generated = generator.G_net(test_real).fake
saver = tf.compat.v1.train.Saver()
sess = tf.compat.v1.Session(config=tf.ConfigProto(allow_soft_placement=True))
# 加载模型
ckpt = tf.train.get_checkpoint_state(
checkpoint_dir) # checkpoint file information
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line
saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
print(" [*] Success to read {}".format(ckpt_name))
else:
print(" [*] Failed to find a checkpoint")
"""=============================================================================="""
def encoder(images, style_size=1, keep_prob=1.0, phase_train=True, weight_decay=0.0, reuse=None, scope='Encoders'):
print("ENCODER INPUT SHAPE: ", images.shape)
with tf.compat.v1.variable_scope(scope, reuse=reuse):
with slim.arg_scope([slim.conv2d, slim.conv2d_transpose, slim.fully_connected],
activation_fn=tf.nn.relu,
# weights_initializer=tf.contrib.layers.xavier_initializer(),
weights_initializer=tf.compat.v1.keras.initializers.VarianceScaling(scale=2.0),
weights_regularizer=tf.keras.regularizers.l2(0.5 * (weight_decay))):
with slim.arg_scope([slim.dropout, slim.batch_norm], is_training=phase_train):
with slim.arg_scope([slim.fully_connected],
normalizer_fn=layer_norm, normalizer_params=None):
print('{} input shape:'.format(scope), [dim.value for dim in images.shape])
batch_size = tf.shape(input=images)[0]
k = 64
# with tf.compat.v1.variable_scope('StyleEncoder'):
# with slim.arg_scope([slim.conv2d, slim.conv2d_transpose, slim.fully_connected],
# normalizer_fn=None, normalizer_params=None):
#
# print('-- StyleEncoder')
#
# net = images
#
# net = conv(net, k, 7, stride=1, pad=3, scope='conv0')
# print('module conv0 shape:', [dim.value for dim in net.shape])
#
# net = conv(net, 2*k, 4, stride=2, scope='conv1')
# print('module conv1 shape:', [dim.value for dim in net.shape])
#
# net = conv(net, 4*k, 4, stride=2, scope='conv2')
# print('module conv2 shape:', [dim.value for dim in net.shape])
#
#
# encoded_style = net
#
# net = slim.avg_pool2d(net, net.shape[1:3], padding='VALID', scope='global_pool')
# net = slim.flatten(net)
#
# style_vec = slim.fully_connected(net, style_size, activation_fn=None, normalizer_fn=None, scope='fc1')
# print('module fc1 shape:', [dim.value for dim in net.shape])
# style_vec = tf.identity(style_vec, name='style_vec')
with tf.compat.v1.variable_scope('AnimeGANGenerator'):
print( "AnimeGan input Image shape: :", [images.shape])
#
# test_real = tf.placeholder(tf.float32, [None, 256, 256, 3], name='test')
#
#
# with tf.variable_scope("generator", reuse=False):
# test_generated = generator.G_net(test_real).fake
# saverA = tf.train.Saver()
#
# print("TEST \n \n")
#
# gpu_options = tf.GPUOptions(allow_growth=False)
#
# with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) as sess2:
# # tf.global_variables_initializer().run()
# # load model
# checkpoint_dir = "pretrained_hayao"
# print("lOADIN ANIME MODEL \n \n \n")
# ckpt = tf.train.get_checkpoint_state(checkpoint_dir) # checkpoint file information
#
# if ckpt and ckpt.model_checkpoint_path:
# ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line
# print(" [*]CHECKPOINT NAME {}".format(ckpt_name))
# saverA.restore(sess2, os.path.join(checkpoint_dir, ckpt_name))
# print(" [*] Success to read {}".format(ckpt_name))
# else:
# print(" [*] Failed to find a checkpoint")
# return
# # print('Processing image: ' + sample_file)
# #sample_image = np.asarray(load_test_data(sample_file, img_size))
# images_rendered = sess2.run(test_generated, feed_dict = {test_real : images})
images_rendered = generator.G_net(images).fake
#K.clear_session()
#images_rendered = testAnimeGan.test(checkpoint_dir="pretrained_hayao", test_dir = "testData",if_adjust_brightness = True,style_name = "H")
print( "rendered Image shape: :", [dim.value for dim in images_rendered.shape])
#imshow the image here
# Transform textures
with tf.compat.v1.variable_scope('ContentEncoder'):
with slim.arg_scope([slim.conv2d, slim.conv2d_transpose, slim.fully_connected],
normalizer_fn=instance_norm, normalizer_params=None):
print('-- ContentEncoder')
net = images
net = conv(net, k, 7, stride=1, pad=3, scope='conv0')
print('module conv0 shape:', [dim.value for dim in net.shape])
net = conv(net, 2*k, 4, stride=2, scope='conv1')
print('module conv1 shape:', [dim.value for dim in net.shape])
net = conv(net, 4*k, 4, stride=2, scope='conv2')
print('module conv2 shape:', [dim.value for dim in net.shape])
for i in range(3):
net_ = conv(net, 4*k, 3, scope='res{}_0'.format(i))
net += conv(net_, 4*k, 3, activation_fn=None, biases_initializer=None, scope='res{}_1'.format(i))
print('module res{} shape:'.format(i), [dim.value for dim in net.shape])
encoded = net
return encoded, images_rendered#, style_vec
