UmVfX1BvaW50's Blog
0%

yolov8模型训练

分类于
  • yolov8相关

Install

ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite (github.com)

数据集架构

train

  • images
  • labels

valid

  • images
  • labels

test

  • images
  • labels

数据集格式

1
2
0 1670.21826171875 2155.037841796875 921.8790283203125 951.6988525390625
[标签 坐标（相对绝对均可，不出边界）]

配置文件data.yaml

1
2
3
4
5
6
7
// 数据集位置
train: D:\Demo\python\pythonProject6\train\images
val: D:\Demo\python\pythonProject7\valid\images
test:

nc: 2	//分类数
names: ['no-injecStor', 'injector']	//类别标签名

模型训练

  • 内置模型yolov8n.yaml / .pt，N M L对应模型复杂度
1
2
3
4
5
6
7
8
9
10
from ultralytics import YOLO

model = YOLO("yolov8n.yaml")	# 导入预训练模型
model.train(data="data.yaml", epochs=130) 
'''
epochs 迭代次数
lr0 初始学习率
lr1 最终学习率
'''

模型预测

1
2
3
4
5
6
7
8
from ultralytics import YOLO

pth_path = r"D:\Demo\python\pythonProject2\runs\detect\train\weights\best.pt"
model = YOLO(pth_path)  # load a custom model

# Predict with the model
# model.predict
results = model("D:\\Demo\\python\\pythonProject2\\test\\images\\0_Concern-In-China-As-Mystery-Virus-Spreads_jpg.rf.5633f5fe7a9b926101b7fc16615dfb6a.jpg", save=True, conf=0.5, iou=0.5) # 保存、置信度、交互比

辅助标注

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import os
import cv2
from ultralytics import YOLO

# 加载预训练的 YOLOv8 模型
model = YOLO('D:\Demo\python\pythonProject6\\runs\detect\\train\weights\last.pt')  # 使用预训练的 YOLOv8 模型


def annotate_images(directory, output_dir):
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 遍历目录中的所有图像文件
    for filename in os.listdir(directory):
        if filename.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
            image_path = os.path.join(directory, filename)
            output_path = os.path.join(output_dir, os.path.splitext(filename)[0] + '.txt')

            # 读取图像
            image = cv2.imread(image_path)

            # 使用模型进行预测
            results = model(image,save=True)

            # 获取预测结果
            predictions = results[0].boxes  # 获取预测框对象

            # 打开输出文件
            with open(output_path, 'w') as f:
                for pred in predictions:
                    # 获取预测结果中的类别、置信度和边界框坐标
                    class_id = int(pred.cls)
                    x_center, y_center, width, height = pred.xywh[0]  # 中心点坐标和宽高
                    confidence = pred.conf  # 置信度

                    # 将结果保存为 YOLO 格式
                    f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

            print(f"Annotated {image_path} and saved to {output_path}")


# 使用示例：将 'input_directory' 替换为包含图像的目录，将 'output_directory' 替换为保存标注结果的目录
annotate_images('C:\\Users\联想\Desktop\IMG3', 'D:\Demo\python\pythonProject7\\train\labels')

else script

同目录 pic & txt 分离

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import os
import shutil

def separate_files(source_folder, jpg_folder, txt_folder):
    # 创建目标文件夹
    os.makedirs(jpg_folder, exist_ok=True)
    os.makedirs(txt_folder, exist_ok=True)

    # 遍历源文件夹中的所有文件
    for filename in os.listdir(source_folder):
        file_path = os.path.join(source_folder, filename)

        # 检查文件是否是 JPG 或 TXT 格式
        if filename.lower().endswith('.jpg'):
            # 将 JPG 文件移动到 jpg_folder
            shutil.move(file_path, os.path.join(jpg_folder, filename))
        elif filename.lower().endswith('.txt'):
            # 将 TXT 文件移动到 txt_folder
            shutil.move(file_path, os.path.join(txt_folder, filename))

# 指定源文件夹路径、JPG 文件夹路径和 TXT 文件夹路径
source_folder = 'D:\gan'
jpg_folder = 'D:\gan\images'
txt_folder = 'D:\gan\labels'

# 调用函数分离文件
separate_files(source_folder, jpg_folder, txt_folder)