yolov8训练记录GPU(linux)
·
训练yolo8算法,记录下过程,避免每次训练都忘记
1.xml格式转txt
import os
import xml.etree.ElementTree as ET
def convert(size, box):
"""
将边界框从Pascal VOC格式转换为YOLO格式。
:param size: 图片尺寸 (width, height)
:param box: Pascal VOC格式的边界框 (xmin, ymin, xmax, ymax)
:return: YOLO格式的中心点坐标和宽高 (center_x, center_y, width, height)
"""
dw = 1. / size[0]
dh = 1. / size[1]
x = (box[0] + box[2]) / 2.0
y = (box[1] + box[3]) / 2.0
w = box[2] - box[0]
h = box[3] - box[1]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return x, y, w, h
def convert_annotation(xml_file_path, classes, output_dir):
"""
解析XML文件并生成YOLO格式的标签。
:param xml_file_path: XML文件路径
:param classes: 类别列表
:param output_dir: 输出目录
"""
tree = ET.parse(xml_file_path)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
out_file_path = os.path.join(output_dir, os.path.basename(xml_file_path).replace('.xml', '.txt'))
out_file = open(out_file_path, 'w')
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
if __name__ == "__main__":
xml_labels_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/xmls' # XML标签所在目录
classes_file = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/classes.txt' # 类别文件路径
output_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/test' # 输出YOLO标签的目录
# 加载类别
with open(classes_file, 'r') as f:
classes = [line.strip() for line in f.readlines()]
# 确保输出目录存在
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 遍历所有XML文件并转换
for xml_file in os.listdir(xml_labels_dir):
if not xml_file.endswith('.xml'):
continue
xml_file_path = os.path.join(xml_labels_dir, xml_file)
convert_annotation(xml_file_path, classes, output_dir)
print("转换完成")
2.训练集、测试集转换代码
import os
import random
import shutil
def split_yolo_dataset(image_dir, label_dir, output_dir, train_ratio=0.8, val_ratio=0.2, test_ratio=0):
"""
Splits the YOLO-formatted dataset into training, validation, and testing sets based on the provided ratios.
"""
if not (train_ratio + val_ratio + test_ratio) == 1.0:
raise ValueError("The sum of train_ratio, val_ratio, and test_ratio must be 1.0")
# List all images without extensions
images = [os.path.splitext(f)[0] for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
total_images = len(images)
print(f"Total images found: {total_images}")
# Shuffle images to ensure randomness
random.shuffle(images)
# Calculate sizes of each split using round to avoid rounding issues
train_count = round(total_images * train_ratio)
val_count = round(total_images * val_ratio)
test_count = total_images - train_count - val_count
# Adjust counts if necessary to ensure no images are left out
if test_count < 0:
val_count += test_count
test_count = 0
elif train_count + val_count + test_count != total_images:
# In case of rounding error, adjust the count of the last set
diff = total_images - (train_count + val_count + test_count)
if test_count > 0:
test_count += diff
elif val_count > 0:
val_count += diff
else:
train_count += diff
# Split the dataset
splits = {
'train': images[:train_count],
'val': images[train_count:train_count + val_count],
'test': images[train_count + val_count:]
}
# Create output directories and copy/move files accordingly
for split_name, split_images in splits.items():
if not split_images: # 如果该集合没有图像,则跳过创建目录
continue
split_image_dir = os.path.join(output_dir, split_name, 'images')
split_label_dir = os.path.join(output_dir, split_name, 'labels')
os.makedirs(split_image_dir, exist_ok=True)
os.makedirs(split_label_dir, exist_ok=True)
for img_base in split_images:
# Find the correct extension for the image file
img_ext = next(
(ext for ext in ['.jpg', '.jpeg', '.png'] if os.path.exists(os.path.join(image_dir, img_base + ext))),
'')
if not img_ext:
print(f"Warning: Could not find an image for base name {img_base}")
continue
src_img_path = os.path.join(image_dir, img_base + img_ext)
dst_img_path = os.path.join(split_image_dir, img_base + img_ext)
shutil.copy(src_img_path, dst_img_path)
src_label_path = os.path.join(label_dir, img_base + '.txt')
if os.path.exists(src_label_path):
dst_label_path = os.path.join(split_label_dir, img_base + '.txt')
shutil.copy(src_label_path, dst_label_path)
else:
print(f"Warning: No corresponding label file found for {img_base}")
if __name__ == "__main__":
image_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/image' # 替换为实际图像所在目录
label_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/非机动车/now/test' # 替换为实际YOLO标签所在目录
output_dir = '/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/file/3.6-now' # 替换为你希望存放分割后数据集的目录
split_yolo_dataset(image_dir, label_dir, output_dir, train_ratio=0.8, val_ratio=0.2, test_ratio=0)
3.训练脚本
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO('/home/columbus/hlz/3.6/yolov8/ultralytics-main/ultralytics/cfg/models/v8/yolov8.yaml') # 指定YOLO模型对象,并加载指定配置文件中的模型配置
model.load('/home/columbus/hlz/3.6/yolov8/ultralytics-main/yolov8s.pt') #加载预训练的权重文件'yolov8s.pt',加速训练并提升模型性能
model.train(data='/home/columbus/hlz/3.6/yolov8/ultralytics-main/data/date.yaml', # 指定训练数据集的配置文件路径,这个.yaml文件包含了数据集的路径和类别信息
epochs=150, # 设置训练的总轮数为200轮
batch=32, # 设置每个训练批次的大小为16,即每次更新模型时使用16张图片
#close_mosaic=10, # 设置在训练结束前多少轮关闭 Mosaic 数据增强,10 表示在训练的最后 10 轮中关闭 Mosaic
#workers=8, # 设置用于数据加载的线程数为8,更多线程可以加快数据加载速度
#patience=50, # 在训练时,如果经过50轮性能没有提升,则停止训练(早停机制)
#device='0', # 指定使用的设备,'0'表示使用第一块GPU进行训练
optimizer='SGD', # 设置优化器为SGD(随机梯度下降),用于模型参数更新
)
更多推荐



所有评论(0)