FFmpeg中基于深度学习模型的图像处理filter：dnn_processing介绍（2）

dnn_processing是FFmpeg中的一个视频filter，可以支持所有基于深度学习模型的图像处理算法，即输入和输出都是AVFrame，而处理过程使用的是深度学习模型。为什么要开发这样一个filter，因为作为FFmpeg DNN模块的maintainer，希望可以有更多的人来使用这个模块，提出意见和批评，乃至发出patch来改善这个模块，dnn_processing就是一个很好的使用者入手功能。至于支持视频分析功能的filter，则还在计划中，不过我不准备太着急开始编码，一来还要考虑如何支持异步建立流水线，如何启用batch size，从而最大化的用好系统的并行计算能力（感谢和pengfei的讨论）；二来则是希望再等待DNN模块一段时间，以更好的夯实基础，完善DNN接口，完善后端支持。

在上一篇点击进入介绍了如何使用dnn_processing来完成针对灰度图的sobel算子的调用，其输入输出的格式是grayf32。本篇将介绍dnn_processing如何完成sr（超分辨率）和derain（去除雨点）这两个filter的功能，从而演示对yuv和rgb格式的支持。

• 重新编译FFmpeg

FFmpeg对深度学习模型的支持有两个后端，一是调用TensorFlow C动态库来完成，二是调用内部C代码来实现。前者被称为tensorflow后端（将简写为tf后端），后者被称为native后端。要启用tf后端，必须在编译FFmpeg的机器上有tensorflow c开发库（包括.h文件和.so文件），而且在configure的时候需要传入相应的参数。而任何情况下的编译，native后端总是被自动启用。这两个后端本文都会涉及，所以，我们必须重新编译FFmpeg。更多细节，请见上一篇介绍。

• 支持sr功能

首先，我们要准备sr需要的深度学习模型，sr filter支持两个模型，分别是srcnn和espcn，具体区别将在后续使用时候介绍。tf模型文件可以通过https://github.com/HighVoltageRocknRoll/sr或https://github.com/XueweiMeng/sr中的脚本生成，得到tf模型后再用FFmpeg目录树中的脚本生成native模型文件，具体如下。

# 演示目的，我们放到/tmp目录下。
# 系统重启后，/tmp目录下的所有内容都会消失。
$ cd /tmp/

$ git clone https://github.com/HighVoltageRocknRoll/sr.git
$ cd sr/

# 下面这条命令会在当前目录下生成tf模型，文件名是srcnn.pb
$ python generate_header_and_model.py --model=srcnn --ckpt_path=checkpoints/srcnn/
$ ls srcnn.pb -s
36 srcnn.pb

# 下面这条命令会在当前目录下生成tf模型，文件名是espcn.pb
$ python generate_header_and_model.py --model=espcn --ckpt_path=checkpoints/espcn 
$ ls espcn.pb -s
88 espcn.pb

# 然后去ffmpeg目录树下
$ cd /path_to_your_ffmpeg_source_tree/
$ cd tools/python/

# 下面命令生成native模型文件srcnn.model
$ python convert.py /tmp/sr/srcnn.pb
$ ls -s srcnn.model 
36 srcnn.model

# 下面命令生成native模型文件espcn.model
$ python convert.py /tmp/sr/espcn.pb
$ ls -s espcn.model 
84 espcn.model

srcnn模型是第一个基于深度学习的超分辨率实现，其原理是首先将一个低分辨率图片放大，然后，针对放大后的高分辨率图片的Y通道基于深度学习做处理，UV通道保持不变，处理前后的分辨率保持不变，最后得到高质量的高分辨率图片。具体命令如下所示。

$ cd /path_to_ffmpeg_build_path/

# 下面命令的输入是一张480p的图片，分辨率为720x480。
# 命令行中使用了三个filter，分别是format，scale和dnn_processing，
# format用来确保到达dnn_processing的格式是YUV格式，
# scale则是将低分辨率图片的长和宽分别放大2倍，
# dnn_processing中的model参数指出模型文件，dnn_backend指出使用tf后端，
# 而input和output则是模型的输入输出变量名字，在用python脚本搭建模型的时候指定。
# 最后，命令行的输出图片文件名是srcnn.jpg，其大小是1440*960。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
scale=w=iw*2:h=ih*2, \
dnn_processing=dnn_backend=tensorflow:model=/tmp/sr/srcnn.pb:input=x:output=y \
srcnn.jpg

# 修改dnn_processing中的dnn_backend和model参数，
# 就可以切换到native后端执行，其他都保持不变。
# 记得一定要用刚才生成的srcnn.model文件。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
scale=w=iw*2:h=ih*2, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/srcnn.model:input=x:output=y \
srcnn.jpg 

espcn模型则比较简洁，不需要提前进行放大，其模型的输入是低分辨率图片，输出就是高分辨率图片，也是只针对Y通道的处理，UV通道调用swscale模块（参数SWS_BICUBIC）进行放大。

$ cd /path_to_ffmpeg_build_path/

# 这里就不需要用到scale了，其他参数解释见前面介绍。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
dnn_processing=dnn_backend=tensorflow:model=/tmp/sr/espcn.pb:input=x:output=y \
espcn.jpg

# 下面命令行用native后端执行。
$ ./ffmpeg -i 480p.jpg -vf \
format=yuv420p, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/espcn.model:input=x:output=y \
espcn.jpg
#

• 支持derain功能

derain的深度学习模型的获取方法如下所示，最后得到 tf模型文件can.pb 和 native模型文件can.model。

$ cd /tmp/
$ git clone https://github.com/XueweiMeng/derain_filter
$ cd derain_filter/scripts/
$ ./export_model.sh
$ ls -s ../can.pb  
148 ../can.pb

$ cd /path_to_your_ffmpeg_source_tree/
$ cd tools/python/
$ python convert.py /tmp/derain_filter/can.pb 
$ ls -s can.model 
108 can.model

derain模型支持的格式是rgb24，所以，FFmpeg命令行如下所示。

# 输入rain.jpg是带雨点的图片，
# 用到了两个filter，format和dnn_processing，
# 具体参数已经在前面有介绍。
$ ./ffmpeg -i rain.jpg -vf \
format=rgb24, \
dnn_processing=dnn_backend=tensorflow:model=can.pb:input=x:output=y \
derain.jpg

# 下面命令使用native后端执行
$ ./ffmpeg -i rain.jpg -vf \
format=rgb24, \
dnn_processing=dnn_backend=native:model=/your_ffmpeg_source_tree/tools/python/can.model:input=x:output=y \
derain.jpg

以上介绍完毕。如果以后出现处理YUV全部通道的深度学习模型成为主流的情况，将会增加相应的支持。

以上内容是本人业余时间兴趣之作，限于水平，差错难免，仅代表个人观点，和本人任职公司无关。

本文首发于微信公众号：那遁去的一