【目标检测】用YOLOv8-Segment训练语义分割数据集（保姆级教学）

2024-11-17 10:31:18

前言

这篇教程会手把手带你用 YOLOv8-Segment 搭建一个属于自己的分割任务项目。从环境配置到数据集准备，再到模型训练和测试，所有步骤都有详细说明，适合初学者使用。你将学会如何安装必要的软件，标注自己的数据，并使用 YOLOv8-Segment 实现自定义目标的实例分割。我们还提供了实用的代码示例，方便你在这个基础上*调整和优化。

准备工作（必看）

安装 Anaconda 并熟悉基本命令
本项目主要基于 Conda 环境进行配置与管理。如果尚未安装 Anaconda，可以参考我的另一篇文章：一步步教你在 Windows 上轻松安装 Anaconda以及使用常用conda命令（超详细）获取详细安装及使用指南。
安装 PyCharm（版本 > 2023.1）
本项目建议使用 PyCharm 加载 Conda 环境进行开发与管理，可以参考相关教程完成 PyCharm 的安装和基础配置。
熟练使用 LabelMe 进行分割标注
本项目要求用户预先准备好 JPG 格式的数据集，并使用 labelme 工具完成目标区域的分割标注。分割标注将生成 JSON 文件，稍后会介绍如何将这些文件转换为 YOLOv8 可用的格式。如果您对 labelme 的使用不熟悉，可以参考我的另一篇教程：小白快速上手 labelme：新手图像标注详解教程。

一、下载 YOLOv8 源码

方法一：从官网直接下载（推荐）

打开 YOLOv8 的 GitHub 官方页面。
点击页面上的 “Code” 按钮，选择 “Download ZIP” 直接下载源码。

下载完成后，请将压缩包解压至 D 盘或 E 盘等非系统盘，并重命名为 yolov8-segment，以减少可能由于系统权限问题导致的错误（解压路径可根据个人需求进行选择，请确保记录路径以便后续使用）。

方法二：通过 Git 克隆源码（适合有一定基础的用户）

打开终端或命令提示符，输入以下命令将 YOLOv8 源码克隆到本地：
```
git clone https://github.com/ultralytics/ultralytics
```
下载完成后，将文件夹重命名为 yolov8-segment：
```
mv ultralytics yolov8-segment
```
进入 yolov8-segment 文件夹：
```
cd yolov8-segment
```

文件结构简介

下载完成后，ultralytics 文件夹的核心结构如下：

ultralytics/
├── ultralytics/
│   ├── yolo/        # YOLOv8 的核心代码，包括检测、分割和分类功能
│   ├── datasets/    # 示例数据集
│   ├── models/      # 模型定义
│   ├── utils/       # 工具函数
├── setup.py         # 安装脚本
└── README.md        # 项目说明
└── ......           # 其它

二、环境安装

1. 创建并激活 YOLOv8 环境

conda create -n yolov8 python==3.9  # 创建python版本3.9名为yolov8的conda环境

conda activate yolov8  # 激活创建的yolov8环境

2. 配置Pytorch环境

YOLOv8 可以在GPU或CPU环境下运行，但推荐在GPU上运行以加快训练速度。

检查NVIDIA GPU及其CUDA支持情况

nvidia-smi是NVIDIA驱动自带的命令行工具，可以查看显卡的CUDA支持情况和驱动信息。打开命令提示符（或CMD），输入：

nvidia-smi

如果系统输出了显卡信息，并显示CUDA版本（例如CUDA Version: 12.7），则表明显卡支持CUDA同时支持GPU，否则~~不支持GPU~~ 。

方式一：使用Conda安装（推荐）

打开 Previous PyTorch Versions | PyTorch 官网，选择<= CUDA Version 的Conda命令并复制（若不支持GPU则选择 CPU 的Conda命令）。
在终端中执行对应的安装命令，以安装PyTorch和CUDA / CPU支持。
安装完成后，可以在终端中输入conda list pytorch检查安装情况：(本教程仅展示CUDA版本的安装结果)

pytorch 2.3.1：这是安装的PyTorch版本，适用于Python 3.9，并支持CUDA 12.1和cuDNN 8。
pytorch-cuda 12.1：这是安装的CUDA支持包，指明CUDA版本为12.1。
pytorch-mutex 1.0：这是一个互斥包，用来确保在同一环境中只启用一种加速方式（如CUDA），避免安装冲突。

方式二：使用Pip安装

