直播流AI应用开发工程技巧(一) 视频帧传输推理优化

一.前言

        最近在做直播流上面的AI应用,所以这个系列主要是记录一些在直播流项目中的工程经验。比如如何基于Flask快速进行视频帧的传输推理,如何用ffmpeg-python里面快速搭建起视音频直播流推理环境,保证音视频时序对齐,如何尽可能多线程的方式加速应用等等。

        这篇我们一起讨论下第一个主题,如何使用python flask来做一个直播流上面的推理引擎。应用场景比如说视频的逐帧审核、逐帧打标等等,当然对于抽帧的逻辑也是同样适用的。在直播应用中使用AI推理,我个人经验是整体的编解码和AI推理的性能都很重要。尤其是在高分辨率上比如4K、8K中,编解码本身的消耗可能高于AI模型推理的消耗。举个简单的例子,比如接受直播流上标注人脸的bbox之后再推流,我们以1080P 50FPS这样的要求为基准,如果用现在比较成熟的检测网络如yolov5,批量推理无论是使用TensorRT还是使用Pytorch、TensorFlow等,检测模型每帧可以达到平均1~2ms的推理性能(yolov5s为例),fps远高于需要的50帧,所以一般而言瓶颈都在前后处理以及编解码上了。

        如果上述的场景在python环境体系下做,最快的方式当然是将AI的模块和直播的模块部署在一起,数据交换全部通过内存,这样即使是使用python这种“较慢”的语言,也能轻松达到实时效果。但如果AI模块和直播模块必须要分开(如分属与两个团队或公司对接),可能中间就要采用HTTP接口的方式进行对接。采用这种方式进行对接开发的话,后端的AI模块和编解码模块中间就涉及到较大的数据传输,因此需要采用合理的视频编码格式进行编解码和压缩,以增大效率。在优化得当的情况下,像检测、分割等场景做到1080P分辨率的实时推理也是可以实现的。

        下面我们来一起探讨下视频帧画面传输的性能优化,作者本身还在这个问题上不断尝试和测试,所以当下发的文章里面可能存在很多改进的点,希望大家包容。

二.原理讨论

2.1 关于批量推理

        如果做视频的推理,为了提高fps,可以批量的进行帧编码,对于后端的AI推理模块而言,一般选择比如16、32、64等批大小(2的n次方更利于发挥cuda并行计算的能力),在一定范围内增大批量推理的batchsize可以获得fps的收益。

2.2 关于视频帧的压缩和编码

        服务端与客户端交互可以使用json的形式来传输数据。对于帧数据的形式,在python体系下一般都是操作numpy数组对象。numpy数组可以直接转换为字节进行传输,但是numpy一般表示图像是三通道的形式,这样传输的字节数量太大,很难实际使用。比如1080P的视频,一帧画面的像素点是 1920 * 1080 * 3 = 6220800 个像素, 按照fp32的精度表示,一个fp32占4个字节。所以理论上一帧画面直接传输(6220800 * 4) / 1024 / 1024 =23MB的数据,实在是太大了。

        为了解决这种情况,必须对numpy的数据进行编码压缩。首先是注意用uint8的精度表达图像数据已经足够,uint8的一个像素点只占1个字节。然后可以采用jpeg压缩,压缩的数据既可以是RGB格式也可以是YUV格式,相比较而言YUV格式的数据经过JPEG编码后占用的空间更小。下面看一下我做的一组测试,其中加载的一张随便的1080P的视频帧。

>>> some_img = cv2.imread("/data/some_1080.jpg")
>>> some_img.shape
(1080, 1920, 3)
>>> len(some_img.tobytes())
6220800
>>> some_img.dtype
dtype('uint8')
>>> encoded_img = cv2.imencode(".jpg", some_img, params=[cv2.IMWRITE_JPEG_QUALITY, 100])[1]
>>> sys.getsizeof(some_img)
6220936
>>> sys.getsizeof(encoded_img)
782868
>>> yuv_img = cv2.cvtColor(some_img, cv2.COLOR_BGR2YUV)
>>> encoded_yuv_img = cv2.imencode(".jpg", yuv_img, params=[cv2.IMWRITE_JPEG_QUALITY, 100])[1]
>>> sys.getsizeof(encoded_yuv_img)
666916
>>> encoded_yuv_img = cv2.imencode(".jpg", yuv_img, params=[cv2.IMWRITE_JPEG_QUALITY, 80])[1]
>>> sys.getsizeof(encoded_yuv_img)
113153
>>> encoded_yuv_img = cv2.imencode(".jpg", yuv_img, params=[cv2.IMWRITE_JPEG_QUALITY, 90])[1]
>>> sys.getsizeof(encoded_yuv_img)
172366
>>> encoded_yuv_img = cv2.imencode(".jpg", yuv_img, params=[cv2.IMWRITE_JPEG_QUALITY, 50])[1]
>>> sys.getsizeof(encoded_yuv_img)
68609
  • 采用JPEG编码,可以直接降低一个数量级的空间占用。
  • 采用yuv格式比rgb格式,占用的字节数量要降低10%。
  • 在用opencv做JPEG encode的时候,可以通过调节压缩比去进一步的降低字节数。从100%调整到90%就有6倍的空间降低。

        为了看的更清晰,我们列个表格。(100%压缩率意味着无压缩)

