前言
YOLOv8 基于深度学习和计算机视觉领域的尖端技术,在速度和准确性方面具有无与伦比的性能。其流线型设计使其适用于各种应用,并可轻松适应从边缘设备到云 API 等不同硬件平台。
一、数据集准备
数据集文件存放规则:
以下是随机划分数据集的代码:
import os
import shutil
import random
# Directories
labels_dir = r"E:\File\VOC2020\output"
images_dir = r"E:\File\VOC2020\images"
total_dir = r"E:\File\VOC2020\fire_detect"
# Output directories
train_images_dir = os.path.join(total_dir, 'train', 'images')
train_labels_dir = os.path.join(total_dir, 'train', 'labels')
valid_images_dir = os.path.join(total_dir, 'valid', 'images')
valid_labels_dir = os.path.join(total_dir, 'valid', 'labels')
test_images_dir = os.path.join(total_dir, 'test', 'images')
test_labels_dir = os.path.join(total_dir, 'test', 'labels')
# Create output directories if they don't exist
os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(valid_images_dir, exist_ok=True)
os.makedirs(valid_labels_dir, exist_ok=True)
os.makedirs(test_images_dir, exist_ok=True)
os.makedirs(test_labels_dir, exist_ok=True)
# Get list of all image paths
image_paths = os.listdir(images_dir)
random.shuffle(image_paths)
# Split data into train, valid, and test sets
num_images = len(image_paths)
train_split = int(0.7 * num_images)
valid_split = int(0.2 * num_images)
train_image_paths = image_paths[:train_split]
valid_image_paths = image_paths[train_split:train_split + valid_split]
test_image_paths = image_paths[train_split + valid_split:]
# Function to copy images and labels
def copy_files(image_paths, dest_images_dir, dest_labels_dir):
for image_path in image_paths:
image_name = os.path.splitext(image_path)[0]
xml_path = os.path.join(labels_dir, image_name + '.txt')
if not os.path.exists(xml_path):
continue
# Copy image
shutil.copy(os.path.join(images_dir, image_path), dest_images_dir)
# Copy label
shutil.copy(xml_path, dest_labels_dir)
# Copy files to respective directories
copy_files(train_image_paths, train_images_dir, train_labels_dir)
copy_files(valid_image_paths, valid_images_dir, valid_labels_dir)
copy_files(test_image_paths, test_images_dir, test_labels_dir)
print("Done")
把图片和标签区分开
def split_image_label(dir_path, image_dir, label_dir):
# 获取目录中的所有文件名
file_names = os.listdir(dir_path)
for file_name in file_names:
# 获取完整文件路径
file_path = os.path.join(dir_path, file_name)
# 获取文件扩展名
_, ext = os.path.splitext(file_path)
ext = ext.lower()
if ext == '.xml' or ext==".txt":
# 如果是XML文件,复制到label_dir目录
dest_path = os.path.join(label_dir, file_name)
shutil.copy(file_path, dest_path)
elif ext == '.jpg' or ext == '.jpeg':
# 如果是JPG文件,复制到image_dir目录
dest_path = os.path.join(image_dir, file_name)
shutil.copy(file_path, dest_path)
else:
# 其他类型文件,跳过处理
pass
xml格式转换成txt格式:
import xml.etree.ElementTree as ET
import os
import shutil
labels = {"fire": 0}
def convert_to_yolo_format(xml_file):
tree = ET.parse(xml_file)
root = tree.getroot()
size = root.find('size')
width = int(size.find('width').text)
height = int(size.find('height').text)
yolo_annotations = []
for obj in root.iter('object'):
# Assuming class names are consistent and mapped to IDs
class_id = obj.find('name').text
class_id = labels.get(class_id)
if class_id is None:
continue # Skip unknown classes
bndbox = obj.find('bndbox')
xmin = int(bndbox.find('xmin').text)
ymin = int(bndbox.find('ymin').text)
xmax = int(bndbox.find('xmax').text)
ymax = int(bndbox.find('ymax').text)
x_center = (xmin + xmax) / 2.0 / width
y_center = (ymin + ymax) / 2.0 / height
bbox_width = (xmax - xmin) / width
bbox_height = (ymax - ymin) / height
yolo_annotations.append(
f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")
return yolo_annotations
def create_yolo_file(yolo_annotations, output_file):
os.makedirs(os.path.dirname(output_file), exist_ok=True)
with open(output_file, 'a+') as f:
for line in yolo_annotations:
f.write(line + '\n')
images_dir = r"E:\File\VOC2020\images"
labels_dir = r"E:\File\VOC2020\labels"
output_dir = r"E:\File\VOC2020\output"
image_paths = os.listdir(images_dir)
for image_path in image_paths:
image_name = os.path.splitext(image_path)[0]
xml_path = os.path.join(labels_dir, image_name+'.xml')
if not os.path.exists(xml_path):
continue
yolo_annotations = convert_to_yolo_format(xml_path)
output_file = os.path.join(output_dir, image_name + '.txt')
create_yolo_file(yolo_annotations, output_file)
print("Done")
二、训练模型
1.写好yaml文件
train: /root/product/fire_detect/train # train images (relative to 'path') 128 images
val: /root/product/fire_detect/valid # val images (relative to 'path') 128 images
test: /root/product/fire_detect/test # test images (optional)
# Classes
names:
0: fire
2.开始训练
from ultralytics import YOLO
import os
if __name__ == "__main__":
# Load a model
data_yaml = r"/root/product/data.yaml"
model = YOLO("yolov8n.yaml") # build a new model from YAML
model = YOLO("yolov8n.yaml").load("/root/product/yolov8n.pt") # build from YAML and transfer weights
# model = YOLO(pre_model, task='detect') # load a pretrained model (recommended for training)
# Train the model
results = model.train(data=data_yaml, epochs=200, imgsz=640, batch=16, workers=8, device=0,
cos_lr=True, close_mosaic=200, warmup_epochs=10)
/root/product/yolov8n.pt是预训练模型需要自己去官网下
3.模型推理
from ultralytics import YOLO
import cv2
# Load a model
model = YOLO(r"/root/product/runs/detect/train3/weights/best.pt") # pretrained YOLOv8n model
image=cv2.imread("/root/product/1.jpg")
# Run batched inference on a list of images
results = model(image) # return a list of Results objects
# Process results list
for result in results:
result.save(filename="result.jpg") # save to disk