打开 Previous PyTorch Versions | PyTorch 官网，选择<= CUDA Version 的Pip命令并复制（若不支持GPU则选择 CPU 的Pip命令）。

在终端中执行对应的安装命令，以安装PyTorch和CUDA / CPU支持。
安装完成后，可以在终端中输入pip list检查安装情况。

torch 2.3.0：核心库（版本2.3.0），用于构建和训练神经网络。
torchaudio 2.3.0：音频处理库（版本2.3.0），为音频数据的加载、预处理和增强提供了工具
torchvision 0.18.0：计算机视觉库（版本0.18.0），包含了常用的图像数据集、数据增强和预训练模型。

下载安装CUDA支持的torch + torchvision + torchaudio (仅GPU版本需要)。
- 打开网址 https://download.pytorch.org/whl/torch_stable.html
- 选择与上面对应版本的torch并下载至本地

cu121：表示此安装包支持 CUDA 12.1，也就是说，这个版本的 PyTorch 可以利用带有 CUDA 12.1 的 NVIDIA GPU 进行计算加速。
	
torch-2.3.1+cu121：说明这是 PyTorch 2.3.1 版本的安装包，并且这个版本支持 CUDA 12.1。在 PyTorch 中，不同 CUDA 版本对应不同的 GPU 支持，确保安装的 PyTorch 版本兼容设备的 CUDA 版本很重要。
	
cp39-cp39：表示该安装包是针对 Python 3.9（“cp39”代表 CPython 3.9）的版本。这里的 cp39-cp39 意味着此包适用于 Python 3.9 的 CPython 解释器（CPython 是 Python 的标准实现）。
	
win_amd64：表示该安装包适用于 Windows 操作系统，并支持 64位架构（AMD64）。如果您的系统是 Windows 64 位，这个包是兼容的。

在终端中输入以下命令以安装指定的 .whl 文件（请将路径替换为您下载的 .whl 文件的实际保存位置）

  pip install D:\torch_gpu\torch-2.3.1+cu121-cp39-cp39-win_amd64.whl

torchvision与torchaudio安装方式同理
安装完毕后再次在终端中输入pip list检查安装情况，如果输出结果中显示类似以下内容，则表明安装成功

3. Pycharm中加载YOLOv8环境

在Pycharm中打开yolov8-segment项目
打开设置 - - > 项目: yolov8-segment- - > python解释器 - - > 添加解释器 - - > 添加本地解释器

选择Conda环境 - - > 浏览Anaconda目录下\Library\bin\conda.bat - - > 加载环境 - - > 使用现有环境 - - > 选择yolov8 - - > 点击确定

4. 安装YOLOv8依赖项

打开Pycharm本地终端

使用以下命令安装 Ultralytics 库和相关依赖：

pip install ultralytics opencv-python matplotlib

三、数据集准备

3.1 数据集初始格式

为了准备 YOLOv8 Segment 模型所需的数据集，请按照以下步骤组织文件夹结构：

在项目根目录 yolov8-segment 下新建 datasets 文件夹。
在 datasets 文件夹内新建一个名为 segment 的子文件夹。
在 segment 文件夹内，新建以下两个子文件夹：
- img 文件夹：用于存放任务所需的图片，格式建议为 JPG 或 PNG。
- json 文件夹：用于存放与图片对应的标注文件，格式为 JSON（推荐使用工具如 Labelme 生成）。

文件夹结构示例

yolov8-segment/
├── datasets/
│   ├── segment/
│       ├── img/        # 存放图片
│       │   ├── image1.jpg
│       │   ├── image2.jpg
│       ├── json/       # 存放标注文件
│           ├── image1.json
│           ├── image2.json

注意事项

图片文件夹 img：包含任务图片，每张图片的文件名应唯一且对应相应的标注文件，例如 image1.jpg。
标注文件夹 json：包含与图片一一对应的 JSON 标注文件，文件名与图片名一致，例如 image1.json。
一致性：确保图片和 JSON 文件的文件名完全一致，以免后续处理出错。

3.2 数据集转换和分割

为了使用 YOLOv8 Segment 模型，我们需要将 JSON 格式的标注文件转换为 YOLO格式的 .txt 文件，并按照训练和验证的比例划分数据集。以下是具体步骤：

1. 新建转换脚本

在项目根目录下新建一个脚本文件 segment_json2txt.py。这个脚本的主要功能是：

