【代码札记】使用ONNX在JVM上调用YoloV8识别鸟类

深度学习、神经网络、大语言模型，这些都是好东西，可唯独Python不是个好东西。催生于整理录像的需求，本文介绍如何在JVM上使用ONNX在GPU上运行YoloV8图像识别。

故事板

最近我开始在二楼的平台上喂鸟了。当然了，以我的心态，活的面包虫肯定是接受不了了，所以只能喂些谷物、坚果和烤干的面包虫，吸引来的鸟类也无外乎麻雀、斑鸠、喜鹊。偶尔能看到灰喜鹊和乌鸦，但前者经常被喜鹊打跑，后者来的较晚，坚果早就被喜鹊藏起来了。有一次周末在后院看到了一群大约四五只戴胜，还没来得及拍下照片，人家就跑了，后来再也没见过了（也可能是我起床就十点十一点了，早就过了鸟类早上活动的时间段了）。

最初为了观察到底有没有鸟来吃，也是为了给退休的二手手机找个活儿干，加上之前正好有买三脚架送的手机架，还有一个特别不好用的绿联三脚架，这不正好闭环了嘛：把手机架在旁边录像，看看都有哪些鸟来吃。录像用的是IP webcam pro，这个App可以在画面内打上时间戳水印，也可以调整录像质量（码率），灵活性比较高。虽然这个App本身也支持运动检测，但由于是户外，它并不能分辨这种运动是鸟来了，还是大风挂的树枝动。所以我就没有用运动检测，从头到尾录制就是了。

这里没提到的一点就是夏天的温度。我的手机是黑色的，人在北京，严格来说算不上太阳直射，但夏天的太阳光威力依旧很大，大约不到半个小时手机就会过热，然后强制结束前台程序。所以一般我都是等下午两点左右开始有树荫了，才把手机放到树荫里。另一个问题是续航，就算是关闭屏幕全靠后台运行，单凭手机的电池坚持不了多久。我一般直接在最开始就挂上充电宝，充电宝没点后再消耗手机的电池。我之前也尝试过先消耗手机的电池，等手机快没电了再挂上充电宝。这样的一个问题就是充电发热，即便手机在阴凉里，如果遇到无风的天气，这种操作也很容易让手机过热。

平均算下来，我下午两点录到晚上八点，每天会产生6个小时的录像。我当然当然可以开倍速找出哪里有鸟，然后把对应的片段剪辑出来，但我这不是懒嘛。于是便产生了一个需求：如果能写一个程序帮我自动做这件事呢？

野路子™解决方案

为什么说用户的话不能全信？如果你按照我的需求走，那可就麻烦大了：首先得想办法从一个视频文件中找出鸟类出现的镜头，找出片段的起点和终点，然后得想办法把这段给弄出来，这里面就涉及到了解码和编码，还得考虑画质会不会有损失，诸如此类的一系列棘手的问题。但如果你追究这个需求的问题，其实归根结底我是想从录像中删掉没有鸟类出现的镜头。

这样一来，问题就好办多了：先把长视频切分成时常较短的视频，然后分别对每个视频进行识别，删掉没有鸟类出现的片段，最后再把剩下的片段拼接起来。由于app录制出来的是mp4格式，并且分辨率什么的都是统一的，视频文件本身也具有时间戳，所以这种办法没什么不妥。最关键的是，这种方法完全不需要重新编码，其中切片和拼接的操作可以直接调用ffmpeg的命令行，将流设置为copy，也完全不会损失质量，而且还快。

但问题来了：图像识别免不了要调用神经网络。像是Yolo这些模型，全都是Python上的，怎么在jvm上跑起来呢？一开始我看到了opencv里面有个dnn的模块，可以用来加载模型做分类。但是在我尝试的时候发现，我将yoloV8导出成为onnx模型，加载到opencv里，它给不出bounding box的坐标，这不出大事了吗。虽然javacpp这个组给了方便好用的javacv和opencv的原始绑定，可是它用不了啊。

ONNX

再后来我贼心不死，深入调查了一下ONNX这个东西，发现它这个框架不错。这个框架做的事情类似于提供了一个指令集或者VM的作用，它给出了一些标准化的运算函数，只要其他各种框架能够将自己的模型使用这一套标准化的函数表达出来，那么无论你的模型原本是用什么框架训练的，都可以在其他地方统一进行运算。而正好，ONNX官方也提供了Java的包，让JVM也可以调用ORT（ONNX Runtime）来进行计算，同时支持GPU加速。

