使用Yolov8进行目标检测并训练自己的数据集

一、前提准备

• 源码下载
  https://github.com/ultralytics/ultralytics
• YOLOv8 文档：
  https://docs.ultralytics.com/

二、下载代码及配置环境

linux可以使用下面命令进行环境配置，当然如果是windows下，直接下载压缩包，解压即可。

代码下载

git clone https://github.com/ultralytics/ultralytics  # clone repo

环境配置

其中requirements.txt 中包含了必要的配置环境：
基本如下:

3.10>=Python>=3.7
torch>=1.7.0
torchvision>=0.8.1

#创建虚拟环境
conda create --name pytorch_gpu python=3.8
activate pytorch_gpu
conda info --envs

#安装gpu版本pytorch
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 -i https://pypi.tuna.tsinghua.edu.cn/simple/

#安装其他依赖
pip install -r requirements.txt  -i https://pypi.tuna.tsinghua.edu.cn/simple/

#安装ultralytics，可以直接使用yolo
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple/

测试环境是否配置成功：

import torch

print(torch.__version__)
print(torch.cuda.is_available())
a = torch.Tensor(5,3)
a=a.cuda()
print(a)

三、下载预训练模型

到yolo官方github下载四个版本的模型

https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt

四、预测

yolo v8官方检测类别

['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 
'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 
'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra','giraffe', 
'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 
'banana', 'apple', 'sandwich', 'orange', 'broccoli','carrot', 'hot dog', 'pizza', 
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 
'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 
'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 
'hair drier', 'toothbrush']

预测命令

yolo predict model=./weights/yolov8l_23_0320.pt source=./input/*.jpg save

会保存图片

yolo predict model=./weights/yolov8l_23_0320.pt source=./a.mp4 save

会保存视频

预测代码

预测单个图片

from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image

from ultralytics import YOLO
from ultralytics.yolo.data.build import load_inference_source
from ultralytics.yolo.utils import ROOT, SETTINGS


MODEL = './weights/yolov8s.pt'
SOURCE = './input/bus.jpg'
#SOURCE = '0'

model = YOLO(MODEL)

img = cv2.imread(str(SOURCE))
output = model(source=img, save=True, conf=0.5,iou=0.7,save_txt=True,show=True)

cv2.waitKey(1000)
cv2.waitKey(0)

参数说明

https://docs.ultralytics.com/cfg/

训练

1、打标签

准备images和labels或者xml

2、 xml转txt（yolo格式）

如果是xml格式的，先转换成yolo格式。

xml_2_yolo.py

"""  "*******************************************************************************************
*函数功能 ：VOC格式数据集转YOLO格式数据集

*输入参数 ：
*返 回 值 ：无
*编写时间 ： 2021.11.1
*作    者 ： diyun
********************************************************************************************"""

import os.path
import xml.etree.ElementTree as ET
 
class_names = ['person']
 
xmlpath='./16_person/person_1000/xml/' #原xml路径
txtpath='./16_person/person_1000/labels/'     #转换后txt文件存放路径
files = []
 
for root, dirs, files in os.walk(xmlpath):
    None
 
number = len(files)
print(number)
i = 0
while i < number:
 
    name = files[i][0:-4]
    xml_name = name + ".xml"
    txt_name = name + ".txt"
    xml_file_name = xmlpath + xml_name
    txt_file_name = txtpath + txt_name
    print(xml_file_name)
    xml_file = open(xml_file_name)
    tree = ET.parse(xml_file)
    root = tree.getroot()
    filename = root.find('filename').text

    image_name = root.find('filename').text
    w = int(root.find('size').find('width').text)
    h = int(root.find('size').find('height').text)
 
    f_txt = open(txt_file_name, 'w+')
    content = ""
 
    first = True
 
    for obj in root.iter('object'):
 
        name = obj.find('name').text
        class_num = class_names.index(name)
 
        xmlbox = obj.find('bndbox')
 
        x1 = int(xmlbox.find('xmin').text)
        x2 = int(xmlbox.find('xmax').text)
        y1 = int(xmlbox.find('ymin').text)
        y2 = int(xmlbox.find('ymax').text)
 
        if first:
            content += str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
            first = False
        else:
            content += "\n" + \
                       str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
 
    # print(str(i / (number - 1) * 100) + "%\n")
    #print(content)
    f_txt.write(content)
    f_txt.close()
    xml_file.close()
    i += 1
 
print("done!")

3、在yolov8/datasets新建make_train_val.py文件，划分训练验证集

'''
*******************************************************************************
函数名称: ReadImage
描    述: yolov8训练，数据集的准备，从yolo数据集txt文件，分为预测训练验证
作    者：狄云
编写时间：2023.09.4


yolo task=detect mode=train model=/home/diyun/Desktop/work/python_project/yolov8/weights/yolov8s.pt \
data=/home/diyun/Desktop/work/python_project/yolov8/datasets/my_train_weizao.yaml batch=2 imgsz=640 epochs=100 pretrained=True mosaic=0.0


*******************************************************************************/
'''

import os
import random
trainval_percent = 0.2
train_percent = 0.8

sets = ['train', 'test','val']

#xmlfilepath = 'data/Annotations'
#txtsavepath = 'data/ImageSets'

ImageSets_path='ImageSets/'
Imgpath = '/home/diyun/1_train_data/1_weizao/images'    #图片文件夹
xmlfilepath = '/home/diyun/1_train_data/1_weizao/labels'  #label文件存放地址

if not os.path.exists('ImageSets/'):
    os.makedirs('ImageSets/')

total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftrainval = open('ImageSets/trainval.txt', 'w')
ftest = open('ImageSets/test.txt', 'w')
ftrain = open('ImageSets/train.txt', 'w')
fval = open('ImageSets/val.txt', 'w')


for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()



for image_set in sets:
    if not os.path.exists('labels/'):
        os.makedirs('labels/')
    image_ids = open(ImageSets_path+'%s.txt' % (image_set)).read().strip().split()
    list_file = open('%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(Imgpath+'/%s.jpg\n' % (image_id))
    list_file.close()

4、my_train.yaml

新建训练配置文件

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]

train: /home/diyun/Desktop/work/python_project/yolov8/datasets/train.txt  # voc_annotation.py生成的train.txt的路径
val: /home/diyun/Desktop/work/python_project/yolov8/datasets/val.txt   # voc_annotation.py生成的val.txt的路径


# Classes
names:
  0: weizao

5、开始训练

yolo task=detect mode=train model=/home/diyun/Desktop/work/python_project/yolov8/weights/yolov8s.pt \
data=/home/diyun/Desktop/work/python_project/yolov8/datasets/my_train_weizao.yaml batch=2 imgsz=640 epochs=100 

yolo task=detect mode=train model=/home/diyun/Desktop/work/python_project/yolov8/weights/yolov8s.pt \
data=/home/diyun/Desktop/work/python_project/yolov8/datasets/my_train_weizao.yaml batch=2 imgsz=640 epochs=100 pretrained=True mosaic=0.0

可能遇到的问题

错误日志：