当前 AI 算法蓬勃发展,但在开源的代码中,基本都是处理图片,原生支持处理视频的算法寥寥无几。究其原因,相比图片的处理,视频的处理不仅需要考虑封装格式的处理(如 MP4、HLS、MKV 等),还要考虑编码格式的处理(如 H264、H265、AV1、VP9 等),这是都是算法开发人员不得不面对的一个障碍。
FFmpeg 作为一个持续了 20 多年的开源项目,号称音视频处理的“瑞士军刀”。在 FFmpeg 中,有一个 AVFilter 模块,支持简单的音视频前处理、后处理,如图像调色、图像叠加等。近几年,随着 AI 技术的发展,FFmpeg 也支持集成了 libtensorflow 的能力,可以支持一些简单的音视频 AI 能力。但开发 FFmpeg 的 AVFilter 模块,仍有一定的门槛。
BabitMF(Babit Multimedia Framework,BMF),是字节跳动最近开源的一个通用的多媒体处理框架。在 BMF 中,AVFilter 对应都是 BMF 模块。从它的开源文档介绍中,看到 BMF 完全兼容 FFmpeg 的功能和标准,而且支持 Python 开发,这可以显著提升 AI 算法在视频处理上的集成效率,对 AI 算法开发人员是一个福音!
那么,BMF 模块真的是 AI 视频处理利器吗?体验一下就知道了。
BMF 安装
BMF 有四种安装方式,具体如下:
- pip 安装:在满足依赖的情况下,安装比较简单
- docker 镜像:无需关注依赖情况,直接拉取镜像即可体验,但 babitmf/bmf_runtime:latest超过 10G
- 预编译二进制文件:需要满足依赖
- 源码构建:需要关注依赖和编译选项,极客玩家必选
我有一台 centos 8 的云服务器,秉承尽量少折腾的原则,先尝试拉取 docker 镜像,但拉取 10G 的镜像实在太慢,遂放弃该安装方式。剩下的三种方法,都需要先处理下依赖,命令如下:
# 安装前置依赖
dnf -y upgrade libmodulemd
dnf -y install glibc-langpack-en epel-release epel-next-release
dnf makecache
dnf update -y
dnf config-manager --set-enabled powertools
dnf -y install make git pkgconfig cmake3 openssl-devel binutils-devel gcc gcc-c++ glog-devel
# 安装 FFmpeg
dnf install -y https://download1.rpmfusion.org/free/el/rpmfusion-free-release-8.noarch.rpm
dnf install -y https://download1.rpmfusion.org/nonfree/el/rpmfusion-nonfree-release-8.noarch.rpm
dnf install -y ffmpeg ffmpeg-devel
因为 dnf 搜索不到 python3.9 版本,因此采用源码安装:
cd /opt
wget https://www.python.org/ftp/python/3.9.13/Python-3.9.13.tgz
tar xvf Python-3.9.13.tgz
cd Python-3.9.13
sudo ./configure --enable-optimizations --enable-shared
sudo make altinstall
设置下环境变量:export LD_LIBRARY_PATH=/usr/local/lib/:$LD_LIBRARY_PATH
,执行 python3 -c 'import bmf'
成功,则表示安装环境已成功。
Python 模块开发
在官方文档有创建 Python 模块的示例,对照着改造了一个超分的算法模块,惊喜地发现并不需要太多的改动!可以在这个代码仓库查看相关的 BMF 模块和测试代码。
开发和管理 BMF Python 模块
BMF 的模块开发,需要关注两个函数:__init__
和process
。其中,__init__
用于初始化模块,process
里包装了对单帧视频或音频的处理逻辑。BMF 提供了模块管理工具 module_manager
,可以方便地安装、管理本地的模块。
接下来,我们使用官网提供的复制流的代码,快速熟悉 BMF 模块的开发和管理流程。
复制流的代码逻辑比较简单,在process
中,直接把输入的视频包直接输出即可,代码参考copy_module.py。接下来,使用module_manager
安装模块,执行命令和输出成功的日志如下:
module_manager install copy_module python copy_module:CopyModule $(pwd)/ v0.0.1
Installing the module:copy_module in "/usr/local/share/bmf_mods/Module_copy_module" success.
接下来,就是对这个模块进行测试,代码如下:
import bmf
import sys
input_file = sys.argv[1]
output_path = 'copy.mp4'
(
bmf.graph()
.decode({'input_path': input_file})['video']
.module('copy_module')
.encode(None, {"output_path": output_path})
.run()
)
代码还是非常直观的,构建graph
,将输入文件进行解码,取其中的视频流,使用我们新建的模块进行处理,最后进行编码输出。运行命令 python3 test_copy_module.py input.jpg
,我们可以看到如下的日志输出如下,可以看到载入了我们新建的copy_module
,将输入的图片处理后,生成了 MP4 文件。
[2023-12-31 11:09:12.655] [info] c_ffmpeg_decoder c++ /root/py3-env/lib/python3.9/site-packages/bmf/lib/libbuiltin_modules.so libbuiltin_modules.CFFDecoder
[2023-12-31 11:09:12.655] [info] Module info c_ffmpeg_decoder c++ libbuiltin_modules.CFFDecoder /root/py3-env/lib/python3.9/site-packages/bmf/lib/libbuiltin_modules.so
[2023-12-31 11:09:12.656] [info] Constructing c++ module
[2023-12-31 11:09:12.658] [error] node id:0 Could not find audio stream in input file 'input.jpg'
Input #0, image2, from 'input.jpg':
Duration: 00:00:00.04, start: 0.000000, bitrate: 497 kb/s
Stream #0:0: Video: mjpeg (Baseline), yuvj420p(pc, bt470bg/unknown/unknown), 128x128 [SAR 1:1 DAR 1:1], 25 tbr, 25 tbn, 25 tbc
[2023-12-31 11:09:12.658] [info] c++ module constructed
[2023-12-31 11:09:12.658] [info] copy_module python /usr/local/share/bmf_mods/Module_copy_module copy_module.CopyModule
[2023-12-31 11:09:12.658] [info] Module info copy_module python copy_module.CopyModule /usr/local/share/bmf_mods/Module_copy_module
[2023-12-31 11:09:12.658] [info] c_ffmpeg_encoder c++ /root/py3-env/lib/python3.9/site-packages/bmf/lib/libbuiltin_modules.so libbuiltin_modules.CFFEncoder
[2023-12-31 11:09:12.658] [info] Module info c_ffmpeg_encoder c++ libbuiltin_modules.CFFEncoder /root/py3-env/lib/python3.9/site-packages/bmf/lib/libbuiltin_modules.so
[2023-12-31 11:09:12.658] [info] Constructing c++ module
[2023-12-31 11:09:12.658] [info] c++ module constructed
[2023-12-31 11:09:12.659] [info] BMF Version: 0.0.9
[2023-12-31 11:09:12.659] [info] BMF Commit: e3c9730
[2023-12-31 11:09:12.659] [info] start init graph
[2023-12-31 11:09:12.659] [info] scheduler count2
debug queue size, node 0, queue size: 5
[2023-12-31 11:09:12.659] [info] node:c_ffmpeg_decoder 0 scheduler 0
debug queue size, node 1, queue size: 5
[2023-12-31 11:09:12.659] [info] node:copy_module 1 scheduler 0
debug queue size, node 2, queue size: 5
[2023-12-31 11:09:12.659] [info] node:c_ffmpeg_encoder 2 scheduler 1
[2023-12-31 11:09:12.660] [info] node id:0 decode flushing
[2023-12-31 11:09:12.660] [info] node id:0 Process node end
[2023-12-31 11:09:12.660] [info] node id:0 close node
[2023-12-31 11:09:12.660] [info] node 0 close report, closed count: 1
[2023-12-31 11:09:12.660] [info] node id:1 eof received
[2023-12-31 11:09:12.660] [info] node id:1 eof processed, remove node from scheduler
[2023-12-31 11:09:12.660] [info] node id:1 process eof, add node to scheduler
[2023-12-31 11:09:12.660] [info] node id:1 Process node end
[2023-12-31 11:09:12.660] [info] node id:1 close node
[2023-12-31 11:09:12.660] [info] node 1 close report, closed count: 2
[2023-12-31 11:09:12.660] [info] node id:2 eof received
[2023-12-31 11:09:12.660] [info] node id:2 eof processed, remove node from scheduler
[libx264 @ 0x7f3308002400] using SAR=1/1
[libx264 @ 0x7f3308002400] using cpu capabilities: MMX2 SSE2Fast SSSE3 SSE4.2 AVX FMA3 BMI2 AVX2 AVX512
[libx264 @ 0x7f3308002400] profile High, level 1.1, 4:2:0, 8-bit
[libx264 @ 0x7f3308002400] 264 - core 157 r2980 34c06d1 - H.264/MPEG-4 AVC codec - Copyleft 2003-2019 - http://www.videolan.org/x264.html - options: cabac=1 ref=3 deblock=1:0:0 analyse=0x3:0x113 me=hex subme=7 psy=1 psy_rd=1.00:0.00 mixed_ref=1 me_range=16 chroma_me=1 trellis=1 8x8dct=1 cqm=0 deadzone=21,11 fast_pskip=1 chroma_qp_offset=-2 threads=3 lookahead_threads=1 sliced_threads=0 nr=0 decimate=1 interlaced=0 bluray_compat=0 constrained_intra=0 bframes=3 b_pyramid=2 b_adapt=1 b_bias=0 direct=1 weightb=1 open_gop=0 weightp=2 keyint=250 keyint_min=25 scenecut=40 intra_refresh=0 rc_lookahead=40 rc=crf mbtree=1 crf=23.0 qcomp=0.60 qpmin=0 qpmax=69 qpstep=4 ip_ratio=1.40 aq=1:1.00
Output #0, mp4, to 'copy.mp4':
Metadata:
encoder : Lavf58.29.100
Stream #0:0: Video: h264 (avc1 / 0x31637661), yuv420p, 128x128 [SAR 1:1 DAR 1:1], q=2-31, 25 fps, 25 tbr, 12800 tbn, 25 tbc
[swscaler @ 0x7f3308125e40] deprecated pixel format used, make sure you did set range correctly
[2023-12-31 11:09:12.663] [info] node id:2 process eof, add node to scheduler
[2023-12-31 11:09:12.664] [info] node id:2 Process node end
[libx264 @ 0x7f3308002400] frame I:1 Avg QP:29.45 size: 795
[libx264 @ 0x7f3308002400] mb I I16..4: 15.6% 64.1% 20.3%
[libx264 @ 0x7f3308002400] 8x8 transform intra:64.1%
[libx264 @ 0x7f3308002400] coded y,uvDC,uvAC intra: 64.8% 81.2% 25.0%
[libx264 @ 0x7f3308002400] i16 v,h,dc,p: 40% 0% 0% 60%
[libx264 @ 0x7f3308002400] i8 v,h,dc,ddl,ddr,vr,hd,vl,hu: 28% 7% 19% 6% 4% 12% 5% 13% 6%
[libx264 @ 0x7f3308002400] i4 v,h,dc,ddl,ddr,vr,hd,vl,hu: 30% 7% 26% 5% 4% 7% 4% 9% 8%
[libx264 @ 0x7f3308002400] i8c dc,h,v,p: 47% 12% 33% 8%
[libx264 @ 0x7f3308002400] kb/s:159.00
[2023-12-31 11:09:12.665] [info] node id:2 close node
[2023-12-31 11:09:12.665] [info] node 2 close report, closed count: 3
[2023-12-31 11:09:12.665] [info] schedule queue 0 start to join thread
[2023-12-31 11:09:12.665] [info] schedule queue 0 thread quit
[2023-12-31 11:09:12.665] [info] schedule queue 0 closed
[2023-12-31 11:09:12.665] [info] schedule queue 1 start to join thread
[2023-12-31 11:09:12.665] [info] schedule queue 1 thread quit
[2023-12-31 11:09:12.665] [info] schedule queue 1 closed
[2023-12-31 11:09:12.665] [info] all scheduling threads were joint
{
"input_streams": [],
"output_streams": [],
"nodes": [ { "module_info": { "name": "c_ffmpeg_decoder", "type": "", "path": "", "entry": "" }, "meta_info": { "premodule_id": -1, "callback_binding": []
},
"option": {
"input_path": "input.jpg"
},
"input_streams": [],
"output_streams": [ { "identifier": "video:c_ffmpeg_decoder_0_1", "stream_alias": "" } ],
"input_manager": "immediate",
"scheduler": 0,
"alias": "",
"id": 0
},
{
"module_info": {
"name": "copy_module",
"type": "",
"path": "",
"entry": ""
},
"meta_info": {
"premodule_id": -1,
"callback_binding": []
},
"option": {},
"input_streams": [ { "identifier": "c_ffmpeg_decoder_0_1", "stream_alias": "" } ],
"output_streams": [ { "identifier": "copy_module_1_0", "stream_alias": "" } ],
"input_manager": "immediate",
"scheduler": 0,
"alias": "",
"id": 1
},
{
"module_info": {
"name": "c_ffmpeg_encoder",
"type": "",
"path": "",
"entry": ""
},
"meta_info": {
"premodule_id": -1,
"callback_binding": []
},
"option": {
"output_path": "copy.mp4"
},
"input_streams": [ { "identifier": "copy_module_1_0", "stream_alias": "" } ],
"output_streams": [],
"input_manager": "immediate",
"scheduler": 1,
"alias": "",
"id": 2
}
],
"option": {},
"mode": "Normal"
}
[2023-12-31 11:09:12.667] [info] schedule queue 0 start to join thread
[2023-12-31 11:09:12.667] [info] schedule queue 0 closed
[2023-12-31 11:09:12.667] [info] schedule queue 1 start to join thread
[2023-12-31 11:09:12.667] [info] schedule queue 1 closed
[2023-12-31 11:09:12.667] [info] all scheduling threads were joint
[2023-12-31 11:09:12.667] [info] node id:0 video frame decoded:1
[2023-12-31 11:09:12.667] [info] node id:0 audio frame decoded:0, sample decoded:0
解决人脸超分算法的依赖
体验完 BMF 模块的使用方式,接下来就是解决算法的依赖问题。因为是 CPU 环境,所以先从 Github 上找一个可以在 CPU 上运行的代码。简单搜索后,决定使用这个人脸超分的代码 ewrfcas/Face-Super-Resolution,首先需要解决依赖问题,让代码在本地可以运行:
- 创建 Python 虚拟环境:
python3.9 -m venv ~/py3_env
- 激活虚拟环境:
source ~/py_env/bin/activate
- 安装 BMF:
pip3 install BabitMF
,执行python3 -c 'import bmf'
确认可运行 - 安装人脸超分代码的依赖:
pip3 install opencv-python scikit-image dlib torch torchvision
- 按照人脸超分代码仓库的
README
,下载依赖的模型,并执行python3 test.py
,确认可执行成功解决了算法依赖问题,就可以开始 BMF Python 模块的改造了。
改造人脸超分模块
我们可以在上面复制流模块的基础上,对算法模块进行改造。具体改造点包括:
init
中进行超分模型的初始化,这样在后续的处理中就可以直接使用了process
中将输入视频流的帧解码并转换成rgb24
的色彩空间,这样可以直接输出numpy
的数组,就可以直接使用原来的超分函数,最后将超分的结果重新编码成视频帧
class FaceSR(Module):
def __init__(self, node, option=None):
self.sr_model = SRGANModel(FaceSROpt(), is_train=False)
self.sr_model.load()
def process(self, task):
input_packets = task.get_inputs()[0]
output_packets = task.get_outputs()[0]
while not input_packets.empty():
pkt = input_packets.get()
if pkt.timestamp == Timestamp.EOF:
Log.log_node(LogLevel.DEBUG, task.get_node(), "Receive EOF")
output_packets.put(Packet.generate_eof_packet())
task.timestamp = Timestamp.DONE
return ProcessResult.OK
if pkt.is_(VideoFrame) and pkt.timestamp != Timestamp.UNSET:
vf = pkt.get(VideoFrame)
frame = ffmpeg.reformat(vf, "rgb24").frame().plane(0).numpy()
sr_frame = self.sr_forward(frame)
rgb = mp.PixelInfo(mp.kPF_RGB24)
video_frame = VideoFrame(mp.Frame(mp.from_numpy(sr_frame), rgb))
video_frame.pts = vf.pts
video_frame.time_base = vf.time_base
out_pkt = Packet(video_frame)
out_pkt.timestamp = video_frame.pts
output_packets.put(out_pkt)
return ProcessResult.OK
def sr_forward(self, img, padding=0.5, moving=0.1):
img_aligned, M = dlib_detect_face(img, padding=padding, image_size=(128, 128), moving=moving)
input_img = torch.unsqueeze(_transform(Image.fromarray(img_aligned)), 0)
self.sr_model.var_L = input_img.to(self.sr_model.device)
self.sr_model.test()
output_img = self.sr_model.fake_H.squeeze(0).cpu().numpy()
output_img = np.clip((np.transpose(output_img, (1, 2, 0)) / 2.0 + 0.5) * 255.0, 0, 255).astype(np.uint8)
rec_img = face_recover(output_img, M * 4, img)
return rec_img
代码改造好后,就可以使用 module_manager
进行本地发布:
~: module_manager install face_sr_module python face_sr_module:FaceSR $(pwd)/ v0.0.1
Installing the module:face_sr_module in "/usr/local/share/bmf_mods/Module_face_sr_module" success.
测试就更简单了,可以使用上面测试复制流的代码,把其中 copy_module
改成 face_sr_module
就可以了。
最后,我们运行测试代码 python3 test_copy_module.py input.jpg
,看下执行效果。
上面左图是输入图片,右图是人脸超分处理后的图片。可以看到,超分效果并不明显,这个后续再排查,不影响本次的 BMF 开发体验。
不知道各位读者是否注意到,测试代码都是使用图片作为输入,这是因为图片也是多媒体格式的一种,那么算法的开发验证,就不需要再把视频转成图片序列了。使用 BMF 开发,就可以做到同时处理图片和视频!
总结与建议
通过一个 Python 人脸超分模块的改造开发,验证了 BMF 多媒体处理框架能让 AI 算法在视频处理上的集成难度下降、集成效率提升。BMF 在字节内部有非常多的应用,支持专业地音视频处理,算法开发人员不再需要担心音视频的封装、编解码、音画同步等复杂情况。只要适配了 BMF 的模块开发要求,就可以做到一次开发,直接支持图片和视频!
当然,体验过程中,也发现一些可以改进的地方,比如:
- 可以提供轻量级的 Docker 镜像,或对现有镜像进行精简,便于下载体验
- 模块开发的示例可以添加一些具体的例子和详细的解释,更加便于理解和上手
本次主要体验了 BMF Python 模块的开发,相信 BMF 内部还有很多值得探索的特性,各位读者如果有兴趣,欢迎留言,一起交流学习。
infoq链接 https://xie.infoq.cn/article/07d52014af4c5bc0bbfbb9aad