检查是否安装好环境
import torch
import torchvision
print(torch.__version__)
print(torchvision.__version__)
print(torch.cuda.is_available())
print(torch.cuda.device_count())
yaml文件的编写
path: Visdrone
train: train/images
val: val/images
test: test/images
nc: 10
names:
0: pedestrian
1: people
2: bicycle
3: car
4: van
5: truck
6: tricycle
7: awning-tricycle
8: bus
9: motor
不需要指定标注文件,只需要填写图片的路径
数据集的划分(v8)
数据集转换(visdrone为例)
import os
from pathlib import Path
from PIL import Image
from tqdm import tqdm
def visdrone2yolo(dir):
def convert_box(size, box):
#Convert VisDrone box to YOLO CxCywh box,坐标进行了归一化
dw = 1. / size[0]
dh = 1. / size[1]
return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh
# (dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory
(dir / 'Annotations_YOLO').mkdir(parents=True, exist_ok=True) # make labels directory
pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
for f in pbar:
img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
lines = []
with open(f, 'r') as file: # read annotation.txt
for row in [x.split(',') for x in file.read().strip().splitlines()]:
if row[4] == '0': # VisDrone 'ignored regions' class 0
continue
cls = int(row[5]) - 1
box = convert_box(img_size, tuple(map(int, row[:4])))
lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'Annotations_YOLO' + os.sep), 'w') as fl:
fl.writelines(lines) # write label.txt
dir = Path(r'E:\DPL\DeepLearnData\目标检测\航空目标检测数据VisDrone\VisDrone2019') #更改为自己的数据集的路径
# Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels
训练
train.py 作用:用于训练
from ultralytics import YOLO
import os
if __name__ == '__main__':
model = YOLO(model="my_lu.yaml")
# Use the model
results = model.train(data="data.yaml", patience=200,epochs=200, device='0', imgsz=640, batch=4, seed=42)
用于得出结果
val.py,作用:根据训练好的结果(pt文件),可以运行出训练结果
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO('runs/detect/train9/weights/best.pt')
metrics = model.val(split='val')
测试多张图片
test.py 作用:可以用训练完的pt文件检测多张图片,将图片放到yolo/assets文件夹下。
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO('runs/detect/train9/weights/best.pt')
results = model.predict(source="drone-yolo/assets", device='0',save=True)
print(results)
断点训练
将未训练好的pt文件指定路径,和数据集的yaml文件
yolo task=detect mode=train model=runs/detect/train19/weights/last.pt data=data.yaml save=True resume=True
本片文章只是自己的笔记,写的很简单,可能有读者看不懂,望海涵。如有错误,请指教,谢谢。