该博客主要介绍了:
1. 如何制作yolo目标检测数据集
2.如何在自己的数据集上训练yolo
3.训练好后的模型如何进行推理
1.数据标注
关于数据如何标注,请查看这篇博文
2.数据集目录结构
重点关注红框内部的结构
images: 图片目录
images/train: 训练集图片
iamges/val: 验证集图片
labels 同理
3.数据配置文件
存放位置
# 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, ..]
path: ./card_data # dataset root dir
train: images/train # train images (relative to 'path') 4 images
val: images/val # val images (relative to 'path') 4 images
test: # test images (optional)
names:
0: name
1: type
2: race
3: attr
4: lv
names 是物体类别,在这个案例中,我自己制作了一个数据集,目的是为了从图片中提取特定文本片段,下图给出了一个直观示例:
4.训练环境配置
python 3.10
pytorch 2.1 # 2.2版本执行的时候报错了
ultralytics 8.1.20
5.训练
from ultralytics import YOLO
# Load a model
model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training)
# Train the model
results = model.train(data='./data.yml', epochs=100, imgsz=640, batch=32)
print(1)
6.预测
from ultralytics import YOLO
# Load a model
# model = YOLO('yolov8n.pt') # load an official model
model = YOLO('./runs/detect/train/weights/best.pt') # load a custom model
# Predict with the model
results = model('./test_imgs') # predict on an image
# Run batched inference on a list of images
# results = model(['im1.jpg', 'im2.jpg']) # return a list of Results objects
# Process results list
for result in results:
boxes = result.boxes # Boxes object for bounding box outputs
masks = result.masks # Masks object for segmentation masks outputs
keypoints = result.keypoints # Keypoints object for pose outputs
probs = result.probs # Probs object for classification outputs
result.show() # display to screen
result.save(filename='result.jpg') # save to disk
7. 数据集
下载地址