def main(argv):
model_file_path = os.path.join(FLAGS.vgg_model, vgg19.MODEL_FILE_NAME)
vgg_net = vgg19.VGG19(model_file_path)
content_images = utils.get_files(FLAGS.train)
style_image = utils.load_image(FLAGS.style)
# create a map for content layers info
CONTENT_LAYERS = {}
for layer, weight in zip(CONTENT_LAYERS_NAME, CONTENT_LAYER_WEIGHTS):
CONTENT_LAYERS[layer] = weight
# create a map for style layers info
STYLE_LAYERS = {}
for layer, weight in zip(STYLE_LAYERS_NAME, STYLE_LAYER_WEIGHTS):
STYLE_LAYERS[layer] = weight
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
trainer = style_transfer_trainer.StyleTransferTrainer(
session=sess,
content_layer_ids=CONTENT_LAYERS,
style_layer_ids=STYLE_LAYERS,
content_images=content_images,
style_image=add_one_dim(style_image),
net=vgg_net,
num_epochs=FLAGS.num_epochs,
batch_size=FLAGS.batch_size,
content_weight=FLAGS.content_weight,
style_weight=FLAGS.style_weight,
tv_weight=FLAGS.tv_weight,
learn_rate=FLAGS.learn_rate,
save_path=FLAGS.output,
check_period=FLAGS.checkpoint_every,
max_size=FLAGS.max_size or None)
trainer.train()
def main():
# parse arguments
args = parse_args()
if args is None:
exit()
# initiate VGG19 model
model_file_path = args.vgg_model + '/' + vgg19.MODEL_FILE_NAME
vgg_net = vgg19.VGG19(model_file_path)
# get file list for training
content_images = utils.get_files(args.trainDB_path)
# load style image
style_image = utils.load_image(args.style)
# create a map for content layers info
CONTENT_LAYERS = {}
for layer, weight in zip(args.content_layers, args.content_layer_weights):
CONTENT_LAYERS[layer] = weight
# create a map for style layers info
STYLE_LAYERS = {}
for layer, weight in zip(args.style_layers, args.style_layer_weights):
STYLE_LAYERS[layer] = weight
# open session
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.5
sess = tf.Session(config=config)
# build the graph for train
trainer = style_transfer_trainer.StyleTransferTrainer(
session=sess,
content_layer_ids=CONTENT_LAYERS,
style_layer_ids=STYLE_LAYERS,
content_images=content_images,
style_image=add_one_dim(style_image),
net=vgg_net,
num_epochs=args.num_epochs,
batch_size=args.batch_size,
content_weight=args.content_weight,
style_weight=args.style_weight,
tv_weight=args.tv_weight,
learn_rate=args.learn_rate,
save_path=args.output,
check_period=args.checkpoint_every,
test_image=args.test,
max_size=args.max_size,
)
# launch the graph in a session
trainer.train()
# close session
sess.close()
for layer, weight in zip(STYLE_LAYERS_NAME, STYLE_LAYER_WEIGHTS):
STYLE_LAYERS[layer] = weight
cap = cv2.VideoCapture(VIDEO_FILE)
fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
video_out = None
with tf.Session() as sess:
trainer = style_transfer_trainer.StyleTransferTrainer(
session=sess,
content_layer_ids=CONTENT_LAYERS,
style_layer_ids=STYLE_LAYERS,
content_images=[],
style_image=add_one_dim(style_image),
net=vgg_net,
num_epochs=0,
batch_size=1,
content_weight=0,
style_weight=0,
tv_weight=0,
learn_rate=0,
save_path=MODEL_FILE,
check_period=0,
max_size=None)
trainer.prepare()
while cap.isOpened():
ret, frame = cap.read()
if not video_out:
video_out = cv2.VideoWriter(VIDEO_OUT_FILE, fourcc, 24,
(frame.shape[1], frame.shape[0]))