EdgeSAM: 边缘设备上的高效分割模型

EdgeSAM

EdgeSAM: 边缘设备上的高效分割模型

随着计算机视觉技术的快速发展,图像分割作为一项基础任务,其应用场景正在不断拓展。然而,许多先进的分割模型由于计算复杂度高,难以在边缘设备上实时运行。为了解决这一问题,南洋理工大学S-Lab和上海人工智能实验室的研究团队提出了EdgeSAM (Edge Segment Anything Model),这是一种针对边缘设备优化的高效分割模型。

EdgeSAM的核心创新

EdgeSAM是对原始Segment Anything Model (SAM)的一种加速变体。其核心创新在于:

将基于ViT的图像编码器蒸馏为纯CNN架构,更适合边缘设备。
采用"提示在循环中"(Prompt-In-the-Loop)的蒸馏策略,保留了SAM的灵活性。
在边缘设备上实现了显著的性能提升,同时保持了较高的分割精度。

这些创新使EdgeSAM成为首个能在iPhone 14上以超过30 FPS运行的SAM变体,为边缘设备上的实时分割任务开辟了新的可能性。

性能对比

根据研究团队的测试,EdgeSAM在多个方面都表现出色:

速度提升: 相比原始SAM,EdgeSAM实现了40倍的速度提升。
边缘设备性能: 在边缘设备上,EdgeSAM比MobileSAM快14倍。
精度保持: 在COCO和LVIS数据集上,EdgeSAM的mIoU分别提高了2.3和3.2。

EdgeSAM性能对比

上图展示了EdgeSAM与SAM和MobileSAM在编码器吞吐量和SA-1K数据集上的mIoU性能对比。可以看出,EdgeSAM在保持较高精度的同时,实现了显著的速度提升。

技术原理

EdgeSAM的核心技术原理包括:

CNN架构: 将原始SAM的ViT编码器替换为纯CNN架构,大幅降低了计算复杂度。
提示在循环中的蒸馏: 在蒸馏过程中引入提示编码器和掩码解码器,使得蒸馏后的模型能够准确捕捉用户输入和掩码生成之间的复杂动态关系。
稳定性得分: 由于EdgeSAM没有对原始SAM的IoU标记进行知识蒸馏,其IoU预测可能不太可靠。因此,EdgeSAM引入了稳定性得分来进行掩码选择。

EdgeSAM架构

实际应用

EdgeSAM的高效性能使其在多个领域都有潜在的应用价值:

移动设备AR/VR: 实时场景分割可以增强AR/VR体验。
智能手机摄影: 实时物体分割可用于高级摄影效果。
自动驾驶: 边缘设备上的实时分割对自动驾驶至关重要。
机器人视觉: 提高机器人在复杂环境中的感知能力。
医疗影像分析: 在便携式医疗设备上实现快速诊断。

安装和使用

要使用EdgeSAM,您需要按照以下步骤进行安装和配置:

克隆GitHub仓库:

git clone https://github.com/chongzhou96/EdgeSAM.git && cd EdgeSAM

安装依赖:

pip install -r requirements.txt

下载预训练模型:

mkdir weights
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam.pth
wget -P weights/ https://huggingface.co/spaces/chongzhou/EdgeSAM/resolve/main/weights/edge_sam_3x.pth

在Python代码中使用EdgeSAM:

from edge_sam import SamPredictor, sam_model_registry
sam = sam_model_registry["edge_sam"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)

模型导出和优化

EdgeSAM支持导出为CoreML和ONNX格式,以便在不同平台上部署:

CoreML导出:
- 编码器: python scripts/export_coreml_model.py [CHECKPOINT]
- 解码器: python scripts/export_coreml_model.py [CHECKPOINT] --decoder --use-stability-score
ONNX导出:
- 编码器: python scripts/export_onnx_model.py [CHECKPOINT]
- 解码器: python scripts/export_onnx_model.py [CHECKPOINT] --decoder --use-stability-score

这些导出的模型可以在边缘设备上实现更高的推理效率。

未来展望

EdgeSAM的成功为边缘设备上的高效图像分割开辟了新的道路。未来的研究方向可能包括:

进一步优化模型结构,以适应更多类型的边缘设备。
探索将EdgeSAM与其他计算机视觉任务结合的可能性。
开发针对特定应用场景的定制版EdgeSAM。
研究如何在不损失性能的情况下进一步压缩模型大小。

总结

EdgeSAM代表了计算机视觉领域的一个重要进展,它成功地将高性能的图像分割能力带到了边缘设备上。通过创新的模型架构和蒸馏策略,EdgeSAM在速度和精度之间取得了优秀的平衡,为实时、高质量的图像分割任务提供了新的解决方案。随着边缘计算和AI技术的不断发展,我们可以期待看到更多像EdgeSAM这样的创新模型,推动计算机视觉技术在更广泛的应用场景中发挥作用。

EdgeSAM的出现不仅推动了学术研究的前沿,也为工业界提供了宝贵的实用工具。它的成功展示了如何将复杂的AI模型有效地部署到资源受限的设备上,这对于物联网(IoT)、智能手机应用程序开发者和嵌入式系统工程师都具有重要的启发意义。随着EdgeSAM的进一步发展和优化,我们可以期待看到更多基于这一技术的创新应用,为用户带来更智能、更高效的图像处理体验。