格式转换：将标注工具（如 Labelme）生成的 JSON 格式标注文件转换为 YOLO 格式的 .txt 文件，方便 YOLOv8 模型直接使用。
数据集划分：将图片和对应的标签文件按照一定比例（如 90% 训练集，10% 验证集）划分为训练集和验证集。
生成配置文件：自动生成一个名为 segment.yaml 的配置文件，该文件包含数据集路径、类别名称和类别数量，供 YOLOv8 模型训练时使用。

# -*- coding: utf-8 -*-
from tqdm import tqdm
import shutil
import random
import os
import argparse
from collections import Counter
import yaml
import json
 
 

def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)
 
 
def convert_label_json(json_dir, save_dir, classes):
    json_paths = os.listdir(json_dir)
    classes = classes.split(',')
    mkdir(save_dir)
 
    for json_path in tqdm(json_paths):
        # for json_path in json_paths:
        path = os.path.join(json_dir, json_path)
        with open(path, 'r') as load_f:
            json_dict = json.load(load_f)
        h, w = json_dict['imageHeight'], json_dict['imageWidth']
 
        # save txt path
        txt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))
        txt_file = open(txt_path, 'w')
 
        for shape_dict in json_dict['shapes']:
            label = shape_dict['label']
            label_index = classes.index(label)
            points = shape_dict['points']
 
            points_nor_list = []
 
            for point in points:
                points_nor_list.append(point[0] / w)
                points_nor_list.append(point[1] / h)
 
            points_nor_list = list(map(lambda x: str(x), points_nor_list))
            points_nor_str = ' '.join(points_nor_list)
 
            label_str = str(label_index) + ' ' + points_nor_str + '\n'
            txt_file.writelines(label_str)
 
 
def get_classes(json_dir):
    '''
    统计路径下 JSON 文件里的各类别标签数量
    '''
    names = []
    json_files = [os.path.join(json_dir, f) for f in os.listdir(json_dir) if f.endswith('.json')]
 
    for json_path in json_files:
        with open(json_path, 'r') as f:
            data = json.load(f)
            for shape in data['shapes']:
                name = shape['label']
                names.append(name)
 
    result = Counter(names)
    return result
 
 
