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/srhttps://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全部通道的深度学习模型成为主流的情况,将会增加相应的支持。

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

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

你可能感兴趣的:(FFmpeg)