不同格式和JPEG编码后的字节数
格式 JPEG压缩 压缩率 字节数
uint8 ndarray 6220936
RGB 100% 782868
YUV 100% 666916
RGB 90% 307097
RGB 80% 214171
RGB 70% 176421
RGB 60% 152427
RGB 50% 137314
YUV 50% 68609

        可以看出来,比如采用YUV格式,进行50%的JPEG编码,就可从原来的传输6MB缩小到传输68K数据,几乎缩小了100倍。因此在思路上,如果对于AI模型的推理效果影响不那么明显的情况下,我们可以测试不同压缩率对于AI服务的识别影响,这一系列操作节省网络IO开销。总结一下就是:

  1. 用uint8代替fp32的精度。
  2. 用yuv格式代替rgb格式进行表示。
  3. 用jpeg编码对图像进行编码和压缩。
  4. 在AI模型上测试,不同的压缩比例对于推理识别效果的影响。通过调整压缩的比例,来平衡模型的推理结果和传输的网络开销。

        一个简单的代码表示就是:

def encode_frame(frame):
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)
    frame = cv2.imencode(".jpg", frame, params=[cv2.IMWRITE_JPEG_QUALITY, 80])[1]
    res = frame.tobytes()
    res = base64.b64encode(res).decode()
    return res

        这个函数接受numpy数组,然后进行YUV格式转换和JPEG编码压缩。

        对应的解码函数为:

def decode_frame_json(data):
    data = base64.b64decode(data.encode())
    image = np.frombuffer(data, np.uint8)
    image = cv2.imdecode(image, cv2.IMREAD_COLOR)
    image = cv2.cvtColor(image, cv2.COLOR_YUV2BGR)
    return image

        对于opencv的imdecode解码函数,无论编码前是yuv格式还是rgb格式,都可以自动识别并且解码。但是因为AI模型常常是在RGB格式上图片训练的,所以不要忘记将YUV图像还原成RGB。

补充:测试针对YoloV5,压缩率是如何影响检测效果的?

一个简单的测试发现,JPEG压缩之后的数据再恢复成ndarray进行yolov5识别(yolov5s),压缩率几乎是对于检测效果无影响的。

第一张图是没有压缩的:

直播流AI应用开发工程技巧(一) 视频帧传输推理优化_第1张图片

 第二张图是压缩率80%的:

直播流AI应用开发工程技巧(一) 视频帧传输推理优化_第2张图片

 第三张是压缩率20%的:

直播流AI应用开发工程技巧(一) 视频帧传输推理优化_第3张图片

 第四张是压缩率为1%的:

直播流AI应用开发工程技巧(一) 视频帧传输推理优化_第4张图片

         可以看到在图像人脸信息比较充分的情况下,压缩率20以上几乎对于检测的结果,无论是位置还是置信度的影响都极小。因此我们可以有较大的信心在直播流上对图像进行较大比例的JPEG编码压缩,而只损害极小的检测性能。(目前简单测试,后续继续补充测试结果)

2.3 关于base64

        在发送数据之前,最好是经过base64编码,这样的好处是对于特殊字符进行处理,防止传输出现错误。base64编码之后的空间占用要较之前略高。

>>> sys.getsizeof(str(encoded_yuv_img.tobytes(), 'latin1'))
68578
>>> import base64
>>> res = base64.b64encode(encoded_yuv_img).decode()
>>> type(res)

>>> sys.getsizeof(res)
91389

        numpy转为字节之后,用UTF-8格式可能会遇到无法解码成字符串的形式,可以尝试使用其他比如latin1的解码格式。

2.4 关于Json格式(如何传输视频帧是最高效的)

        目前我的策略是通过json的格式收发视频帧的数据。这样的好处是json简单,都是字符串,且可以有现成的库可以直接转换。但是json传输视频帧数据可能因为base64编码的原因,造成实际传输空间的占用。对于1080P的视频帧画面而言,base64编码可能比原始增加30%左右的空间占用。

        或许更好的方式是直接传输字节数据本身,放弃使用json格式传输。或者用其他的编码方式,比如latin1进行。

        这里占坑,我在做过更多实验或者有了新进展之后,再回来补充。

  

你可能感兴趣的:(计算机视觉,人工智能,音视频,AI工程化,AI服务)