1安装 Ultralytics
pip install ultralytics
2 准备数据
https://docs.ultralytics.com/datasets/detect/?h=data+format#ultralytics-yolo-format
yolo的每张图片都要有一个对应label 的txt文件,如果图片中没有label,则也不需要txt文件,txt文件中数据格式为:
class x_center y_center width height
数据需要进行标准化,如果打标的box是像素坐标,则 x_center,width 要除以图片width,y_center ,height要除以图片height。
iw,ih = image.size
x_center = (x + w/2) / iw
y_center = (y + h/2) /ih
box_w = w / iw
box_h = h / ih
Class类别的序号必须从0开始Class类别的序号必须从0开始
label 文件格式如下:
目录是这样
配置data.yaml
import yaml
yaml_dict = {
'train': 'train_data/images/train',
'val':'train_data/images/val',
'nc':3, #类别数
'names':['logo0','logo1','logo2']
}
with open('data.yaml','w') as file:
yaml.dump(yaml_dict,file)
Ultralytics Settings
ultralytics使用settings管理配置,其中比较重要的是datasets_dir,训练数据的存储路径,修改为自己的路径
from ultralytics import settings
print(settings)
settings.update({'datasets_dir': '/path/to/data'})
settings.reset()
3 训练
yolo提供了各种不同用途不同大小的模型,按需使用
https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes
from ultralytics import YOLO
model = YOLO('yolov8n.pt')
model.train(data="data.yaml",epochs=20,imgsz=640,batch=16)
训练好的模型默认保存在 'runs/detect/train'
预测
from PIL import Image
import cv2
img = Image.open()
img = cv2.imread()
model = YOLO('model_path')
results = model.predict(img,imgsz=640) #conf=0.5
#results[0].boxes 中返回了各种格式的结果
boxes = results[0].boxes.xyxy.cpu()
boxes = np.array(boxes).astype(np.int32)
print(boxes)