我选用的模型是YoloV8，这个Yolo系列分了好多，我觉得V8算是表现不错，而且支持文档比较丰富，它的命令行可以按照你给定的分辨率和batch大小直接产生对应的onnx模型。我这里选用的是1920x1080，至于batch大小，这个要看你GPU能吃多少。我的笔记本内置的3070使用yolov8x模型只能给到3，再大的话就会交换到系统内存，导致在传输数据上浪费时间。后来我去租GPU，RTX4090的话同样使用yolov8x，能给到10到12，不过超过10的话瓶颈就不在于显存了，而在于计算上。由于onnx使用的是通用的运算函数，所以它的计算效率并不像pytorch那些模型那么高效，不过jvm上有啥就用啥，能用就不错了，还要什么自行车啊。

准备好了模型，下一步就是调用了。

YoloV8

一开始我以为Yolo是个系列，类似llama1和llama2那种。后来发现不是，yolo系列的每个模型都是不同的人或公司搞出来的。并且根据模型的特性和功能不通，他们的输出也不尽相同。这里我简单介绍一下V8的输入输出吧。我的分辨率是1920x1080，由于高度必须得是32的整数，所以命令行自动帮我扩展到了1920x1088，batch大小是10，所以我的输入shape就是：[10, 3, 1088, 1920]。由于ONNX的向量并不提供多维访问，所以我们需要手动将这个四位坐标转换成数组索引。好在ONNX用的顺序和我们平时用的差不多，记输入向量的形状为[B, C, H, W]，那么我们要访问的坐标[b, c, y, x]对应的一维数组索引就是b * C * H * W + c * H * W + y * W+ x，这部分没啥技术含量，就是算索引，只要知道这个多维数组是怎么在内存里排列的就好了。

由于我对Yolo系列并不熟悉，以下内容也是现学现卖。如果有说错的地方还请大家在评论区指正。

关于输出，形状是[B, 4 + L, A]，这里的B还是表示batch大小，L表示标签数目，这里我用的coco数据集，所以标签数目是80，这一维度的输出有84个。最后的A表示的是锚点（anchors），这个根据你的输入分辨率不同，数目也不同。虽然我也不太懂这个锚点是什么，但根据我从网络上搜索到的信息，这个锚点就是将你的输入画面分割成好多个网格（grid），你可以理解为把一个画面按照20像素分割成网格，每一个网格就是一个锚点，然后再将画面按照50像素分成网格，每个网格又是一个锚点，最后再将画面按照100像素分割，将这些大小不同的网格对应的锚点加起来，就是输出的所有锚点。这样做的目的在于Yolo模型可以根据这些大小网格分割出来的画面来识别大小不同、位置不同的物体。我们在解析数据时，也需要遍历每一个锚点，

至于标签那个大小为什么还要加4，这是因为网络除了给出每个标签的置信度，还会给出这个物理对应的bounding box。对于给定的批次索引b和锚点索引a，data[b, 0, a]给出的是碰撞盒的中心x坐标，类似地，data[b, 1, a]给出的是中心y坐标，而data[b, 2, a]给出的是盒子的宽度，最后data[b, 3, 0]是盒子的高度。从data[b, 4, 0]开始才是第一个标签的置信度。

所以综上所述，处理输出的伪代码如下：

输入：批次索引`b`，输出形状`shape`，输出数据`data`
val (_, labels, anchors) = shape // 模式匹配，提取shapes和anchor的个数，这里的labels包含了4个额外的数据
val result = mutableList<Detection>() // 使用List来存储识别结果

for (a in 0 until anchors) { // 遍历每个anchor
    val (x,y,w,h) = FloatArray(4) { data[b, it, a] } // 类似kotlin的写法，通过模式匹配将label的0到3提取出来并赋值给yxwh四个变量
    var maxScore = data[b, 4, a] // 将第一个label作为最小值开始比较
    var maxIndex = 0
    for (l in 5 until labels) { // 从第二个开始比
        val s = data[b, l, a]
        if (s > maxScore) {
            maxScore = s
            maxIndex = l - 4
        }
    } // 这里一般认为一个锚点只会有一个结果，因此我们要选取一个置信度最高的标签
    // 根据最大值判断
    if (canAccept(maxIndex, maxScore)) { // 如果当前标签的置信度可以被接受，那就添加到结果里
        result.add(Detection(x, y, w, h, indexToLabelString(maxIndex), maxScore))
    }
}

输出：result

这里有两点要说明：首先是那个canAccept。在最早的版本中就是一个简单的对比，后来我发现不同的标签给出的置信度不一样，所以用户可能会想针对不同的标签设置不同的阈值，这个函数就是来干这个的。其次就是这里的result并不能直接使用。因为yoloV8的原始输出只是告诉你每个锚点上最有可能的物体是什么。但两个锚点之间可能会有重叠，所以我们还需要额外的处理。

