RMBG-2.0轻量模型实战Jetson Nano边缘设备部署实时视频流抠图演示想不想在树莓派大小的设备上实现实时视频抠图把人物从背景里干净利落地“摘”出来今天我们就来聊聊怎么在Jetson Nano这块小小的边缘计算板上部署RMBG-2.0这个轻量级AI抠图神器并让它流畅地处理摄像头视频流。RMBG-2.0是个专为“背景去除”而生的AI模型。它最大的特点就是“小而强”模型本身很轻不需要动辄几十GB的显存在Jetson Nano这种资源有限的设备上也能跑得起来同时它的抠图精度又很高像头发丝、玻璃杯、纱巾这些传统算法很难处理的细节它都能应付得不错。无论是给电商商品换背景还是制作短视频素材它都能派上大用场。这篇文章我就手把手带你走通整个流程从在Jetson Nano上准备好环境到把RMBG-2.0模型跑起来最后实现连接USB摄像头进行实时抠图演示。你会发现在边缘端玩转AI并没有想象中那么复杂。1. 项目与环境准备在开始敲代码之前我们得先把“舞台”搭好。这里的主角是Jetson Nano我们需要确保它的系统环境能支撑起RMBG-2.0模型的运行。1.1 硬件与系统要求首先确认一下你的设备是否满足基本要求硬件NVIDIA Jetson Nano开发板2GB或4GB版本均可。建议使用4GB版本以获得更流畅的体验。系统官方推荐的JetPack SDK包含Ubuntu 18.04或20.04 LTS、CUDA、cuDNN等。你可以通过命令cat /etc/nv_tegra_release来查看JetPack版本。外设一个USB摄像头用于视频流输入。一个显示器或通过SSH远程连接。如果你的Jetson Nano系统是全新的第一步最好是做一次全面的系统更新和升级这能避免很多后续的依赖库问题。sudo apt update sudo apt upgrade -y1.2 安装Python与核心依赖RMBG-2.0模型通常通过Python来调用。Jetson Nano自带了Python3但我们还需要安装一些关键的库特别是深度学习框架和图像处理库。安装pip和基础工具sudo apt install -y python3-pip python3-dev安装PyTorch这是运行大多数AI模型的基础。对于Jetson平台NVIDIA提供了预编译的PyTorch wheel包安装起来最方便。你需要根据你的JetPack版本去NVIDIA官方论坛找到对应的下载链接。例如对于JetPack 4.6Python 3.6命令可能如下# 示例命令请替换为对应你JetPack版本的正确URL wget https://developer.download.nvidia.com/compute/redist/jp/v461/pytorch/torch-1.10.0-cp36-cp36m-linux_aarch64.whl pip3 install torch-1.10.0-cp36-cp36m-linux_aarch64.whl安装后可以运行python3 -c import torch; print(torch.__version__)来验证。安装其他必要库pip3 install opencv-python-headless pillow numpyopencv-python-headless用于摄像头捕获、视频处理和图像显示headless版本不含GUI功能更轻量适合服务器或远程SSH场景。如果需要本地显示窗口可以安装opencv-python。pillow(PIL)一个常用的图像处理库。numpy科学计算基础库。2. 获取与部署RMBG-2.0模型环境准备好后接下来就是把主角——RMBG-2.0模型请到我们的设备上。2.1 下载模型文件RMBG-2.0的模型文件通常是一个.onnx格式的文件。ONNXOpen Neural Network Exchange是一种开放的模型格式可以被多种推理引擎支持部署起来比较灵活。你可以从模型的官方发布页面例如Hugging Face Model Hub或GitHub Releases找到并下载这个文件。这里我们假设你已经下载好了名为rmbg-2.0.onnx的模型文件并将其放在了Jetson Nano的某个目录下比如~/projects/rmbg/。如果没有现成的ONNX文件你可能需要从PyTorch或TensorFlow格式转换过来这个过程需要一台性能更强的机器来完成不在本文讨论范围内。2.2 安装ONNX Runtime推理引擎为了在Jetson Nano上高效运行.onnx模型我们使用ONNX Runtime。幸运的是ONNX Runtime提供了针对Jetson ARM架构的预编译版本。访问 ONNX Runtime GitHub Release页面寻找标签为jetson的版本。例如下载适用于JetPack 4.6和Python 3.6的wheel文件# 示例命令请根据你的实际Python版本和JetPack版本调整文件名 wget https://github.com/microsoft/onnxruntime/releases/download/v1.14.1/onnxruntime_gpu-1.14.1-cp36-cp36m-linux_aarch64.whl pip3 install onnxruntime_gpu-1.14.1-cp36-cp36m-linux_aarch64.whl安装GPU版本可以利用Jetson Nano的CUDA核心加速推理速度会比CPU版本快很多。3. 编写图片抠图测试脚本在处理视频流之前我们先写一个简单的脚本来测试模型对单张图片的抠图效果确保核心功能是正常的。创建一个名为test_image.py的文件import cv2 import numpy as np from PIL import Image import onnxruntime as ort import time def load_and_preprocess_image(image_path, target_size(1024, 1024)): 加载图片并预处理成模型需要的格式。 RMBG-2.0通常要求输入为RGB格式尺寸为1024x1024或特定尺寸数值归一化到[0,1]。 # 使用PIL打开图片并确保为RGB img_pil Image.open(image_path).convert(RGB) # 调整大小 img_pil img_pil.resize(target_size, Image.Resampling.LANCZOS) # 转换为numpy数组并调整维度顺序为 (C, H, W) img_np np.array(img_pil).astype(np.float32) / 255.0 img_np img_np.transpose(2, 0, 1) # HWC to CHW # 添加批次维度 (N, C, H, W) img_np np.expand_dims(img_np, axis0) return img_np, img_pil.size # 返回处理后的图像和原始尺寸用于后处理还原 def postprocess_mask(mask_output, original_size): 对模型输出的掩码进行后处理。 模型输出可能是 (1, 1, H, W) 的浮点数矩阵值在0~1之间。 # 移除批次和通道维度得到 (H, W) mask mask_output[0, 0] # 将概率值转换为0-255的灰度图二值化阈值可调例如0.5 mask_uint8 (mask 0.5).astype(np.uint8) * 255 # 将掩码缩放到原始图片尺寸 mask_resized cv2.resize(mask_uint8, original_size, interpolationcv2.INTER_NEAREST) return mask_resized def apply_mask_to_image(original_img_path, mask): 将生成的掩码应用到原图上生成抠图结果透明背景或自定义背景。 original_img cv2.imread(original_img_path) original_img cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB) # 确保颜色通道一致 # 将掩码转换为三通道 mask_3ch cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB) # 创建一个带透明通道RGBA的结果图 result_rgba np.zeros((original_img.shape[0], original_img.shape[1], 4), dtypenp.uint8) result_rgba[..., :3] original_img result_rgba[..., 3] mask # 透明度通道由掩码决定 return Image.fromarray(result_rgba, RGBA) def main(): # 1. 初始化ONNX Runtime会话 model_path ./rmbg-2.0.onnx # 修改为你的模型路径 providers [CUDAExecutionProvider, CPUExecutionProvider] # 优先使用GPU session ort.InferenceSession(model_path, providersproviders) input_name session.get_inputs()[0].name output_name session.get_outputs()[0].name # 2. 处理单张图片 test_image_path ./test_photo.jpg # 准备一张测试图片 processed_img, orig_size load_and_preprocess_image(test_image_path) # 3. 模型推理 start_time time.time() outputs session.run([output_name], {input_name: processed_img}) inference_time time.time() - start_time print(f单张图片推理耗时: {inference_time:.3f} 秒) # 4. 后处理并保存结果 mask postprocess_mask(outputs[0], orig_size) result_image apply_mask_to_image(test_image_path, mask) result_image.save(./output_result.png) print(抠图结果已保存为 output_result.png) # 5. 简单显示如果本地有GUI环境 # cv2.imshow(Mask, mask) # cv2.waitKey(0) # cv2.destroyAllWindows() if __name__ __main__: main()运行这个脚本前请确保在同一目录下准备好了rmbg-2.0.onnx模型文件和一张名为test_photo.jpg的测试图片。运行命令python3 test_image.py如果一切顺利你会看到控制台打印出推理时间并生成一张背景透明的output_result.png图片。检查一下头发丝边缘等细节处理得如何这能直观感受模型的精度。4. 实现实时视频流抠图单张图片测试通过后我们就可以挑战更酷的实时视频流处理了。核心思路是从摄像头逐帧读取图片对每一帧执行抠图然后将原始帧和抠图结果或合成效果实时显示出来。创建一个新的文件realtime_video.pyimport cv2 import numpy as np import onnxruntime as ort import time from threading import Thread, Lock from queue import Queue class VideoBGRemover: def __init__(self, model_path, camera_id0, target_size(512, 512)): 初始化视频背景移除器。 target_size 可以调小以提升速度但可能会降低精度。 self.target_size target_size self.running False # 初始化模型 print(正在加载RMBG-2.0模型...) providers [CUDAExecutionProvider, CPUExecutionProvider] self.session ort.InferenceSession(model_path, providersproviders) self.input_name self.session.get_inputs()[0].name self.output_name self.session.get_outputs()[0].name print(模型加载完毕。) # 初始化摄像头 self.cap cv2.VideoCapture(camera_id) if not self.cap.isOpened(): raise IOError(f无法打开摄像头 {camera_id}) # 设置一个较低的分辨率以提升性能 self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) print(摄像头已就绪。) # 用于线程间通信的队列和锁 self.frame_queue Queue(maxsize2) # 缓冲队列避免堆积 self.result_queue Queue(maxsize2) self.lock Lock() # 性能统计 self.fps 0 self.inference_time 0 def preprocess_frame(self, frame): 预处理摄像头捕获的帧 # 转换颜色空间 BGR - RGB frame_rgb cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 调整尺寸 frame_resized cv2.resize(frame_rgb, self.target_size, interpolationcv2.INTER_LINEAR) # 归一化并转换维度 (H,W,C) - (C,H,W) - (N,C,H,W) img_np frame_resized.astype(np.float32) / 255.0 img_np img_np.transpose(2, 0, 1) img_np np.expand_dims(img_np, axis0) return img_np, frame.shape[:2] # 返回处理后的张量和原始帧尺寸 def inference_worker(self): 推理工作线程 while self.running: try: # 从队列获取帧数据超时避免死锁 data self.frame_queue.get(timeout1.0) if data is None: break processed_img, orig_size data # 执行模型推理 start_infer time.time() outputs self.session.run([self.output_name], {self.input_name: processed_img}) self.inference_time time.time() - start_infer # 后处理生成掩码并缩放到原始尺寸 mask_output outputs[0] mask mask_output[0, 0] mask_uint8 (mask 0.5).astype(np.uint8) * 255 mask_resized cv2.resize(mask_uint8, (orig_size[1], orig_size[0]), interpolationcv2.INTER_NEAREST) # 将结果放入结果队列 self.result_queue.put(mask_resized) except Exception as e: if self.running: print(f推理线程出错: {e}) continue def capture_worker(self): 摄像头捕获线程 last_time time.time() frame_count 0 while self.running: ret, frame self.cap.read() if not ret: print(无法从摄像头读取帧。) break # 计算FPS frame_count 1 current_time time.time() if current_time - last_time 1.0: self.fps frame_count frame_count 0 last_time current_time # 预处理并放入队列如果队列未满 if not self.frame_queue.full(): processed_data self.preprocess_frame(frame) self.frame_queue.put((processed_data[0], processed_data[1])) # 从结果队列获取最新的掩码 mask None while not self.result_queue.empty(): mask self.result_queue.get() # 只保留最新的结果 # 显示结果 if mask is not None: # 将掩码应用到原帧这里演示生成绿色背景的合成效果 green_bg np.zeros_like(frame) green_bg[:] [0, 255, 0] # 绿色背景 mask_3ch cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0 composite frame * mask_3ch green_bg * (1 - mask_3ch) composite composite.astype(np.uint8) # 在画面上叠加性能信息 info_text fFPS: {self.fps} | Infer: {self.inference_time*1000:.1f}ms cv2.putText(composite, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) cv2.imshow(Real-time Background Removal, composite) else: # 如果没有掩码显示原始帧 cv2.imshow(Real-time Background Removal, frame) # 按q键退出 if cv2.waitKey(1) 0xFF ord(q): self.running False break def run(self): 启动实时抠图演示 self.running True # 启动推理线程 infer_thread Thread(targetself.inference_worker, daemonTrue) infer_thread.start() print(开始实时视频抠图演示按 q 键退出。) # 在主线程中运行捕获和显示循环 self.capture_worker() def cleanup(self): 清理资源 self.running False time.sleep(0.5) # 等待线程结束 if self.cap.isOpened(): self.cap.release() cv2.destroyAllWindows() print(资源已释放。) if __name__ __main__: remover VideoBGRemover(model_path./rmbg-2.0.onnx, camera_id0, target_size(512, 512)) try: remover.run() except KeyboardInterrupt: print(程序被用户中断。) finally: remover.cleanup()这个脚本做了几件关键的事情多线程设计将耗时的模型推理放在一个独立的线程中避免阻塞摄像头帧的捕获和显示从而保证视频流的相对流畅。动态分辨率将输入模型的图片尺寸从1024x1024降低到512x512这能显著提升推理速度在Jetson Nano上实现接近实时的效果可能达到3-10 FPS取决于场景复杂度。你可以在速度和精度之间做权衡。实时合成显示不是显示透明的PNG而是将抠出的人物与一个绿色背景实时合成这样视觉效果更直观。性能监控在画面上实时显示帧率FPS和单次推理耗时方便你评估性能。运行这个脚本python3 realtime_video.py将你的USB摄像头对准自己或一个物体屏幕上应该会弹出窗口显示原始画面以及抠图后与绿色背景合成的效果。观察头发、手指等边缘的处理质量并留意左上角显示的FPS和推理时间。5. 效果评估与优化建议跑起来之后你可能会发现一些可以改进的地方。这里提供一些效果评估的角度和优化思路。5.1 效果评估维度边缘精度这是抠图的核心。仔细观察视频中人物头发梢、半透明的眼镜片、快速挥动的手部边缘。RMBG-2.0在这些复杂边缘上的表现是衡量其“轻量且强大”的关键。推理速度关注屏幕上显示的FPS和推理时间。在Jetson Nano上能达到5 FPS以上通常就算可用如果能稳定在10 FPS左右体验就比较流畅了。这很大程度上取决于你设置的target_size。资源占用你可以打开另一个终端运行jtop如果已安装或tegrastats命令查看GPU、CPU和内存的使用情况。一个优秀的轻量模型应该在提供足够精度的同时保持较低的资源消耗。5.2 性能优化技巧如果感觉速度不够快可以尝试以下方法调整输入尺寸这是最有效的杠杆。在VideoBGRemover初始化时将target_size从(512, 512)进一步下调到(384, 384)甚至(256, 256)速度会大幅提升但精度也会相应下降。你需要为你的具体应用找到一个平衡点。使用TensorRT加速这是为Jetson平台带来最大性能提升的途径。你可以将ONNX模型转换为TensorRT引擎。TensorRT是NVIDIA的高性能深度学习推理优化器和运行时能针对Jetson的硬件进行极致优化。转换过程需要一些步骤但网上有丰富的教程。模型量化尝试将模型从FP32精度转换为FP16甚至INT8精度。这能减少模型大小和内存占用并提升推理速度。ONNX Runtime和TensorRT都支持量化。代码层面优化确保没有不必要的内存拷贝。在我们的示例中队列通信和图像缩放都是潜在的性能瓶颈对于极高帧率要求可能需要更精细的设计。5.3 实际应用扩展这个演示只是一个起点你可以基于此构建更实用的应用更换背景把固定的绿色背景换成一张图片或另一个视频流就能实现虚拟背景或简单的场景合成。保存结果修改代码将处理后的视频帧序列保存为视频文件如MP4用于制作短视频素材。集成到服务中将抠图功能封装成一个Flask或FastAPI网络服务。这样你可以通过手机或其他设备上传图片或视频流让Jetson Nano在后台处理并返回结果。触发式处理不一定需要连续处理每一帧。可以设置为检测到人物移动或特定手势时才触发一次抠图并保存节省算力。6. 总结通过这篇文章我们完成了一次从零到一的边缘AI应用部署实践。在资源受限的Jetson Nano上我们成功部署了轻量级的RMBG-2.0抠图模型并实现了实时视频流的背景移除演示。整个过程的关键点在于平衡在模型精度、推理速度和设备资源之间找到最佳结合点。RMBG-2.0本身的设计就体现了这种平衡而我们的部署方案如调整输入尺寸、使用ONNX Runtime GPU推理则是在具体硬件上实现这一平衡的操作。边缘AI的魅力就在于它能将智能带到数据产生的地方。想象一下一个内置了此功能的小型设备可以放在直播间实时为主播抠像可以在线下门店为顾客自动生成带虚拟背景的趣味照片甚至可以集成到无人机上进行实时场景分析。希望这次实战能为你打开一扇门让你看到在小型化、低功耗设备上运行先进AI应用的巨大潜力。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。