直播流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），压缩率几乎是对于检测效果无影响的。

第一张图是没有压缩的：

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

 第三张是压缩率20%的：

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

         可以看到在图像人脸信息比较充分的情况下，压缩率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进行。

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