引言
你这个模型, 它复现起来难吗?
我一水博客的, 能给你看复现不出来的算法?
效果展示
将梵高的代表作<星空>的风格, 迁移到视频上
准备工作
- 华强买瓜原版视频, 不需要音频. 毕竟处理的图像不包含音频信息. 最好也不要包含字幕, 字幕会影响整体的效果. 音频和字幕可以后期用PR剪辑上去
- 风格图像, 随便挑几幅, 最好是印象派的画作, 非常具有风格
- Python3.6.7
- Tensorflow2.0.0
- OpenCV-Python3.4.2.16
具体步骤
让我们新建一个HuaQiangBuyWatermelon的项目…
业务逻辑
原来我计划的流程是这样的
原视频 = 读取视频
写入视频 = 视频写入对象
风格图像 = 读取风格图像
for each 当前帧 in 原视频
转换帧 = 转换风格(当前帧, 风格图像)
写入视频(转换帧)
end for
但是过程中OpenCV的VideoWriter出了点问题, 写不进去. 无奈只能采用以下的逻辑
for each 当前帧 in 原视频
保存(当前帧.jpg)
end for
for each 当前图像 in 所有保存的帧
转换帧 = 转换风格(当前图像, 风格图像)
保存(转换帧)
end for
最后得到了每一帧的经过风格转换的图像, 再借助PR将这些所有的转换帧剪成视频
日志
整个模型虽然简单, 但是非常耗时, 我截取的华强买瓜视频一共6900帧. 为了DEBUG方便, 日志是一个好习惯.
新建logger.py文件
为了方便DEBUG, 我会在每行日志的最前面显示输出的类型, 比如INFO, WARNING, ERROR
import time
logger = open('./log.txt', 'a+')
def log_string(out_str, dtype='INFO', to_log_file=True):
"""
保存日志
:param to_log_file:
:param dtype:
:param out_str:
"""
global logger
if not isinstance(out_str, str):
out_str = str(out_str)
if '\\r' == out_str:
print()
if to_log_file:
logger.write('\n')
logger.flush()
else:
local_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
out_str = '[{:^7s}] [{}] '.format(dtype, local_time) + out_str
if 'INFO' == dtype:
print(out_str)
elif 'WARNING' == dtype:
print("\033[0;34m{}\033[0m".format(out_str))
else:
print("\033[0;31m{}\033[0m".format(out_str))
if to_log_file:
logger.write(out_str + '\n')
logger.flush()
保存的日志大概如下:
转换风格
这里用到了神经网络, 简单来说就是通过一个深层的VGG网络去提取图像更深层的信息, 难度在于损失函数的设计, 这里不赘述, 感兴趣的读者请查阅其它文献.
train()接受4个参数:
- content_image 待转换的图像
- style_image 风格图像, 将这张图像的风格迁移到content_image上
- num_epochs epoch的数量, 训练的周期数
- remain 剩余的待处理的图像数量, 用来估计剩余时间ETA
该函数返回风格转换完毕的图像
新建train.py文件
def train(content_image, style_image, num_epochs, remain):
# 截取0-1的浮点数，超范围部分被截取
def clip_0_1(image_):
return tf.clip_by_value(image_, clip_value_min=0.0, clip_value_max=1.0)
# 损失函数
def style_content_loss(outputs):
style_outputs = outputs['style']
content_outputs = outputs['content']
# 风格损失值，就是计算方差
style_loss = tf.add_n([tf.reduce_mean((style_outputs[name] - style_targets[name]) ** 2)
for name in style_outputs.keys()])
# 权重值平均到每层，计算总体风格损失值
style_loss *= style_weight / num_style_layers
# 内容损失值，也是计算方差
content_loss = tf.add_n([tf.reduce_mean((content_outputs[name] - content_targets[name]) ** 2)
for name in content_outputs.keys()])
content_loss *= content_weight / num_content_layers
# 总损失值
loss = style_loss + content_loss
return loss
# 一次训练
@tf.function()
def train_step(image_):
with tf.GradientTape() as tape:
# 抽取风格层、内容层输出
outputs = extractor(image_)
# 计算损失值
loss = style_content_loss(outputs)
# 梯度下降
grad = tape.gradient(loss, image_)
# 应用计算后的新参数，注意这个新值不是应用到网络
# 作为训练完成的vgg网络，其参数前面已经设定不可更改
# 这个参数实际将应用于原图
# 以求取，新图片经过网络后，损失值最小
opt.apply_gradients([(grad, image_)])
# 更新图片，用新图片进行下次训练迭代
image_.assign(clip_0_1(image_))
# 定义最能代表内容特征的网络层
content_layers = ['block5_conv2']
# 定义最能代表风格特征的网络层
style_layers = ['block1_conv1',
'block2_conv1',
'block3_conv1',
'block4_conv1',
'block5_conv1']
# 神经网络层的数量
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)
# 使用自定义模型建立一个抽取器
extractor = StyleContentModel(style_layers, content_layers)
# 设定风格特征的目标，即最终生成的图片，希望风格上尽量接近风格图片
style_targets = extractor(style_image)['style']
# 设定内容特征的目标，即最终生成的图片，希望内容上尽量接近内容图片
content_targets = extractor(content_image)['content']
# 内容图片转换为张量
image = tf.Variable(content_image)
# 优化器
opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
# 预定义风格和内容在最终结果中的权重值，用于在损失函数中计算总损失值
style_weight = 1e-2
content_weight = 1e4
start = time.time()
epochs = num_epochs
steps_per_epoch = 50
step = 0
for n in range(epochs):
start2 = time.time()
for m in range(steps_per_epoch):
step += 1
train_step(image)
end2 = time.time()
log_string('time of epoch {}: {:.2f}s'.format(n+1, end2-start2))
end = time.time()
log_string("Total time: {:.2f}s, ETA: {:.2f}min".format(end - start, remain*(end - start)/60))
########################################
return image[0].numpy()
其中用到的各种函数如下
from __future__ import absolute_import, division, print_function, unicode_literals
import time
from abc import ABC
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from logger import log_string
# 设置绘图窗口参数，用于图片显示
mpl.rcParams['figure.figsize'] = (13, 10)
mpl.rcParams['axes.grid'] = False
def load_img(path_to_img):
max_dim = 512
# 读取二进制文件
img = tf.io.read_file(path_to_img)
# 做JPEG解码，这时候得到宽x高x色深矩阵，数字0-255
img = tf.image.decode_jpeg(img)
# 类型从int转换到32位浮点，数值范围0-1
img = tf.image.convert_image_dtype(img, tf.float32)
# 减掉最后色深一维，获取到的相当于图片尺寸（整数），转为浮点
shape = tf.cast(tf.shape(img)[:-1], tf.float32)
# 获取图片长端
long = max(shape)
# 以长端为比例缩放，让图片成为512x???
scale = max_dim / long
new_shape = tf.cast(shape * scale, tf.int32)
# 实际缩放图片
img = tf.image.resize(img, new_shape)
# 再扩展一维，成为图片数字中的一张图片（1，长，宽，色深）
img = img[tf.newaxis, :]
return img
# 定义一个工具函数，帮助建立得到特定中间层输出结果的新模型
def vgg_layers(layer_names):
""" Creates a vgg model that returns a list of intermediate output values."""
# 定义使用ImageNet数据训练的vgg19网络
vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
# 已经经过了训练，所以锁定各项参数避免再次训练
vgg.trainable = False
# 获取所需层的输出结果
outputs = [vgg.get_layer(name).output for name in layer_names]
# 最终返回结果是一个模型，输入是图片，输出为所需的中间层输出
model = tf.keras.Model([vgg.input], outputs)
return model
# 定义函数计算风格矩阵，这实际是由抽取出来的5个网络层的输出计算得来的
def gram_matrix(input_tensor):
result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
input_shape = tf.shape(input_tensor)
num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
return result / num_locations
# 自定义keras模型
class StyleContentModel(tf.keras.models.Model, ABC):
def __init__(self, style_layers_, content_layers_):
super(StyleContentModel, self).__init__()
# 自己的vgg模型，包含上面所列的风格抽取层和内容抽取层
self.vgg = vgg_layers(style_layers_ + content_layers_)
self.style_layers = style_layers_
self.content_layers = content_layers_
self.num_style_layers = len(style_layers_)
# vgg各层参数锁定不再参数训练
self.vgg.trainable = False
def call(self, input_, **kwargs):
# 输入的图片是0-1范围浮点，转换到0-255以符合vgg要求
input_ = input_ * 255.0
# 对输入图片数据做预处理
preprocessed_input = tf.keras.applications.vgg19.preprocess_input(input_)
# 获取风格层和内容层输出
outputs = self.vgg(preprocessed_input)
# 输出实际是一个数组，拆分为风格输出和内容输出
style_outputs, content_outputs = (
outputs[:self.num_style_layers],
outputs[self.num_style_layers:])
# 计算风格矩阵
style_outputs = [gram_matrix(style_output)
for style_output in style_outputs]
# 转换为字典
content_dict = {content_name: value
for content_name, value
in zip(self.content_layers, content_outputs)}
# 转换为字典
style_dict = {style_name: value
for style_name, value
in zip(self.style_layers, style_outputs)}
# 返回内容和风格结果
return {'content': content_dict, 'style': style_dict}
读取视频并保存每一帧
新建generate.py文件
这里使用OpenCV的VideoCapture读取视频, 然后一帧一帧保存下来
这个函数只要运行一次就行了, 然后就可以注释掉了
def split_frames():
cap = cv.VideoCapture('./video/video.mp4')
cnt = 0
while True:
ret, frame = cap.read()
if not ret:
break
cnt += 1
log_string('saving frame {}'.format(cnt))
cv.imwrite('./contents/{}.jpg'.format(cnt), frame)
cap.release()
log_string('end')
log_string('\\r')
log_string('\\r')
最后, 再将保存好的这些原视频的帧, 一张一张传递给train()函数即可.
但是这里需要注意, os.listdir()返回的目录下的所有文件, 是按文件名的字典序排序的, 也就是说, 假如有1.jpg, 2.jpg, …, 20.jpg这20个图像, 它返回的结果是
1.jpg, 11.jpg, 12.jpg, 13.jpg, 14.jpg, 15.jpg, 16.jpg, 17.jpg, 18.jpg, 19.jpg,
2.jpg, 21.jpg, 22.jpg, 23.jpg, 24.jpg, 25.jpg, 26.jpg, 27.jpg, 28.jpg, 29.jpg
因此保险起见, 最好对它重新排个序(当然不排序的话, 只要按原文件名去保存转换好的图像, 也是可以的, 但是为啥推荐做排序, 假设处理某一帧的时候出错了, 那么排序后能够快速定位到这一帧的索引, 方便后续DEBUG)
files = os.listdir(content_dir)
num_frames = len(files)
files.sort(key=lambda x: int(x[:-4]))
这里需要将x.jpg中的最后4个字符.jpg去掉, 只留下一个int型的x, 然后升序排序.
def generate():
style_img = load_img('./styles/denoised_starry.jpg')
content_dir = './contents/'
files = os.listdir(content_dir)
num_frames = len(files)
files.sort(key=lambda x: int(x[:-4]))
num_epochs = 5
begin = args.begin-1
end = args.end
cnt = begin
for cur in range(begin, end):
file = files[cur]
log_string('-' * 60)
cnt += 1
save_name = './generates/{}.jpg'.format(cnt)
log_string('cnt: {}/{}, current image: {}'.format(cnt, num_frames, file))
content_img = load_img('{}/{}'.format(content_dir, file))
img_styled = train(content_img, style_img, num_epochs, num_frames - cnt) * 255
log_string('saving to {}'.format(save_name))
img_styled = cv.cvtColor(img_styled, cv.COLOR_RGB2BGR)
cv.imwrite(save_name, img_styled)
log_string('-' * 60)
log_string('\\r')
if __name__ == '__main__':
split_frames()
问题
由于整个模型需要的时间非常长, 过程中万一出现异常, 那就直接终止了, 比如你睡前开始运行…结果半小时后它碰到异常停了…那一晚上就浪费了…为了避免重新跑带来的麻烦, 需要提高点鲁棒性.
最容易发生的问题其实是OOM爆显存. 一旦发生了任何问题, 整个程序就停止了, 这时候就体现出来日志的重要性了, 查看日志可以快速地定位到模型处理到哪一帧, 然后后期再补上这一帧.
其实只需要给generate.py用try catch捕捉异常即可. 等模型跑完之后, 查看日志, 看看有多少个ERROR, 重新跑一遍这些帧即可, 不用把所有的6900帧都跑一遍.
新建main.py
import os
from logger import log_string
if __name__ == '__main__':
steps = 50 # 每次处理50帧
for begin in range(135, 6835, steps):
cmd = 'python generate.py -begin {} -end {}'.format(begin, begin+steps-1)
try:
log_string('run {}'.format(cmd))
os.system(cmd)
except:
log_string('cmd {} failed'.format(cmd), 'ERROR')
最后运行的时候只需要运行main.py即可. 由于加了try catch, 碰到异常也不会终止进程. 发生异常的时候会在日志里面记录, 最后只要等它跑完, 去日志里面定位ERROR, 然后把这些缺失的帧补上就行.