《博主简介》

小伙伴们好，我是阿旭。专注于人工智能、AIGC、python、计算机视觉相关分享研究。
✌更多学习资源，可关注公-仲-hao:【阿旭算法与机器学习】，共同学习交流~
👍感谢小伙伴们点赞、关注！

《------往期经典推荐------》

一、AI应用软件开发实战专栏【链接】

项目名称	项目名称
1.【人脸识别与管理系统开发】	2.【车牌识别与自动收费管理系统开发】
3.【手势识别系统开发】	4.【人脸面部活体检测系统开发】
5.【图片风格快速迁移软件开发】	6.【人脸表表情识别系统】
7.【YOLOv8多目标识别与自动标注软件开发】	8.【基于YOLOv8深度学习的行人跌倒检测系统】
9.【基于YOLOv8深度学习的PCB板缺陷检测系统】	10.【基于YOLOv8深度学习的生活垃圾分类目标检测系统】
11.【基于YOLOv8深度学习的安全帽目标检测系统】	12.【基于YOLOv8深度学习的120种犬类检测与识别系统】
13.【基于YOLOv8深度学习的路面坑洞检测系统】	14.【基于YOLOv8深度学习的火焰烟雾检测系统】
15.【基于YOLOv8深度学习的钢材表面缺陷检测系统】	16.【基于YOLOv8深度学习的舰船目标分类检测系统】
17.【基于YOLOv8深度学习的西红柿成熟度检测系统】	18.【基于YOLOv8深度学习的血细胞检测与计数系统】
19.【基于YOLOv8深度学习的吸烟/抽烟行为检测系统】	20.【基于YOLOv8深度学习的水稻害虫检测与识别系统】
21.【基于YOLOv8深度学习的高精度车辆行人检测与计数系统】	22.【基于YOLOv8深度学习的路面标志线检测与识别系统】
23.【基于YOLOv8深度学习的智能小麦害虫检测识别系统】	24.【基于YOLOv8深度学习的智能玉米害虫检测识别系统】
25.【基于YOLOv8深度学习的200种鸟类智能检测与识别系统】	26.【基于YOLOv8深度学习的45种交通标志智能检测与识别系统】
27.【基于YOLOv8深度学习的人脸面部表情识别系统】	28.【基于YOLOv8深度学习的苹果叶片病害智能诊断系统】
29.【基于YOLOv8深度学习的智能肺炎诊断系统】	30.【基于YOLOv8深度学习的葡萄簇目标检测系统】
31.【基于YOLOv8深度学习的100种中草药智能识别系统】	32.【基于YOLOv8深度学习的102种花卉智能识别系统】
33.【基于YOLOv8深度学习的100种蝴蝶智能识别系统】	34.【基于YOLOv8深度学习的水稻叶片病害智能诊断系统】
35.【基于YOLOv8与ByteTrack的车辆行人多目标检测与追踪系统】	36.【基于YOLOv8深度学习的智能草莓病害检测与分割系统】
37.【基于YOLOv8深度学习的复杂场景下船舶目标检测系统】	38.【基于YOLOv8深度学习的农作物幼苗与杂草检测系统】
39.【基于YOLOv8深度学习的智能道路裂缝检测与分析系统】	40.【基于YOLOv8深度学习的葡萄病害智能诊断与防治系统】
41.【基于YOLOv8深度学习的遥感地理空间物体检测系统】	42.【基于YOLOv8深度学习的无人机视角地面物体检测系统】
43.【基于YOLOv8深度学习的木薯病害智能诊断与防治系统】	44.【基于YOLOv8深度学习的野外火焰烟雾检测系统】
45.【基于YOLOv8深度学习的脑肿瘤智能检测系统】	46.【基于YOLOv8深度学习的玉米叶片病害智能诊断与防治系统】
47.【基于YOLOv8深度学习的橙子病害智能诊断与防治系统】	48.【车辆检测追踪与流量计数系统】
49.【行人检测追踪与双向流量计数系统】	50.【基于YOLOv8深度学习的反光衣检测与预警系统】
51.【危险区域人员闯入检测与报警系统】	52.【高压输电线绝缘子缺陷智能检测系统】

二、机器学习实战专栏【链接】，已更新31期，欢迎关注，持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】，持续更新中~~
六、YOLO性能对比专栏【链接】，持续更新中~

《------正文------》

引言

本文主要介绍一个使用YOLOv8 + MobileSAM自动生成结构化分割数据标注的方法。这个脚本是一个完整的从图像检测到分割再到生成标注文件的工作流，利用YOLOv8进行快速物体检测，再借助SAM模型实现高精度的分割，并自动化生成结构化的标注数据，非常适合于大规模图像数据的分割数据集标注。

实现步骤

下面是详细的实现步骤。

导入所需库和模块

from pathlib import Path
from ultralytics import SAM, YOLO
import torch

