Using Yolov8 for object detection and training your own dataset
1. Prerequisite preparation
- Source code download
https://github.com/ultralytics/ultralytics - YOLOv8 documentation:
https://docs.ultralytics.com/
2. Download the code and configure the environment
Linux can use the following command to configure the environment. Of course, if it is under Windows, just download the compressed package and unzip it.
Code download
git clone https://github.com/ultralytics/ultralytics # clone repo
Environment configuration
The requirements.txt contains the necessary configuration environment:
basically as follows:
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/
Check whether the test environment is configured successfully:
import torch
print(torch.__version__)
print(torch.cuda.is_available())
a = torch.Tensor(5,3)
a=a.cuda()
print(a)
3. Download the pre-trained model
Go to yolo official github to download four versions of the model
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
4. Forecast
yolo v8 official detection category
['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']
prediction command
yolo predict model=./weights/yolov8l_23_0320.pt source=./input/*.jpg save
Pictures will be saved
yolo predict model=./weights/yolov8l_23_0320.pt source=./a.mp4 save
Video will be saved
prediction code
Predict a single image
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)
Parameter Description
https://docs.ultralytics.com/cfg/
train
1. Tag
Prepare images and labels
2. Create a new make_train_val.py file in yolov8/datasets and divide the training and verification sets.
'''
*******************************************************************************
函数名称: 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()
3、my_train.yaml
Create a new training configuration file
# 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
4. Start training
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