def main(image_dir, json_dir, txt_dir, save_dir):
    # 创建文件夹
    mkdir(save_dir)
    images_dir = os.path.join(save_dir, 'images')
    labels_dir = os.path.join(save_dir, 'labels')
 
    img_train_path = os.path.join(images_dir, 'train')
    img_val_path = os.path.join(images_dir, 'val')
 
    label_train_path = os.path.join(labels_dir, 'train')
    label_val_path = os.path.join(labels_dir, 'val')
 
    mkdir(images_dir)
    mkdir(labels_dir)
    mkdir(img_train_path)
    mkdir(img_val_path)
    mkdir(label_train_path)
    mkdir(label_val_path)
 
    # 数据集划分比例，训练集75%，验证集15%，测试集15%，按需修改
    train_percent = 0.90
    val_percent = 0.10
 
    total_txt = os.listdir(txt_dir)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)  # 范围 range(0, num)
 
    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
 
    train = random.sample(list_all_txt, num_train)
    # 在全部数据集中取出train
    val = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    # val = random.sample(list_all_txt, num_val)
 
    print("训练集数目：{}, 验证集数目：{}".format(len(train), len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]
 
        srcImage = os.path.join(image_dir, name + '.jpg')
        srcLabel = os.path.join(txt_dir, name + '.txt')
 
        if i in train:
            dst_train_Image = os.path.join(img_train_path, name + '.jpg')
            dst_train_Label = os.path.join(label_train_path, name + '.txt')
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
        elif i in val:
            dst_val_Image = os.path.join(img_val_path, name + '.jpg')
            dst_val_Label = os.path.join(label_val_path, name + '.txt')
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
 
    obj_classes = get_classes(json_dir)
    classes = list(obj_classes.keys())
 
    # 编写yaml文件
    classes_txt = {i: classes[i] for i in range(len(classes))}  # 标签类别
    data = {
        'path': os.path.join(os.getcwd(), save_dir),
        'train': "images/train",
        'val': "images/val",
        'names': classes_txt,
        'nc': len(classes)
    }
    with open(save_dir + '/segment.yaml', 'w', encoding="utf-8") as file:
        yaml.dump(data, file, allow_unicode=True)
    print("标签：", dict(obj_classes))
 
 
if __name__ == "__main__":
    """
    python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_rings/txts --classes "cat,dogs"
    """
    classes_list = 'platemetal,label,Dust_plug,screws'  # 类名
 
    parser = argparse.ArgumentParser(description='json convert to txt params')
    parser.add_argument('--image-dir', type=str,default='datasets/segment/img', help='图片地址')
    parser.add_argument('--json-dir', type=str,default='datasets/segment/json', help='json地址')
    parser.add_argument('--txt-dir', type=str,default='datasets/segment/txt' , help='保存txt文件地址')
    parser.add_argument('--save-dir', default='datasets/segment/seg', type=str, help='保存最终分割好的数据集地址')
    parser.add_argument('--classes', type=str, default=classes_list, help='classes')
    args = parser.parse_args()
    json_dir = args.json_dir
    txt_dir = args.txt_dir
    image_dir = args.image_dir
    save_dir = args.save_dir
    classes = args.classes
    # json格式转txt格式
    convert_label_json(json_dir, txt_dir, classes)
    # 划分数据集，生成yaml训练文件
    main(image_dir, json_dir, txt_dir, save_dir)

2. 修改路径参数

数据集划分比例（`segment_json2txt.py`第90行左右）

# 数据集划分比例
train_percent = 0.90  # 训练集占比
val_percent = 0.10    # 验证集占比

默认设置为训练集 90%，验证集 10%，没有单独测试集。如果需要，可以按需修改，例如训练集 75%，验证集 15%，测试集 10%。

参数说明及调整

在运行脚本时，需要指定以下参数，根据你的数据集路径和类别定义进行调整：

--image-dir：图片文件夹路径，例如 datasets/segment/img。
--json-dir：JSON 标注文件夹路径，例如 datasets/segment/json。
--txt-dir：生成的 TXT 文件存放路径，例如 datasets/segment/txt。
--save-dir：最终分割数据集存放路径
--classes：分割任务中的类别名称（必改）

3. 运行脚本

运行脚本后（终端输出类似上图），会生成两个文件夹：

txt：包含转换为 YOLO 格式的标签文件。
seg：分割好的最终数据集，包括 train 和 val 子文件夹，以及一个用于训练的 segment.yaml 配置文件。

以下是运行脚本后生成的示例目录结构：

datasets/
├── segment/
│   ├── txt/                    # 包含转换为 YOLO 格式的标签文件
│   │   ├── image1.txt
│   │   ├── image2.txt
│   │   ├── image3.txt
│   │   └── ...                # 其他标签文件
│   ├── seg/                   # 分割好的最终数据集
│       ├── images/            # 存放划分后的图像
│       │   ├── train/         # 训练集图片
│       │   │   ├── image1.jpg
│       │   │   ├── image2.jpg
│       │   │   ├── image3.jpg
│       │   │   └── ...       # 更多训练图片
│       │   ├── val/           # 验证集图片
│       │       ├── image4.jpg
│       │       ├── image5.jpg
│       │       ├── image6.jpg
│       │       └── ...       # 更多验证图片
│       ├── labels/            # 存放划分后的标签
│           ├── train/         # 训练集标签
│           │   ├── image1.txt
│           │   ├── image2.txt
│           │   ├── image3.txt
│           │   └── ...       # 更多训练标签
│           ├── val/           # 验证集标签
│               ├── image4.txt
│               ├── image5.txt
│               ├── image6.txt
│               └── ...       # 更多验证标签
│       ├── segment.yaml       # 用于 YOLOv8 训练的配置文件

3.3 下载预训练权重

3.3.1 YOLOv8-Segment 预训练权重（已停用）

在 yolov8-segment 根目录下创建一个文件夹 weights
打开开源工具包/YOLOv8分割预训练权重，点击"下载zip"

解压其目录下的"YOLOv8分割预训练权重.zip"内容至weights文件夹

3.3.2 YOLOv11-Segment 预训练权重

在 yolov8-egment 根目录下创建一个文件夹 weights ，打开 Ultralytics 官方文档中关于 YOLOv11 分割任务的说明页面，下滑找到预训练权重，点击下载至 weights 文件夹（下载预训练权重）

以下是weights的示例目录结构：

yolov8-segment/
├── weights/
	├── yolo11n-segment.pt
	├── yolo11s-segment.pt
	├── yolo11m-segment.pt
	├── ...

四、训练模型

训练命令

Pycharm本地终端使用以下命令启动训练：

yolo train task=segment data=datasets/segment/seg/segment.yaml model=weights/yolo11s-seg.pt epochs=50 imgsz=640 batch=8 workers=0 device=0

该命令表示使用 YOLOv8 的分割任务 (segment)，加载自定义权重 yolo11s-seg.pt，基于配置文件 segment.yaml 中的数据集，训练 50 个轮次，每次训练 8 张图片，输入图像尺寸为 640，使用第 0 块 GPU 进行训练。
提示：您可以根据设备配置或训练需求，灵活调整如 epochs（训练轮次）、batch（批次大小）和 imgsz（图像尺寸）等参数，以优化训练效率和性能。

参数	含义	示例值
`train`	操作模式，指定当前命令为训练模式。	必须明确指定 `train`
`task`	指定任务类型。	`segment`（实例分割）、`detect`（目标检测）
`data`	数据集配置文件路径，定义训练集、验证集路径及类别信息。	`datasets/segment/seg/segment.yaml`
`model`	预训练权重文件路径，用于微调训练。	`weights/yolo11s-seg.pt`、`yolov8n-seg.pt`
`epochs`	训练的总轮数，表示数据集被遍历的次数。	`50`（快速训练）、`100`（标准训练）
`imgsz`	输入图像的尺寸（宽和高），所有图片会被缩放到这个大小进行训练。	`640`（默认值，推荐）、`512`（较小尺寸）
`batch`	每次训练的图片数量（批次大小），越大需要的显存越多。	`8`（低显存设备）、`16`（推荐值）、`32`
`workers`	数据加载线程数，`0` 表示使用主线程，适合 Windows 系统或资源受限的设备。	`0`（默认值）、`4`（推荐）、`8`（多线程加载）
`device`	指定训练设备。可以使用 GPU 或 CPU，数字代表 GPU 编号。	`0`（第 0 块 GPU）、`1`（第 1 块 GPU）、`cpu` （使用CPU）

运行结果

训练完成后，模型权重和日志文件将保存在默认的 runs/segment/train 文件夹中，可以通过以下方式查看：

权重文件：runs/segment/train/weights/ 下包含最新的训练权重（last.pt）和最佳权重（best.pt）。
训练日志：包含训练精度、损失值等信息。

五、验证测试

1、推理命令

使用 YOLOv8-Segment 进行实例分割推理测试：

yolo predict task=segment model=runs/segment/train/weights/best.pt source=/path/to/image.jpg conf=0.25 show=True

task=segment：指定任务类型为实例分割。
model：训练生成的模型权重路径。
source：指定推理的输入源，支持以下几种格式：
- 单张图片：/path/to/image.jpg
- 文件夹（多张图片）：/path/to/images/
- 视频文件：/path/to/video.mp4
- 摄像头：使用 0 表示默认摄像头。
- 网络流：支持 RTSP/HTTP URL，例如 rtsp://username:password@ip_address:port
conf：置信度阈值，默认为 0.25，设置更高的值可减少低置信度结果。
iou：IOU 阈值，默认为 0.45，调整后可以控制重叠目标的处理。
save：是否保存预测结果，默认开启。
show：是否直接显示推理结果。
device：指定运行设备
- GPU：0（第 0 块 GPU）、1（第 1 块 GPU）。
- CPU："cpu"。

2、自定义脚本

在 yolov8-segment 根目录下创建 segment_detect.py，内容如下：

from ultralytics import YOLO

# 加载训练好的模型
model = YOLO("runs/segment/train/weights/best.pt")  # 替换为实际的权重路径

# 执行推理
model.predict(
    task="segment",                # 指定任务类型为实例分割
    source="path/to/image.jpg",    # 输入源：图片、视频、文件夹或摄像头（如 0 表示摄像头）
    conf=0.25,                     # 置信度阈值，过滤低置信度目标
    iou=0.45,                      # IOU 阈值，控制目标框的重叠过滤
    save=True,                     # 是否保存预测结果，默认保存到指定目录
    save_txt=False,                # 是否保存预测结果为文本
    save_conf=False,               # 是否保存预测框的置信度值
    show=True,                     # 是否实时显示预测结果
    device=0                       # 使用设备（0 表示第 0 块 GPU，或者 'cpu' 表示使用 CPU）
)

运行后，分割结果会保存到 runs/segment/predict 目录中，结果图片中将包含分割区域的掩码和轮廓，便于观察分割效果。

码农公寓

前言