这几天做了给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 YOLOdef 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 osdef 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 = boxdraw.rectangle([x1, y1, x2, y2], outline="red", width=2)draw.text((x1, y1), f"{label} {confidence:.2f}", fill="red", font=font)return imagedef 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() # 获取类别 IDconfidences = 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的源代码逻辑。