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 convert_image(img, img_size):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = preprocessing(img, img_size)
img = np.expand_dims(img, axis=0)
img = np.asarray(img)
return img