NMS

NMS全称Non-Maximum Suppression，非极大值抑制算法。这个算法主要用来消除多余的/重叠的检测框。这个算法的伪代码如下：

输入：一堆Bounding Box（`input`），阈值`threshold`
输出：一小堆BoundingBox

val candidate = LinkedList(input.sortedByDescending { it.confidence }) // 首先按照置信度降序排列，置信度高的在前面
val result = mutableListOf<Detection>()
if (candidate.isEmpty()) return result // 确保输入不为空
result.add(candidate.removeFirst()) // 把第一个候选者加入结果列表

while (candidate.isNotEmpty()) {
    val c = candidate.removeFirst() // 取出一个候选者
    val maxIoU = result.maxOf { it.calculateIoU(c) } // 计算和已知结果最大的IoU
    if (maxIoU <= threshold) result.add(c) // 如果没超过阈值就接受这个候选
}

return result

这里的IoU是Intersection over Union，所谓Intersection就是两个bounding box交叉的面积，而Union就是两个bounding box叠起来的总面积，这个总面积不包含重复的。我们可以按照如下方法计算：

fun Detection.calculateIoU(other: Detection): Double {
    val ourArea = width.toDouble() * height
    val otherArea = other.width.toDouble() * height
    val intersection = intersectionArea(other)
    val unionArea = ourArea + otherArea - intersection
    return intersection / unionArea
}

fun Detection.intersectionArea(other: Detection): Double {
    val ourX1 = centerX - width / 2.0
    val ourY1 = centerY - height / 2.0
    val ourX2 = ourX1 + width
    val ourY2 = ourY1 + height

    val otherX1 = other.centerX - other.width / 2.0
    val otherY1 = other.centerY - other.height / 2.0
    val otherX2 = otherX1 + other.width
    val otherY2 = otherY1 + other.height

    // use the leftest x2 minus rightest x1, clamp to 0
    val width = (min(ourX2, otherX2) - max(ourX1, otherX1)).coerceAtLeast(0.0)
    // use the lowest y2 minus the highest y1, clamp to 0
    val height = (min(ourY2, otherY2) - max(ourY1, otherY1))
    return width * height
}

这个也没什么难度。经过NMS处理过的结果大体可以保证不会有好多个检测框堆在一起的情况。

视频解码和结果处理

核心说完了，下面简单说说怎么把视频里的画面输入给ONNX模型。虽然说之前那个方便的opencv绑定没用上，可是javacv还是用得上的。其中的org.bytedeco.javacv.FFmpegFrameGrabber就可以让我们使用FFmpeg从视频文件中抓取Frame，配合org.bytedeco.javacv.Java2DFrameConverter还可以把抓出来的Frame转化成Java的BufferedImage，这样一来我们后续就可以通过访问raster来直接抓取像素信息了。

为了方便考虑，我在读入视频的时候选择了全部读入加抽帧。因为本身设计上来讲，要识别的片段都是1分钟以内的片段，而视频中相邻的帧变化不大，所以可以采用每N帧抽出1帧的方式进行识别。减少了内存占用，同时还能加快检测速度，毕竟检测树叶晃动属于无效计算，能省就省。

此外，由于加载视频、准备向量、推理，这些使用的完全是不同的部件，所以我用kotlin协程来协调他们。通过调整Channel的大小，可以决定要事先加载多少视频到内存里。在进行推理时，根据batch大小讲输入数据分组，然后使用async来派发将图像转换成向量的任务，这样一个线程在等待GPU运算的同时，其他线程可以先将后面的数据集准备好，当前批次算好了，直接开始计算下一批次。计算完成后的IO操作也是一样，快速解析出推理结果之后，根据结果决定这个视频片段的去留，而实际的IO操作则派发到其他协程上，主线程直接开始下一个视频片段的处理。

经过这样的优化，在RTX4090上，使用batch大小为10，分辨率为1920x1088（多出来的8行直接留黑），能够达到15FPS的处理效果（不考虑抽帧）。如果考虑抽帧的话，比如我每6帧抽1帧，那四舍五入就是90FPS的处理速度。实际上考虑到视频的观赏性，如果鸟只在镜头前一闪而过，则完全没有观赏性。假设使用30FPS的帧率录制影片，我们完全可以使用8帧、10帧或者15帧来抽取。对于难以识别的场景（例如树叶阴影），我们可以使用8帧，多一些画面就多一些成功识别的机会。对于比较简单的场景，我们可以使用12或15帧来提高检测速度。

