在计算机视觉领域,YOLO(You Only Look Once)系列模型因其实时目标检测的能力而广受欢迎。其中,YOLOv8是最新版本,具备更强大的特征学习和推理能力。本文将介绍如何使用YOLOv8自定义数据集,并创建一个简单的前端界面来展示检测结果。
一、数据集准备
要使用YOLOv8进行训练,首先需要准备自己的数据集。数据集需包含图片及相应的标签。YOLO格式的标签文件通常是每个图像对应一个.txt文件,文件中每一行对应一个目标,格式为:
class_id center_x center_y width height
其中,class_id
为类别编号,center_x
、center_y
、width
、height
均为相对于图像宽高的归一化值。
将数据集分为训练集和验证集,并创建一个配置文件data.yaml
,内容示例如下:
train: path/to/train/images
val: path/to/val/images
nc: 2 # 类别数
names: ['class1', 'class2'] # 类别名称
二、模型训练
安装YOLOv8所需的依赖,可以使用以下命令:
pip install torch torchvision torchaudio
pip install -U ultralytics
接下来可以使用以下代码进行训练:
from ultralytics import YOLO
# 加载预训练模型
model = YOLO('yolov8n.pt') # 使用YOLOv8的预训练模型
# 训练模型
results = model.train(data='data.yaml', epochs=50, imgsz=640, batch=16)
在训练过程中,模型会生成一系列日志,包括损失值和精度等。训练完成后,模型会保存于runs/train/exp/weights/best.pt
。
三、模型推理
训练完成后,可以进行推理以检测新图像中的目标。可以使用以下代码:
import cv2
# 加载训练好的模型
model = YOLO('runs/train/exp/weights/best.pt')
# 读取图像
img = cv2.imread('path/to/your/image.jpg')
# 进行目标检测
results = model(img)
# 处理检测结果
for result in results:
boxes = result.boxes.xyxy # 边界框坐标
confidences = result.boxes.conf # 置信度
classes = result.boxes.cls # 类别
for box, conf, cls in zip(boxes, confidences, classes):
print(f'Class: {cls}, Confidence: {conf}, Box: {box}')
四、前端界面
为了便于使用和展示检测结果,我们可以构建一个简单的前端界面。这里我们使用Flask作为后端,HTML+JavaScript作为前端。
首先安装Flask:
pip install flask
然后创建一个简单的Flask应用,app.py
:
from flask import Flask, request, render_template
import cv2
from ultralytics import YOLO
app = Flask(__name__)
model = YOLO('runs/train/exp/weights/best.pt')
@app.route('/')
def index():
return render_template('index.html')
@app.route('/predict', methods=['POST'])
def predict():
file = request.files['file']
file.save('uploaded_image.jpg')
img = cv2.imread('uploaded_image.jpg')
results = model(img)
# 处理检测结果
# (此处省略结果处理代码,可按照前面推理的代码进行处理)
# 返回前端结果
if __name__ == '__main__':
app.run(debug=True)
index.html
的简单示例:
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>YOLOv8 目标检测</title>
</head>
<body>
<h1>上传图片进行目标检测</h1>
<form action="/predict" method="post" enctype="multipart/form-data">
<input type="file" name="file">
<input type="submit" value="上传">
</form>
</body>
</html>
五、总结
通过以上步骤,我们成功地训练了YOLOv8模型,并构建了一个简单的Web界面以进行目标检测。用户可以通过上传图片快速查看检测结果,这为目标检测的应用提供了良好的基础。随着技术的进步与完善,期待未来能有更多更强大的功能出现。