当前位置: 首页 > article >正文

yolov8环境搭建+训练自己的数据集

新建conda环境

conda create -n env_yolov8_bicycle python=3.10
conda activate env_yolov8_bicycle

安装pytorch和ultralytics

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
pip install ultralytics

yolov8源码下载下载

https://github.com/ultralytics/ultralytics

在这里插入图片描述

准备数据集

首先在根目录下,新建datasets文件夹
这里用的是随便下载的数据集(绝缘子缺陷检测),将下载的insulator数据集文件夹复制到datasets文件夹中。但是这个数据集是voc格式,需要将其转换为yolo格式

https://aistudio.baidu.com/datasetdetail/122549

在这里插入图片描述

voc转yolo

在datasets文件夹下,新建xml2txt.py。注意,需要修改imgpath 、xmlpath、txtpath。

# xml2txt.py
import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
from os import listdir
from os.path import join

classes = []

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')
        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            bb = convert((w, h), b)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


if __name__ == "__main__":
    postfix = 'jpg'
    imgpath = 'insulator/JPEGImages' #图片路径,需修改 
    xmlpath = 'insulator/Annotations' #voc标注文件路径,,需修改 
    txtpath = 'insulator/txt' #yolo标注文件路径,需修改 
    
    if not os.path.exists(txtpath):
        os.makedirs(txtpath, exist_ok=True)
    
    list = os.listdir(xmlpath)
    error_file_list = []
    for i in range(0, len(list)):
        try:
            path = os.path.join(xmlpath, list[i])
            if ('.xml' in path) or ('.XML' in path):
                convert_annotation(path, list[i])
                print(f'file {list[i]} convert success.')
            else:
                print(f'file {list[i]} is not xml format.')
        except Exception as e:
            print(f'file {list[i]} convert error.')
            print(f'error message:\n{e}')
            error_file_list.append(list[i])
    print(f'this file convert failure\n{error_file_list}')
    print(f'Dataset Classes:{classes}')

划分train、test、vaild

将数据集从voc转换为yolo格式化后,按照0.7 0.15 0.15的比例,进行分割。在datasets文件夹下,新建xml2txt.py。

# split.py
import os
import random
import shutil


def copy_files(src_dir, dst_dir, filenames, extension):
    os.makedirs(dst_dir, exist_ok=True)
    missing_files = 0
    for filename in filenames:
        src_path = os.path.join(src_dir, filename + extension)
        dst_path = os.path.join(dst_dir, filename + extension)

        # Check if the file exists before copying
        if os.path.exists(src_path):
            shutil.copy(src_path, dst_path)
        else:
            print(f"Warning: File not found for {filename}")
            missing_files += 1

    return missing_files


def split_and_copy_dataset(image_dir, label_dir, output_dir, train_ratio=0.7, valid_ratio=0.15, test_ratio=0.15):
    # 获取所有图像文件的文件名(不包括文件扩展名)
    image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]

    # 随机打乱文件名列表
    random.shuffle(image_filenames)

    # 计算训练集、验证集和测试集的数量
    total_count = len(image_filenames)
    train_count = int(total_count * train_ratio)
    valid_count = int(total_count * valid_ratio)
    test_count = total_count - train_count - valid_count

    # 定义输出文件夹路径
    train_image_dir = os.path.join(output_dir, 'train', 'images')
    train_label_dir = os.path.join(output_dir, 'train', 'labels')
    valid_image_dir = os.path.join(output_dir, 'valid', 'images')
    valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
    test_image_dir = os.path.join(output_dir, 'test', 'images')
    test_label_dir = os.path.join(output_dir, 'test', 'labels')

    # 复制图像和标签文件到对应的文件夹
    train_missing_files = copy_files(image_dir, train_image_dir, image_filenames[:train_count], '.jpg')
    train_missing_files += copy_files(label_dir, train_label_dir, image_filenames[:train_count], '.txt')

    valid_missing_files = copy_files(image_dir, valid_image_dir, image_filenames[train_count:train_count + valid_count],
                                     '.jpg')
    valid_missing_files += copy_files(label_dir, valid_label_dir,
                                      image_filenames[train_count:train_count + valid_count], '.txt')

    test_missing_files = copy_files(image_dir, test_image_dir, image_filenames[train_count + valid_count:], '.jpg')
    test_missing_files += copy_files(label_dir, test_label_dir, image_filenames[train_count + valid_count:], '.txt')

    # Print the count of each dataset
    print(f"Train dataset count: {train_count}, Missing files: {train_missing_files}")
    print(f"Validation dataset count: {valid_count}, Missing files: {valid_missing_files}")
    print(f"Test dataset count: {test_count}, Missing files: {test_missing_files}")


# 使用例子
image_dir = 'bicycle_dataset/images'
label_dir = 'bicycle_dataset/labels'
output_dir = './bicycle_dataset_split'

split_and_copy_dataset(image_dir, label_dir, output_dir)


新建yaml文件

在datasets文件夹中,创建 main.yaml 文件。注意,这里要修改path路径

path: C:\AI\demo_yolov8_1\datasets\insulator_split
train: train/images
val: valid/images
test: test/images
# number of classes
nc: 1

# class names
names: ['insulator']

配置Yolov8

切换至ultralytics/cfg/models/v8 目录,修改yolov8.yaml文件,将nc改为1(这里主要修改的是类别数nc,数据集有几种类型,就填几)。

模型训练

下载预训练模型yolo8s.pt,放在datasets下
在Yolov8的GitHub开源库中下载对应版本的模型。我选的是yolo8s.pt。

https://pan.baidu.com/s/1tsOxijvq9auwYw0v5jrXWA?pwd=8888
  1. 运行
yolo task=detect mode=train model=datasets/yolov8s.pt epochs=100 batch=10 data=datasets/main.yaml

http://www.kler.cn/a/371056.html

相关文章:

  • ZooKeeper的应用场景:深入探讨分布式系统中的多样化应用
  • CheckPointUtilsTest
  • Kafka 物理存储机制
  • 【使用Flask构建RESTful API】从零开始开发简单的Web服务!
  • Python 爬虫的寻宝大冒险:如何捕获 API 数据的宝藏
  • 重学SpringBoot3-怎样优雅停机
  • Notepad++如何同时检索多个关键字
  • 单目相机标定
  • PG数据库之视图详解
  • React写关键字高亮的三个方案
  • Node.js:内置模块
  • Docker | 通过commit操作实例来认识镜像底层实现的原理以及学会打包镜像
  • 前端零基础入门到上班:【Day3】从零开始构建网页骨架HTML
  • 【AIGC】2024-arXiv-Lumiere:视频生成的时空扩散模型
  • static函数和友元
  • Linux云计算 |【第五阶段】ARCHITECTURE-DAY5
  • 【华为HCIP实战课程十八】OSPF的外部路由类型,网络工程师
  • df -h 和df -i 的区别
  • pip命令行安装pytest 一直报错
  • springboot075电影评论网站系统设计与实现(论文+源码)_kaic
  • Bash、sh 和 Shell都弄混了?
  • 软考架构师论文范文(游戏后端)
  • 为 Python 项目配置 Git 忽略文件的实用指南
  • SSA-CNN-LSTM-MATT多头注意力机制多特征分类预测
  • vue2项目在发布后更新,提示用户刷新页面
  • MySQL史上最全总结