同时为了方便检查，在输出时还会把被检测的物体画出来，单独将那一帧保存成图片。这样我们无需查看一堆视频碎片，我们只需要过一遍检测出来的图片，确保每一个视频都正确识别即可。

FFmpeg切片与合并

核心和周边都说完了，最后说说怎么做前期准备吧。我这里使用命令行来调用ffmpeg，但其实方法都差不多。

关于切割视频，我们需要使用segment demuxer：

ffmpeg -i input.mp4 -f segment \
             -segment_time 30 \
             -vcodec copy -acodec copy \
             -reset_timestamps 1 -map 0 \
             output_%d.mp4

别问我为什么要这么写，我也不知道，我从stackvoerflow上看来的。总之这个命令就是把input.mp4按照30秒的长度切分成小块。处理的时候音频流和视频流全都采用copy，不重新编码，输出就是output_%d.mp4，第一个视频就是output_0.mp4，以此类推。

合并的话我们要用concat muxer，调用命令行之前，首先我们要准备一个文件来列出我们需要拼接的文件：

file '文件1的绝对路径.mp4'
file '文件2的绝对路径.mp4'
file '文件3的绝对路径.mp4'
file '文件4的绝对路径.mp4'

然后调用命令行：

ffmpeg -f concat -safe 0 \
             -i 刚才那个清单文件的路径 \
             -c copy \
             output.mp4

这样一来ffmpeg就会将清单中的文件合并成一个文件，并且输出成output.mp4。同样，别问为什么，我也是StackOverflow上看来的。此外我还发现，这个命令行的选项顺序至关重要。如果顺序错了，命令行可能就直接失败了。

租用GPU

由于我的笔记本散热比较糟糕，夏天我这屋子里也没有空调。所以后来我将目光转向了云GPU。使用云GPU的一个问题就是要将待处理的视频文件传输到云端。得益于注诸如S3、B2这些对象存储，我们可以事先将切片上传到桶里，然后租用服务器，直接在服务器内将桶中的数据下载到服务器上开始推理，然后将推理结果存回桶里。

AWS

最开始我的候选是AWS，但是默认情况下AWS不给你GPU实例的权限，你得找客服申请。墨迹了一周，申请下来了。一个小时0.526美元的g4dn.xlarge实例，跑起来还没我自己的电脑快，还那么贵，呸，恶心，我都关着灯。

vast.ai

再后来我想起来之前用过的vast.ai，这里也顺便挂一个我的referral链接：https://cloud.vast.ai/?ref_id=61326

由于我没有提现的需求（vast.ai允许将奖励金的75%提现到paypal），我只是想少花点钱跑GPU负载，所以如果大家想要注册vast.ai的话，不妨考虑使用我的链接，这样我就可以获得你消费金额的3%作为回馈。比如，你消费100美元，我就可以获得3美元的免费额度。

关于Docker image的选择，经过我多次测试，发现不光要cuda，还得要cudnn那一套，针对目前onnx 1.18的Java支持，最好还是选用nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04（点击连接可以直接在vast ai上选用对应的模板）。虽然gradle在构建时可以选择额外下载打包好的cudnn，但我发现有时候这些机器的网速比较慢，与其自己花钱等着下载，不如让平台直接下载万事俱备的docker镜像。

下载完成后就可以构建使用了。为了方便，我给我这个程序准备了一个cli命令行，大概就是spc：

-s duration用来split（分割），duration的格式就是Java的Duration的格式，比如PT30S就是30秒
-p model用来process（处理）视频
-c就是用来concat（合并）视频

就是单纯的三个主要命令行选项，跟鲨鱼殴打中心没有一点关系。

关于这部分我就不多介绍了。感兴趣的话可以移步hurui200320/yolo-bird-kt，我在这个项目的简介中详细记载了如何使用。

后记

至此这一篇暂且告一段落。天气热了，我也不天天去录像了。无外乎就是那几种鸟，真想看不同的，还是得走出家门到不同的地方去看。不过这个项目做起来是挺有意思的，从最开始对yolo和onnx一无所知，到最后写出一个能运行，甚至是一个相对不错的工作流，我认为这个可能是为数不多我觉得很有趣的个人项目了。

总之这篇文章就到此为止了，这个项目我之后还会继续维护，感兴趣的话可以移步hurui200320/yolo-bird-kt，也欢迎大家在issue区提出不同的想法。

-全文完-

【代码札记】使用ONNX在JVM上调用YoloV8识别鸟类由天空 Blond 采用知识共享署名 - 非商业性使用 - 相同方式共享 4.0 国际许可协议进行许可。
本许可协议授权之外的使用权限可以从 https://skyblond.info/about.html 处获得。

Catalog