在本文中,将探讨如何利用结合了SAM 2和Florence-2的Grounded SAM 2模型来标记计算机视觉数据。通过使用Autodistill框架,可以利用大型基础模型自动标记数据,以便训练更小的、经过微调的模型,例如YOLOv8或YOLOv10。
首先,需要准备一个图像数据集。以本指南为例,将标记一个集装箱的数据集。可以通过以下命令下载数据集:
wget https://media.roboflow.com/containers.zip
该数据集包含了多个集装箱的场地图片。将使用结合了Florence-2(定位模型)和SAM 2(分割模型)的Grounded SAM 2来为数据集中的每个螺丝生成分割掩码。
接下来,将安装GroundedSAM 2模型。这是一个在Autodistill生态系统中可用的模型。Autodistill允许使用大型基础模型来自动标记数据,以便训练更小的、经过微调的模型。
pip install autodistill-grounded-sam-2
安装完模型后,可以开始测试哪些提示在数据集上表现最佳。使用Grounded SAM 2,可以提供文本提示,并为所有匹配提示的对象生成分割掩码。在自动标记完整数据集之前,测试提示以确保找到了一个能够准确识别感兴趣对象的提示是很重要的。
让从测试提示“shipping container”开始。创建一个新的Python文件,并添加以下代码:
from autodistill_grounded_sam_2 import GroundedSAM2
from autodistill.detection import CaptionOntology
from autodistill.utils import plot
import cv2
import supervision as sv
base_model = GroundedSAM2(
ontology=CaptionOntology(
{
"screw": "screw"
}
)
)
results = base_model.predict("image.png")
image = cv2.imread("image.png")
mask_annotator = sv.MaskAnnotator()
annotated_image = mask_annotator.annotate(
image.copy(), detections=results
)
sv.plot_image(image=annotated_image, size=(8, 8))
在上面的代码中,导入了所需的依赖项,然后初始化了GroundedSAM2模型的一个实例。设置了一个本体,其中“screw”是提示,保存了模型的结果,并带有类名“screw”。要自定义提示,更新第一个“screw”文本。要自定义将保存到最终数据集的标签,更新第二个“screw”实例。
然后对名为“image.png”的图像运行推理。使用supervision Python包来显示模型返回的分割掩码。
让在以下图像上运行脚本:
脚本返回了以下可视化结果:
紫色显示了集装箱被识别的位置。系统成功识别了两个集装箱!如果提示不适用于数据,更新上面的“shipping container”提示,并探索不同的提示。可能需要几次尝试才能找到适用于数据的提示。
请注意,有些对象可能无法用Grounded SAM 2识别。Grounded SAM 2最适合通用对象。例如,该模型可以识别集装箱,但可能很难找到集装箱ID的位置。
有了成功识别对象的提示,现在可以标记完整数据集了。创建一个新文件夹,并将想要标记的所有图像添加到文件夹中。
base_model.label("containers", extension="jpg")
在上面的代码中,将data替换为想要标记的图像所在的文件夹的名称。将“jpg”替换为文件夹中图像的扩展名。当准备好时,运行代码。
图像将使用Grounded SAM 2进行标记,并保存到一个新的文件夹中。当标记过程完成后,将拥有一个可以导入到Roboflow的数据集。
有了标记好的数据集,下一步是检查标签的质量并训练模型。Roboflow为这两个步骤都提供了工具。使用Roboflow,可以审查并修改注释,然后使用数据集来训练模型。
要检查数据集并训练模型,首先创建一个免费的Roboflow账户。然后,点击Roboflow仪表板上的“创建项目”。为项目设置一个名称,然后选择“实例分割”作为项目类型。
将被带到一个页面,可以在该页面上上传图像进行审查。将由GroundedSAM 2生成的标记图像文件夹拖放到页面上。图像将在浏览器中处理。当图像准备好上传时,点击“保存并继续”。
现在可以审查图像并修改注释。点击上传的数据集批次,然后选择一个图像进行检查和修改注释。
可以添加、调整和删除多边形。要了解更多关于使用Roboflow进行注释以及可用的功能,请参阅Roboflow Annotate文档。
添加新多边形的快速方法是使用Segment Anything支持的标记工具。这个工具允许点击图像的一个区域,并使用Segment Anything分配一个多边形。
要启用此功能,请单击右侧任务栏中的魔术棒功能,然后单击“增强智能多边形”。这将启用Segment Anything支持的Smart Polygon注释工具。使用这个工具,可以将鼠标悬停在图像的一部分上,查看将创建的掩码,然后单击以创建掩码。
一旦标记了图像,可以从Roboflow的“生成”标签生成数据集版本。在此标签上,可以为数据集应用增强和预处理步骤。阅读预处理和增强指南,了解如何通过预处理和增强来帮助提高模型性能。
注意:如果数据集中测试集有0张图像,请点击训练/测试拆分并重新平衡数据集,以确保测试集中有图像。需要测试集中的图像来训练模型。
当准备好时,点击数据集生成页面上的“创建”。这将创建一个带有任何选定的预处理和增强步骤的数据集版本。将被带到数据集页面,可以在那里训练模型。
点击“使用Roboflow训练”按钮来训练模型。将能够自定义训练作业。对于第一次训练作业,选择所有“推荐”选项:快速训练,并从COCO检查点训练。
当训练开始时,将能够在Roboflow仪表板上监控训练进度。
一旦有了训练好的模型,可以从Roboflow仪表板的可视化标签预览它。以下是集装箱检测模型运行的示例:
模型成功地识别了数据集中集装箱的位置。有了训练好的模型,下一步是部署模型。可以使用Roboflow API将模型部署到云端,或者部署到自己的硬件上。如果选择将模型部署到自己的硬件上,可以使用Roboflow Inference运行部署。Inference是高性能自托管计算机视觉推理服务器。