语义分割-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)
相关推荐
Elastic 中国社区官方博客1 天前
Elasticsearch:使用 Agent Builder 的 A2A 实现 - 开发者的圣诞颂歌
大数据·数据库·人工智能·elasticsearch·搜索引擎·ai·全文检索
chools1 天前
【AI超级智能体】快速搞懂工具调用Tool Calling 和 MCP协议
java·人工智能·学习·ai
郝学胜-神的一滴1 天前
深度学习必学:PyTorch 神经网络参数初始化全攻略(原理 + 代码 + 选择指南)
人工智能·pytorch·python·深度学习·神经网络·机器学习
leobertlan1 天前
好玩系列:用20元实现快乐保存器
android·人工智能·算法
笨笨饿1 天前
#58_万能函数的构造方法:ReLU函数
数据结构·人工智能·stm32·单片机·硬件工程·学习方法
jr-create(•̀⌄•́)1 天前
从零开始:手动实现神经网络识别手写数字(完整代码讲解)
人工智能·深度学习·神经网络
冬奇Lab1 天前
一天一个开源项目(第78篇):MiroFish - 用群体智能引擎预测未来
人工智能·开源·资讯
冬奇Lab1 天前
你的 Skill 真的好用吗?来自OpenAI的 Eval 系统化验证 Agent 技能方法论
人工智能·openai
数智工坊1 天前
Transformer 全套逻辑:公式推导 + 原理解剖 + 逐行精读 - 划时代封神之作!
人工智能·深度学习·transformer
GreenTea1 天前
AI 时代,工程师的不可替代性在哪里
前端·人工智能·后端