语义分割-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)
相关推荐
测试小张z2 小时前
AI智能客服测试点
人工智能·ai
昨夜见军贴06162 小时前
AI审核驱动效率跃迁:IACheck如何让环境数据审核提速5倍并覆盖90%常规问题
人工智能
昨夜见军贴06162 小时前
AI审核驱动路径优化:IACheck智能决策树重塑环境监测数据矛盾点的人机协同处置体系
人工智能
AI街潜水的八角2 小时前
YOLO26手语识别项目实战1-三十五种手语实时检测系统数据集说明(含下载链接)
python·深度学习
春日见2 小时前
从底层思维3分钟彻底弄清卷积神经网络CNN
人工智能·深度学习·神经网络·计算机视觉·docker·cnn·计算机外设
AC赳赳老秦2 小时前
自媒体博主:OpenClaw多Agent协同,实现选题-创作-审核全流程自动化
运维·服务器·开发语言·人工智能·自动化·媒体·openclaw
报错小能手2 小时前
初识transformer《Attention Is All You Need》论文解析
人工智能·深度学习·transformer
福客AI智能客服2 小时前
智能客服平台:AI人工智能客服机器人正在重构企业服务体系
人工智能·机器人
IT大师兄吖2 小时前
fashn-vton-1.5 一键换装 懒人整合包
人工智能