pathlib.Path：用于处理文件和目录路径的库。
ultralytics.YOLO 和 ultralytics.SAM：分别用于加载YOLOv8对象检测模型和Segment Anything Model (SAM)语义分割模型。
torch：PyTorch深度学习框架。

定义变量

# 定义图像数据路径
img_data_path = 'ultralytics/assets'
# 定义检测模型和SAM模型的路径
det_model="yolov8n.pt"
sam_model="MyModels/mobile_sam.pt"
# 根据CUDA是否可用选择设备
device = '0' if torch.cuda.is_available() else 'cpu'
# 定义输出目录，默认为None
# 输出路径
output_dir = None

img_data_path：需要进行自动分割标注的图像数据路径。
det_model 和 sam_model：分别为YOLO检测模型和SAM分割模型的文件路径。
device：根据系统中是否可用CUDA（NVIDIA GPU加速），选择运行模型的设备（GPU或CPU）。

模型初始化

# 初始化检测模型和SAM模型
det_model = YOLO(det_model)
sam_model = SAM(sam_model)

使用给定路径加载YOLOv8检测模型和SAM分割模型到内存中，准备进行图像处理。
这里使用的是yolov8n.pt目标检测和mobile_sam.pt分割一切模型，如果对于精度要求更高的，可以使用yolov8x.pt或者sam_b.pt等更大的模型。

图像检测

# 对图像数据进行检测
det_results = det_model(data, stream=True, device=device)

使用YOLO模型对img_data_path路径下的所有图像进行物体检测，stream=True允许逐个处理图像，减少内存占用；device=device指定计算设备。

结果处理与分割

# 遍历检测结果
for result in det_results:
    # 获取类别ID
    class_ids = result.boxes.cls.int().tolist()  # noqa
    # 如果有检测到物体
    if len(class_ids):
        # 获取检测框坐标
        boxes = result.boxes.xyxy  # Boxes object for bbox outputs
        # 使用SAM模型进行分割
        sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)

遍历检测结果：对于每个检测到物体的图像，提取其类别ID和边界框坐标。
应用SAM模型：对每个检测到的物体使用SAM模型进行精细化分割，获取更精确的分割掩码。

生成标注文件

# 获取分割结果
        segments = sam_results[0].masks.xyn  # noqa
        # 为每个图像生成标注文件
        with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f:
            # 遍历每个分割区域
            for i in range(len(segments)):
                s = segments[i]
                # 如果分割区域为空，则跳过
                if len(s) == 0:
                    continue
                # 将分割区域坐标转换为字符串格式
                segment = map(str, segments[i].reshape(-1).tolist())
                # 写入标注信息
                f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")

对于每张图像及其检测到的每个物体，创建一个文本文件用于存储标注信息。
文件命名规则为原图像文件名加上.txt后缀，存储于之前创建的output_dir目录下。

标注格式：每行代表一个物体的标注信息，包括类别ID后跟分割区域的坐标点序列，坐标点以空格分隔。

生成的标注文件为YOLO分割数据格式，如下图：

分割图片结果如下所示：

完整源码

from pathlib import Path
from ultralytics import SAM, YOLO
import torch

# 定义图像数据路径
img_data_path = 'ultralytics/assets'

# 定义检测模型和SAM模型的路径
det_model="yolov8n.pt"
sam_model="MyModels/mobile_sam.pt"

# 根据CUDA是否可用选择设备
device = '0' if torch.cuda.is_available() else 'cpu'

# 定义输出目录，默认为None
# 输出路径
output_dir = None

# 初始化检测模型和SAM模型
det_model = YOLO(det_model)
sam_model = SAM(sam_model)

# 获取图像数据路径
data = Path(img_data_path)

# 如果输出目录未定义，则生成默认的输出目录
if not output_dir:
    output_dir = data.parent / f"{data.stem}_auto_annotate_labels"
    # 创建输出目录
    Path(output_dir).mkdir(exist_ok=True, parents=True)

# 对图像数据进行检测
det_results = det_model(data, stream=True, device=device)

# 遍历检测结果
for result in det_results:
    # 获取类别ID
    class_ids = result.boxes.cls.int().tolist()  # noqa
    # 如果有检测到物体
    if len(class_ids):
        # 获取检测框坐标
        boxes = result.boxes.xyxy  # Boxes object for bbox outputs
        # 使用SAM模型进行分割
        sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
        # 获取分割结果
        segments = sam_results[0].masks.xyn  # noqa
        # 为每个图像生成标注文件
        with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f:
            # 遍历每个分割区域
            for i in range(len(segments)):
                s = segments[i]
                # 如果分割区域为空，则跳过
                if len(s) == 0:
                    continue
                # 将分割区域坐标转换为字符串格式
                segment = map(str, segments[i].reshape(-1).tolist())
                # 写入标注信息
                f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")

好了，这篇文章就介绍到这里，感谢点赞关注，更多精彩内容持续更新中~

关注文末名片G-Z-H：【阿旭算法与机器学习】，可获取更多干货学习资源