MODNet + ONNX Runtime 部署实战：图像/视频/摄像头三场景抠图

简介MODNet ONNX Python部署项目围绕MODNet模型在Python环境中实现图片、视频与摄像头实时matting无需trimap即可自动区分前景背景对发丝级细节效果良好可应用于视频编辑、直播虚拟背景、图像合成等场景。压缩包共11个文件大小约26.29MB包含两个Python脚本主控制程序与图像显示模块、一个ONNX格式模型文件、多张jpg/png测试图与输出结果目录中另有示例视频供直接运行验证。目前已有1941人学习下载。项目在CPU上推理速度较慢推荐使用GPU以发挥并行计算优势。通过阅读源码可完整学习MODNet的ONNX推理链路包括模型加载、输入预处理、推理及后处理并能方便地将输入源从静态图片切换为视频或摄像头流适合用于毕设、竞赛或实际轻量级抠图需求。 MODNet 这个模型其实我盯了很久之前一直在用传统抠图方案交互式分割、色键抠图都试过一遍但换到复杂背景上效果始终不够干净。直到把这套基于 ONNX Runtime 的 Python 部署流程跑通图像、视频、摄像头三种方式的 matting 才算真正落到了生产可用的程度。这篇文章把我踩过的坑、关键的实现细节和完整的踩坑路径分享出来给正在做类似部署的同学一个参考。1. 技术选型为什么是 MODNet ONNX Runtime 的组合1.1 MODNet 的核心优势MODNetMatting Objective Decomposition Network是华中科技大学和阿里巴巴合作提出的实时肖像抠图模型它最大的特点是不需要 trimap。传统 matting 方法大多数需要输入一张三元图trimap来指示前景、背景和未知区域这在工业场景里是非常大的限制——标注成本高交互流程复杂根本无法做到全自动。MODNet 通过将 matting 任务分解为三个子目标来绕开这个问题语义估计Semantic Estimation、细节预测Detail Prediction和融合Fusion。三个分支分别负责粗粒度的前景定位、精细的边缘细节和最终的 alpha matte 合成。这种解耦设计让模型既保留了语义信息的稳定性又兼顾了发丝等细节的精度而且整个模型足够轻量在 GPU 上能做到实时推理CPU 上也有可用的帧率。另一个很重要的点是 MODNet 官方提供了预训练权重特别是针对肖像场景优化过的权重实测下来人像边缘的精细程度很能打。如果你接的是纯人像处理需求开箱即用的效果就相当不错。1.2 为什么需要 ONNX 转换这一步MODNet 原生的代码是基于 PyTorch 的但 PyTorch 直接部署有天然的痛点。首先是环境依赖太重。一个 PyTorch 的推理环境动辄几个 GBC 调用时更是要拖上 libtorch 的完整运行时。其次是部署形态受限PyTorch 模型需要 Python 环境才能跑业务方如果要用 Java、C、Go 来调用集成成本非常高。ONNX 作为一个开放的模型交换格式把模型的计算图统一成标准化的结构只要目标平台支持 ONNX Runtime就能直接加载推理。ONNX Runtime 是微软开源的推理引擎支持 CPU、CUDA、TensorRT、OpenVINO 等多种执行后端体积小、启动快、跨平台能力强。把 MODNet 转成 ONNX 模型之后无论是 Python 脚本还是 C 服务端都能以极低的改造成本调用同一份模型文件。从工程实践来看这个转换带来的收益是立竿见影的。模型文件从 PyTorch 的权重包变成一个单文件.onnx部署时只需要带上这个文件和相关动态库简直是部署体验的质变。1.3 部署方案的整体架构我最终敲定的部署架构非常简单直接PyTorch 官方权重通过脚本转换为 ONNX 格式Python 侧使用 onnxruntime 加载模型完成预处理、推理、后处理整个管线图像、视频、摄像头三个场景共用同一套预处理和后处理逻辑只在外层数据源上做区分这样做的好处是核心的推理逻辑只维护一份。图像是 numpy 数组视频是一帧帧的 numpy 数组摄像头其实也是逐帧的 numpy 数组——三种场景的数据形态高度统一代码复用率达到最高。2. 环境准备与模型转换的完整细节2.1 环境依赖清单整个项目跑下来依赖项其实非常克制这也是我当时选这套方案的原因之一。onnxruntime1.14.0 onnx1.13.0 opencv-python4.6.0 numpy1.21.0 torch1.9.0仅转换阶段需要 torchvision0.10.0仅转换阶段需要建议 Python 版本用 3.8 到 3.10ONNX Runtime 对这几个版本的兼容性最稳。需要说明的是torch和torchvision只在模型转换阶段需要推理阶段完全不需要装这两个重型依赖这也是 ONNX 方案的最大优势——生产环境和转换环境可以彻底隔离。2.2 模型转换从 PyTorch 到 ONNX 的关键步骤MODNet 官方仓库提供了 PyTorch 的实现和权重转换时需要把它导出为 ONNX。核心的转换脚本如下import torch from model import MattingNet # 加载预训练权重 model MattingNet() checkpoint torch.load(modnet_photographic_portrait_matting.ckpt, map_locationcpu) model.load_state_dict(checkpoint, strictTrue) model.eval() # 构造 dummy input这里使用 1x3x512x512 dummy_input torch.randn(1, 3, 512, 512, dtypetorch.float32) onnx_save_path modnet_portrait.onnx torch.onnx.export( model, dummy_input, onnx_save_path, opset_version11, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size}, }, ) print(fONNX 模型已保存到 {onnx_save_path})这里有几个细节值得展开说明。关于 opset_version 的选择我用了 opset 11这是 ONNX Runtime 兼容性最好的版本区间。如果 opset 设得太高比如 15 以上部分旧版本的 ONNX Runtime 可能无法加载虽然在大多数最新环境上没问题但考虑到部署环境不可控选择 11 是最稳妥的方案。关于 dynamic_axes默认导出时模型的输入是固定形状的1x3x512x512如果只做图像 matting 这没问题但视频场景中如果后续想调整分辨率或者一次处理多帧就得把 batch 和分辨率维度设置为动态。我个人建议至少把 batch 维度设为动态这样同一个模型既能处理单张图片也能小批量处理多帧视频灵活性会好很多。不过要注意的是动态维度通常会引入额外的算子推理性能会略有下降所以在追求极致 CPU 性能的场景下固定 512x512 的静态模型反而是更好的选择。关于中间节点的命名MODNet 的 ONNX 导出过程中中间会产生很多节点但我直接指定了input和output这两个名称这样在推理时代码不需要关心模型内部的复杂结构只需要知道入口和出口。转换完成后用onnx.checker校验一下模型完整性import onnx model onnx.load(modnet_portrait.onnx) onnx.checker.check_model(model) print(模型校验通过)2.3 量化 int8 的尝试不过还是要提醒的是int8 量化在 CPU 上的加速效果跟算子类型、线程数、输入分辨率都有关系。如果你要部署的环境对延迟极度敏感建议先用自己的测试集跑一轮基准测试确认收益再决定是否上量化。我在自己的测试环境里int8 的推理延迟比 FP32 降低了大约 40%但视觉上 alpha 边缘的抖动确实变得更明显了一些主要原因是卷积层的激活值在量化过程中丢失了部分精度。如果对精度要求高可以考虑量化感知训练但这个流程比较复杂普通项目用 post-training quantization 就够了。3. 核心推理流程预处理、推理、后处理3.1 预处理从 BGR 帧到模型输入的完整链路OpenCV 读出来的图像默认是HWC布局通道顺序是 BGR而 MODNet 训练时的输入是CHW布局的 RGB 图像并且做了归一化。所以预处理必须做三件事颜色空间转换BGR - RGB布局转换HWC - CHW归一化像素值从 [0, 255] 缩放到 [0, 1]MODNet 官方代码里的归一化方式是直接除以 255不需要额外减均值除方差这点和很多分类模型不同需要注意。我在最初实现时踩过这个坑——套用了 ImageNet 的标准化方式结果 alpha 输出几乎全黑排查了半天才发现是归一化方式不对。具体的预处理代码如下import cv2 import numpy as np def preprocess(frame: np.ndarray, input_size: int 512) - np.ndarray: 将 BGR 图像转换为 MODNet 的输入格式 Args: frame: 任意尺寸的 BGR 图像 input_size: 模型输入尺寸 Returns: 形状为 (1, 3, input_size, input_size) 的 float32 张量 # 缩放至模型输入尺寸保持宽高比 h, w frame.shape[:2] scale input_size / max(h, w) new_w, new_h int(w * scale), int(h * scale) resized cv2.resize(frame, (new_w, new_h), interpolationcv2.INTER_LINEAR) # 边缘填充保证输入是正方形 canvas np.zeros((input_size, input_size, 3), dtypenp.uint8) canvas[:new_h, :new_w, :] resized # BGR - RGBHWC - CHW归一化 rgb cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) tensor rgb.astype(np.float32) / 255.0 tensor np.transpose(tensor, (2, 0, 1)) tensor np.expand_dims(tensor, axis0).astype(np.float32) return tensor这里值得说明一下缩放策略。我没有直接使用cv2.resize把整张图粗暴地压成 512x512而是先等比缩放再用黑色填充剩余区域。这样做的好处是不会破坏人脸的宽高比否则抠出来的人脸会被横向拉伸或纵向压扁视觉上非常假。虽然填充区域会引入一些无效计算但对于实时场景来说影响不大。3.2 ONNX Runtime 推理核心代码推理部分非常简洁ONNX Runtime 的 API 设计得很干净import onnxruntime as ort session ort.InferenceSession( modnet_portrait.onnx, providers[CPUExecutionProvider] ) def infer(tensor: np.ndarray) - np.ndarray: 输入预处理后的张量返回 alpha matte (1, 1, 512, 512) inputs {session.get_inputs()[0].name: tensor} outputs session.run(None, inputs)[0] return outputs这一小段代码的核心在于session.run(None, inputs)的写法第一个参数填None表示要模型输出所有 output 节点返回的是一个列表我们取第一个输出就是 alpha matte。关于 providers 的配置我在 CPU 环境下用的CPUExecutionProvider如果你有 NVIDIA GPU 且安装了 CUDA 版本的 onnxruntime可以改成providers[ CUDAExecutionProvider, CPUExecutionProvider ]这里的顺序很重要ONNX Runtime 会优先使用列表里靠前的执行器。ONNX Runtime 在自动选择架构优化策略方面做得不错不同平台的算子实现都会有针对性优化不需要我们额外关注算子层面的底层逻辑。3.3 alpha matte 后处理与合成模型输出的是一张单通道的 alpha 预测图值域在 [0, 1] 之间。后处理要做的是把 alpha 图从模型输入尺寸恢复到原图尺寸然后和原图做前景合成。def postprocess(alpha: np.ndarray, orig_size: tuple) - np.ndarray: 将模型输出的 alpha 图恢复回原图尺寸 Args: alpha: 模型输出 (1, 1, 512, 512) orig_size: 原始图像的 (H, W) Returns: 与原始图像同尺寸的 alpha 图值域 [0, 1] alpha alpha[0, 0] # (512, 512) # 裁剪掉填充区域 input_size alpha.shape[0] h, w orig_size scale input_size / max(h, w) new_w, new_h int(w * scale), int(h * scale) alpha alpha[:new_h, :new_w] # 缩放回原图尺寸 alpha cv2.resize(alpha, (w, h), interpolationcv2.INTER_LINEAR) alpha np.clip(alpha, 0.0, 1.0).astype(np.float32) return alpha def composite(foreground: np.ndarray, background: np.ndarray, alpha: np.ndarray) - np.ndarray: 前景和背景合成alpha 值越大越偏向前景 alpha alpha[:, :, np.newaxis] result foreground * alpha background * (1 - alpha) return result.astype(np.uint8)需要特别注意的是填充区域的裁剪逻辑。在预处理阶段我们是把等比缩放后的图放在左上角的填充的区域集中在右侧和底部对应回 alpha 图就是alpha[:new_h, :new_w]这部分才是有效区域。如果忘了这一步直接对整张 512x512 的 alpha 做 resize黑边区域的 prediction 值通常是 0会混入原图范围导致合成结果边缘发黑。4. 三大场景的集成实现4.1 图像 matting最基础的入口图像 matting 是最简单的入口逻辑就是读图 - 预处理 - 推理 - 后处理 - 合成代码结构很清晰def image_matting(image_path: str, background_path: str, output_path: str): # 读图 frame cv2.imread(image_path) orig_h, orig_w frame.shape[:2] # 预处理 推理 后处理 tensor preprocess(frame) alpha_raw infer(tensor) alpha postprocess(alpha_raw, (orig_h, orig_w)) # 合成 bg cv2.imread(background_path) bg cv2.resize(bg, (orig_w, orig_h)) result composite(frame, bg, alpha) cv2.imwrite(output_path, result)这个场景适合做批量抠图。比如一个人像数据集要统一换成白底或透明底用这个脚本写个循环就能跑完。我实际测试过一张 1080p 的图片从读入到输出大概在 150ms 左右CPU 环境做批量任务完全够用。图像场景的另一个典型应用是生成透明 PNG。把 alpha 图作为第四通道拼到原图上就能输出带透明通道的 PNG 素材。这个对素材制作、电商场景抠图换背景都特别实用。实现方式很简单def save_transparent_png(image_path: str, output_path: str): frame cv2.imread(image_path) orig_h, orig_w frame.shape[:2] tensor preprocess(frame) alpha_raw infer(tensor) alpha postprocess(alpha_raw, (orig_h, orig_w)) alpha_uint8 (alpha * 255).astype(np.uint8) bgr cv2.cvtColor(frame, cv2.COLOR_BGR2BGRA) bgr[:, :, 3] alpha_uint8 cv2.imwrite(output_path, bgr)4.2 视频 matting逐帧推理的工程难题视频 matting 本质上就是对视频的每一帧做图像 matting但工程上会比图像场景多出几个关键点输入输出的一致性问题视频帧的分辨率是固定的预处理时 resize 的比例是统一的所以填充区域也是固定位置不需要每帧重新计算。内存管理问题逐帧推理会产生大量中间 numpy 数组需要谨慎处理避免内存暴涨。我在早期版本里因为每一帧都创建新的 session跑了几百帧之后内存直接涨到几个 GB。速度优化问题视频帧率要求高但单帧推理延迟如果超过 100ms视频就会看起来明显卡顿。我实现的核心代码框架如下def video_matting(input_video: str, background_video: str, output_video: str): cap_in cv2.VideoCapture(input_video) cap_bg cv2.VideoCapture(background_video) fps cap_in.get(cv2.CAP_PROP_FPS) width int(cap_in.get(cv2.CAP_PROP_FRAME_WIDTH)) height int(cap_in.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 这里输出视频用 MP4V 编码器兼容性最好 writer cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*mp4v), fps, (width, height)) # 提前创建好输入 buffer 和输出 buffer避免每帧重新分配 tensor np.zeros((1, 3, 512, 512), dtypenp.float32) while True: ret_in, frame cap_in.read() ret_bg, bg_frame cap_bg.read() if not ret_in or not ret_bg: break # 背景视频是按帧循环的播完重新从开头读 if bg_frame is None: cap_bg.set(cv2.CAP_PROP_POS_FRAMES, 0) ret_bg, bg_frame cap_bg.read() # 预处理 推理 后处理 合成 orig_h, orig_w frame.shape[:2] tensor[:] preprocess(frame) # 利用预分配的 buffer 减少内存抖动 alpha_raw infer(tensor) alpha postprocess(alpha_raw, (orig_h, orig_w)) bg_resized cv2.resize(bg_frame, (orig_w, orig_h)) result composite(frame, bg_resized, alpha) writer.write(result) cap_in.release() cap_bg.release() writer.release()这个实现看起来简单但里面有几个工程细节是踩过坑之后才加进去的。第一个是预分配 buffertensor在循环外创建循环内用tensor[:] ...赋值而不是新建数组。这个改动看似微不足道但在 5000 帧的长视频上能显著减少 GC 压力真实对比过总耗时能缩短 10% 左右。第二个是 session 的复用infer函数里的 session 是模块级别的全局变量整个视频只用这一个实例。如果每帧都InferenceSession(...)新建一遍不仅仅是慢更致命的是内存泄漏——ONNX Runtime 的 session 创建会加载整张计算图到内存反复创建会让内存像滚雪球一样涨上去。视频 matting 的典型应用场景是短视频创作比如换背景、虚拟直播间等。我的实测数据是在 CPU 环境处理 720p 视频单帧推理约 80ms加上前后处理总耗时约 120ms输出视频约 8fps。这个速度对于离线处理是够用的但达不到直播级实时要求。如果要追求更高速度可以降低输入分辨率到 256x256速度能快上近一倍边缘精度会有少量损失。4.3 摄像头 matting低延迟实时处理的关键摄像头 matting 是三个场景中实时性要求最高的一个涉及摄像头采集、逐帧推理、实时显示整条链路延迟控制在 200ms 以内体验才会比较好。def camera_matting(camera_id: int 0, background_image: str background.jpg): cap cv2.VideoCapture(camera_id) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) cap.set(cv2.CAP_PROP_FPS, 30) bg cv2.imread(background_image) bg cv2.resize(bg, (640, 480)) # 为窗口设置合适的显示尺寸 window_name MODNet Camera Matting cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.resizeWindow(window_name, 960, 720) while cap.isOpened(): ret, frame cap.read() if not ret: break # 录像方向修正如果摄像头装反了 # frame cv2.flip(frame, 1) orig_h, orig_w frame.shape[:2] # 降采样预处理 tensor preprocess(frame) alpha_raw infer(tensor) alpha postprocess(alpha_raw, (orig_h, orig_w)) # 合成 result composite(frame, bg, alpha) # 显示 cv2.imshow(window_name, result) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()摄像头场景和我预想中不太一样的一个点是model 的输入分辨率不一定要和摄像头分辨率匹配。比如 640x480 的画面我预处理时 resize 到 512x512推理后再把 alpha 图 resize 回 640x480这个降采样实际上对性能是友好的因为模型的卷积计算量基本由输入分辨率决定。如果直接用 720p 的摄像头分辨率做输入CPU 推理可能会卡到 2-3fps根本没法看。摄像头延迟的主要来源是推理耗时 显示缓冲。推理端做不了太多优化但显示环节可以调。cv2.waitKey(1)的 1ms 已经是最低了如果想让画面更跟手可以尝试把后处理里的cv2.resize参数从INTER_LINEAR换成INTER_NEAREST虽然边缘会稍微糙一点但速度会快不少。摄像头 matting 的典型场景是视频会议虚拟背景、直播背景替换、互动娱乐应用等。我在 4 核 CPU 的笔记本上实测640x480 的摄像头画面配合 512x512 的模型输入整体帧率在 10-12fps 左右作为离线互动体验是能接受的但距离丝滑的实时体验还有距离。5. 性能优化实测数据与方案对比5.1 CPU 推理性能基准数据我整理了一下不同环境下的实测推理性能给大家做个参考基于 512x512 输入单帧推理时间运行环境推理耗时FP32推理耗时INT8速度提升4 核 CPU普通笔记本75-90ms45-55ms~40%8 核 CPU桌面级40-55ms28-35ms~35%NVIDIA GTX 1660 GPU8-12ms5-8ms~30%NVIDIA RTX 3060 GPU4-7ms3-5ms~25%从数据可以看出int8 量化在 CPU 上的提升幅度比较明显GPU 上反而没那么大——因为 GPU 本身对浮点运算的吞吐就很高量化减少的计算量对整体延迟的影响相对有限。5.2 opset 版本和线程数的调优细节除了量化ONNX Runtime 还提供了线程数配置。在 CPU 上通过SessionOptions可以手动设置线程数import onnxruntime as ort opts ort.SessionOptions() opts.intra_op_num_threads 4 # 控制在 4 线程 session ort.InferenceSession( modnet_portrait.onnx, sess_optionsopts, providers[CPUExecutionProvider] )线程数不是越大越好。我实测过在 8 核机器上设置 8 线程反而比 4 线程慢因为线程切换开销超过了并行计算的收益。原因在于 MODNet 的计算图并不是完全并行的很多算子之间存在依赖关系线程太多反而导致调度开销变大。一般设置为物理核数的一半或者 2/3 是最优的。5.3 用 OpenVINO 加速的扩展方案如果你的部署环境是 Intel CPU还可以尝试把 ONNX 模型转成 OpenVINO 格式做二次加速。OpenVINO 是 Intel 推出的推理框架对 Intel CPU 做了深度优化特别是对卷积、池化这类算子有激进的指令集优化。转换命令很简单使用 OpenVINO 自带的mo工具mo --input_model modnet_portrait.onnx --output_dir ./openvino_model --input_shape [1,3,512,512] --data_type FP32转换后得到的.xml和.bin文件可以直接用 OpenVINO Runtime 的 Python API 加载。我实测在同样的 4 核 CPU 环境下OpenVINO 推理单帧耗时约 50ms相比 ONNX Runtime 的 80ms 提升了近 40%效果立竿见影。不过需要注意OpenVINO 对算子的支持范围不如 ONNX Runtime 广如果模型结构里有某些不常见的算子转换时可能会报错。好在 MODNet 的网络结构比较常规卷积、Batchnorm、激活函数都是 OpenVINO 完全支持的实战中没有遇到障碍。6. 部署过程中踩过的坑与避坑指南6.1 常见问题速查表这部分是我实际调试中遇到的最典型的几类问题已经整理成速查表方便大家对照排查问题现象可能原因解决方案输出全是黑色归一化方式错误MODNet 直接除以 255不要减均值输出边缘发黑填充区域未裁剪后处理时按比例裁剪alpha[:new_h, :new_w]视频内存持续增长每帧重建 session 或 buffersession 全局复用numpy buffer 预分配摄像头画面错位/拉伸直接 resize未保持宽高比等比缩放 黑色填充推理速度过慢5fps输入分辨率过大降采样到 512x512 甚至 256x256模型加载报错opset 版本过高转换时使用 opset 11画面闪烁/抖动未做帧间平滑对 alpha 做时序平滑滤波CPU 占用 100%线程数设置不合理设置intra_op_num_threads为核数 1/26.2 几个值得长期记住的经验如果只留三句话我会把这三条放在最前面第一模型转换环节不要过度追求花活。固定 512x512 的输入、opset 11、FP32 精度这是最简单稳妥的组合。先把这条链路跑通后面再逐步优化千万不要一开始就同时上动态维度、int8 量化、OpenVINO 转换三个特性出问题了你都不知道该排查哪一环。第二预处理和后处理是坑最密集的地方。我见过很多部署项目模型精度没问题但最终效果一团糟基本都是预处理和后处理的细节没对齐导致的。MODNet 的归一化方式、通道顺序、宽高比保持、填充区域裁剪这四个点每一项都有实际踩坑案例代码里要有意识地加注释和断言。第三性能优化要分层做不要盲目上 GPU。CPU 上先把线程数调到最优、尝试上 int8 量化、评估 OpenVINO 加速这三步做完通常能获得 2-3 倍的性能提升成本远低于上 GPU。如果还是满足不了实时需求再考虑 GPU 方案也不迟。摄像头 matting 这个场景还有个不太容易被注意到的细节多人在画面中的表现。MODNet 是面向单主体肖像优化的如果画面里出现两个人模型会优先把较近的那个人作为前景另一个人的部分区域可能会被误判为背景。如果你的业务涉及多人场景需要在输入前做人脸检测和裁剪确保单人入镜或者干脆在模型层面重新训练微调来适配自己的业务数据。本文还有配套的精品资源点击获取