这几天做了给Yolov8检测物体的小任务,今天来做下总结。
首先介绍下整个Yolov8检测的步骤吧,安装库那些就不讲了。
这是我的文件包的对象树。
有images包,里面装了训练和验证的图像。
labels包,装了标注好的labels的txt文件,一开始标注好是json文件,要经过脚本转化,变成yolov能识别的txt文件。
这里我之前遇到一个麻烦,就是我已经开始是用多边形标注的,结果yolov8训练出来一坨,我查了查yolov8好像在目标识别这块不能用多边形。分割的话可以用多边形做。然后就改成矩形做了。
dataset.yaml文件
path: C:/Users/ren/Desktop/gm/guardrail_monitoring/datasets/guardrail
train: C:/Users/ren/Desktop/gm/guardrail_monitoring/datasets/guardrail/images/train
val: C:/Users/ren/Desktop/gm/guardrail_monitoring/datasets/guardrail/images/val
# 类别信息
nc: 1
names: ['fence']
script文件里装的是运行脚本
train.py
from ultralytics import YOLO
def train_model():
"""
训练 YOLOv8 目标检测模型。
"""
# 加载预训练的目标检测模型
model = YOLO("yolov8n.pt") # 使用 YOLOv8 的目标检测模型
# 训练模型
results = model.train(
data="C:/Users/ren/Desktop/gm/guardrail_monitoring/datasets/dataset.yaml", # 数据集配置文件路径
epochs=100, # 训练轮数
imgsz=640, # 图像大小
batch=16, # 批量大小
name="guardrail_detection", # 训练任务名称
patience=10, # 早停机制,如果 10 轮验证集性能没有提升,则停止训练
device="0", # 使用 GPU 训练(如果有 GPU)
workers=4, # 数据加载的线程数
optimizer="auto", # 自动选择优化器
lr0=0.01, # 初始学习率
lrf=0.01, # 最终学习率
weight_decay=0.0005, # 权重衰减
save=True, # 保存训练结果
save_period=10, # 每 10 轮保存一次模型
)
print("训练完成!模型权重保存在 runs/detect/guardrail_detection/weights/ 目录下。")
if __name__ == "__main__":
train_model()
device = “0”就代表我用的GPU训练,记得要装pytorch的GPU版本。
训练完了之后权值会在run/../weight文件里面,选里面最好的best,吧它的地址换进,infer_image文件里面
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import os
def draw_boxes(image, boxes, labels, confidences):
"""
在图像上绘制边界框和标签。
:param image: PIL 图像对象
:param boxes: 边界框坐标列表
:param labels: 类别标签列表
:param confidences: 置信度列表
:return: 绘制后的图像
"""
draw = ImageDraw.Draw(image)
font = ImageFont.load_default()
for box, label, confidence in zip(boxes, labels, confidences):
x1, y1, x2, y2 = box
draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
draw.text((x1, y1), f"{label} {confidence:.2f}", fill="red", font=font)
return image
def infer_image(image_path, output_dir):
"""
对单张图像进行推理。
:param image_path: 图像路径
:param output_dir: 输出目录
"""
# 加载训练好的模型
model = YOLO("C:/Users/ren/Desktop/gm/runs/detect/guardrail_detection2/weights/best.pt")
# 推理(降低置信度阈值)
results = model(image_path, conf=0.1) # 设置置信度阈值为 0.1
# 处理结果
for result in results:
boxes = result.boxes.xyxy.tolist() # 获取边界框坐标
labels = result.boxes.cls.tolist() # 获取类别 ID
confidences = result.boxes.conf.tolist() # 获取置信度
names = result.names # 获取类别名称
# 将类别 ID 转换为类别名称
labels = [names[int(cls)] for cls in labels]
# 打开图像
image = Image.open(image_path)
# 绘制标注
if boxes: # 如果有检测到目标
image = draw_boxes(image, boxes, labels, confidences)
else:
print(f"{image_path} 未检测到目标!")
# 创建输出目录(如果不存在)
os.makedirs(output_dir, exist_ok=True)
# 保存图像
image_name = os.path.basename(image_path) # 获取图像文件名
output_path = os.path.join(output_dir, image_name)
image.save(output_path)
print(f"推理完成!结果保存在 {output_path}")
def infer_validation_set(validation_dir, output_dir):
"""
对验证集进行批量推理。
:param validation_dir: 验证集目录
:param output_dir: 输出目录
"""
# 遍历验证集目录中的所有图像
for image_name in os.listdir(validation_dir):
if image_name.lower().endswith((".jpg", ".jpeg", ".png")): # 仅处理图像文件
image_path = os.path.join(validation_dir, image_name)
infer_image(image_path, output_dir)
if __name__ == "__main__":
validation_dir = "C:/Users/ren/Desktop/gm/guardrail_monitoring/datasets/guardrail/images/val" # 验证集目录
output_dir = "C:/Users/ren/Desktop/outputs/images" # 输出目录
infer_validation_set(validation_dir, output_dir)
最后就训练好了Yolov8模型
明天来解析一下yolov8的源代码逻辑。