在计算机视觉领域,YOLO(You Only Look Once)系列模型因其实时目标检测的能力而广受欢迎。其中,YOLOv8是最新版本,具备更强大的特征学习和推理能力。本文将介绍如何使用YOLOv8自定义数据集,并创建一个简单的前端界面来展示检测结果。

一、数据集准备

要使用YOLOv8进行训练,首先需要准备自己的数据集。数据集需包含图片及相应的标签。YOLO格式的标签文件通常是每个图像对应一个.txt文件,文件中每一行对应一个目标,格式为:

class_id center_x center_y width height

其中,class_id为类别编号,center_xcenter_ywidthheight均为相对于图像宽高的归一化值。

将数据集分为训练集和验证集,并创建一个配置文件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界面以进行目标检测。用户可以通过上传图片快速查看检测结果,这为目标检测的应用提供了良好的基础。随着技术的进步与完善,期待未来能有更多更强大的功能出现。

点赞(0) 打赏

微信小程序

微信扫一扫体验

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部