语义分割-yolo26seg

写在前面的话：

在复现过程中主要参考下边这篇文章，在此表示感谢！语义分割：YOLO26的分割模型训练自己的数据集（从代码下载到实例测试）

一、程序下载

官方程序地址如下：ultralytics:打造前沿SOTA YOLO模型，支持检测、分割、分类与姿态估计，快速精准易用 - AtomGit | GitCode

下载程序界面向下滑，下载预训练模型，下载后直接放ultralytics-main路径下

二、环境配置

1.创建虚拟环境

conda create -n yolo26seg python=3.11 -y

2.激活虚拟环境

conda activate yolo26seg

3.安装PyTorch

python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

4.安装 YOLOv26（ultralytics）

python -m pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

三、数据划分

在ultralytics-main目录下建文件夹datasets\images和datasets\labels，将图片数据放在images里，标签放在labels里，运行下边程序进行数据划分。

# coding: utf8
"""
数据存放在images文件夹下，标签存放在labels文件夹下，运行程序后目录结构如下：
datasets/
├── images/          # 原始文件已清空
│   ├── train/
│   ├── val/
│   └── test/
├── labels/          # 原始文件已清空
│   ├── train/
│   ├── val/
│   └── test/
├── train.txt
├── val.txt
└── test.txt
"""

import os
import shutil
import random

# ====================== 【只需修改这里】 ======================
DATASET_ROOT = r"C:\Users\ultralytics-main\datasets"  # 你的数据集根目录
IMAGE_DIR = "images"  # 原始图片文件夹
LABEL_DIR = "labels"  # 原始标签文件夹
SPLIT_RATIO = [0.8, 0.1, 0.1]  # 训练/验证/测试比例


# ==============================================================

def create_folders(root):
    """自动创建 train/val/test 文件夹"""
    subsets = ['train', 'val', 'test']
    for sub in subsets:
        os.makedirs(os.path.join(root, IMAGE_DIR, sub), exist_ok=True)
        os.makedirs(os.path.join(root, LABEL_DIR, sub), exist_ok=True)


def get_paired_files(img_dir, lab_dir):
    """自动配对图片和任意格式标签"""
    # 支持的图片格式
    img_exts = {'jpg', 'jpeg', 'png', 'bmp', 'tiff', 'tif', 'JPG', 'PNG'}
    # 支持的标签格式（txt/json/png/jpg全支持）
    lab_exts = {'png', 'jpg', 'jpeg', 'txt', 'json', 'bmp', 'PNG', 'JPG'}

    img_map = {}
    for img_file in os.listdir(img_dir):
        ext = img_file.split('.')[-1].lower()
        if ext in img_exts:
            name = os.path.splitext(img_file)[0]
            img_map[name] = img_file

    lab_map = {}
    for lab_file in os.listdir(lab_dir):
        ext = lab_file.split('.')[-1].lower()
        if ext in lab_exts:
            name = os.path.splitext(lab_file)[0]
            lab_map[name] = lab_file

    pairs = []
    for name in img_map:
        if name in lab_map:
            img_path = os.path.join(img_dir, img_map[name])
            lab_path = os.path.join(lab_dir, lab_map[name])
            pairs.append((img_path, lab_path))

    print(f"✅ 成功配对图片+标签：{len(pairs)} 组")
    if len(pairs) == 0:
        raise Exception("❌ 没有匹配到图片和标签，请检查文件名！")

    return pairs


def split_dataset(pairs, root, ratio):
    """划分数据集 + 剪切移动 + 生成txt"""
    random.seed(42)
    random.shuffle(pairs)
    total = len(pairs)

    train_num = int(total * ratio[0])
    val_num = int(total * ratio[1])

    splits = [
        ('train', 0, train_num),
        ('val', train_num, train_num + val_num),
        ('test', train_num + val_num, total)
    ]

    for subset, start, end in splits:
        count = end - start
        print(f"\n📂 生成 {subset} 集：{count} 张")

        img_sub_folder = os.path.join(root, IMAGE_DIR, subset)
        lab_sub_folder = os.path.join(root, LABEL_DIR, subset)
        txt_save_path = os.path.join(root, f"{subset}.txt")

        with open(txt_save_path, 'w', encoding='utf-8') as f:
            for i in range(start, end):
                img_src, lab_src = pairs[i]
                img_name = os.path.basename(img_src)
                lab_name = os.path.basename(lab_src)

                shutil.move(img_src, os.path.join(img_sub_folder, img_name))
                shutil.move(lab_src, os.path.join(lab_sub_folder, lab_name))

                # 写入txt
                img_line = f"{IMAGE_DIR}/{subset}/{img_name}".replace('\\', '/')
                lab_line = f"{LABEL_DIR}/{subset}/{lab_name}".replace('\\', '/')
                f.write(f"{img_line} {lab_line}\n")


if __name__ == '__main__':
    img_dir = os.path.join(DATASET_ROOT, IMAGE_DIR)
    lab_dir = os.path.join(DATASET_ROOT, LABEL_DIR)

    print("🚀 开始划分数据集...")
    create_folders(DATASET_ROOT)
    paired = get_paired_files(img_dir, lab_dir)
    split_dataset(paired, DATASET_ROOT, SPLIT_RATIO)

    print("\n🎉 全部完成！")

四、配置文件

在ultralytics-main路径下创建data.yaml，修改数据集路径，分割的类别及数量。

path: ../datasets  # 数据集所在路径
train: train.txt  # 数据集路径下的train.txt
val: val.txt  # 数据集路径下的val.txt
test: test.txt  # 数据集路径下的test.txt

is_segment: true

# 类别数量（必须和下面 names 数量一致）
nc: 5

# Classes
names:
  0: 汽车
  1: 货车
  2: 非机动车
  3: 道路
  4: 行人
  5: 护栏

五、模型训练、验证、推理

1.模型训练

新建ultralytics-main\train.py，程序如下修改相关参数即可，训练结果路径为：ultralytics-main\runs\segment\train

from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(r'yolo26-seg.yaml')  # 默认n m只需写yolo26m即可定位到m模型
    model.load('yolo26n-seg.pt') # 是否加载预训练权重
    model.train(data=r'data.yaml',
                task='segment',  # 任务类别
                imgsz=640,
                epochs=300,
                single_cls=False,  # 多类别设置False
                batch=8,
                workers=0,
                device='0',
                amp=True,
                optimizer="AdamW"
                )

2.模型验证

新建ultralytics-main\val.py，程序如下修改相关参数即可。

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('runs/segment/train/weights/best.pt')
    model.val(data='data.yaml',
              task='segment',
              imgsz=640,
              batch=16,
              split='test',   #使用测试集评估
              workers=10,
              device='0',
              )

3.推理

新建ultralytics-main\detect.py，程序如下修改相关参数即可，推理结果路径为runs\segment\predict。

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('runs/segment/train/weights/best.pt')
    model.predict(
        source='datasets/images/test',
        imgsz=512,
        device='0',
        save=True,
        boxes=False  # 关闭检测框绘制
    )

六、模型导出

新建ultralytics-main\export.py，将模型导出为onnx格式，程序如下修改相关参数即可， 运行程序后在runs\segment\train\weights里生成best.onnx。

from ultralytics import YOLO

if __name__ == "__main__":

    pth_path = r"runs/segment/train/weights/best.pt"
    model = YOLO(pth_path)  # 加载自定义训练模型

    # Export the model
    model.export(format='onnx', opset=12)