首页 > 其他分享 >如何训练——草原牛羊马目标检测数据集 数据集拥有3个类别、总计2400张图片 支持YOLO、VOC格式 已经划分为训练集、验证集、测试集 可直接进行YOLOv8使用

如何训练——草原牛羊马目标检测数据集 数据集拥有3个类别、总计2400张图片 支持YOLO、VOC格式 已经划分为训练集、验证集、测试集 可直接进行YOLOv8使用

时间:2024-11-10 21:18:31浏览次数:5  
标签:py 训练 self YOLO cv2 YOLOv8 train dataset image

在这里插入图片描述在这里插入图片描述

在这里插入图片描述
如何使用YOLOv8进行草原牛羊马的目标检测,并提供详细的训练代码和数据集准备步骤。假设你已经有一个包含2400张图片的数据集,并且这些图片已经标注了YOLO格式的标签,且已经分好训练集、验证集和测试集。

项目结构
深色版本
grassland_animal_detection/
├── dataset/
│ ├── images/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
├── weights/
│ └── yolov8s.pt
├── src/
│ ├── train.py
│ ├── predict.py
│ ├── utils.py
│ ├── ui/
│ │ └── main_window.ui
├── requirements.txt
└── README.md

  1. 安装依赖
    首先,确保你已经安装了必要的库。创建一个requirements.txt文件,内容如下:

深色版本
opencv-python
numpy
pillow
PyQt5
ultralytics
然后,使用以下命令安装依赖:

bash
深色版本
pip install -r requirements.txt
2. 数据集准备
确保你的数据集已经按照以下结构组织:

深色版本
dataset/
├── images/
│ ├── train/
│ ├── val/
│ └── test/
└── labels/
├── train/
├── val/
└── test/
每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件都是.jpg格式,而标签文件是.txt格式(YOLO格式),并且它们的名字与对应的图像文件相同。

  1. 数据集配置文件
    创建一个数据集配置文件(如grassland_animal_dataset.yaml),该文件定义了数据集的基本信息,包括路径、类别等。示例配置如下:

yaml
深色版本

训练、验证和测试的数据集路径

train: dataset/images/train
val: dataset/images/val
test: dataset/images/test
在这里插入图片描述

标签路径

labels_train: dataset/labels/train
labels_val: dataset/labels/val
labels_test: dataset/labels/test

类别名称

names:
0: 牛
1: 羊
2: 马
在这里插入图片描述

类别数量

nc: 3
4. 训练代码
4.1 train.py
这是训练模型的主程序文件。

python
深色版本
import subprocess

def train_model(dataset_config, pretrained_weights, output_dir, epochs=200, batch_size=16):
# 训练模型
train_command = [
“python”, “ultralytics/yolo/v8/detect/train.py”,
“–data”, dataset_config,
“–cfg”, “yolov8.yaml”,
“–weights”, pretrained_weights,
“–batch-size”, str(batch_size),
“–epochs”, str(epochs),
“–project”, output_dir,
“–name”, “exp”
]
subprocess.run(train_command)

if name == “main”:
dataset_config = “grassland_animal_dataset.yaml”
pretrained_weights = “weights/yolov8s.pt”
output_dir = “runs/train”

# 训练YOLOv8
train_model(dataset_config, pretrained_weights, output_dir)
  1. 模型评估
    训练完成后,可以通过验证集和测试集来评估模型的性能。YOLOv8提供了方便的命令来进行模型评估。示例如下:

bash
深色版本
python ultralytics/yolo/v8/detect/val.py --data grassland_animal_dataset.yaml --weights runs/train/exp/weights/best.pt --task val
python ultralytics/yolo/v8/detect/val.py --data grassland_animal_dataset.yaml --weights runs/train/exp/weights/best.pt --task test
6. 模型推理
你可以使用训练好的模型进行推理,检测新的图像中的目标。示例命令如下:

bash
深色版本
python ultralytics/yolo/v8/detect/predict.py --source path/to/your/image.jpg --weights runs/train/exp/weights/best.pt --img 640 --conf 0.25 --iou 0.45
7. UI设计
7.1 使用Qt Designer设计UI
使用Qt Designer设计一个简单的UI界面,包含图片检测、视频检测和摄像头实时检测的功能。保存为ui/main_window.ui。

7.2 转换UI文件
使用pyuic5工具将.ui文件转换为Python代码:

bash
深色版本
pyuic5 -x ui/main_window.ui -o src/ui_main_window.py
8. 主程序代码
8.1 main.py
这是主程序文件,负责启动UI并处理用户输入。

python
深色版本
import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QTimer
from ultralytics import YOLO
from src.ui_main_window import Ui_MainWindow
from src.utils import predict_image, predict_video, predict_camera

加载模型

model = YOLO(‘runs/train/exp/weights/best.pt’)

class MainWindow(QMainWindow, Ui_MainWindow):
def init(self):
super().init()
self.setupUi(self)
self.initUI()

def initUI(self):
    self.btn_image.clicked.connect(self.open_image)
    self.btn_video.clicked.connect(self.open_video)
    self.btn_camera.clicked.connect(self.open_camera)
    self.timer = QTimer(self)
    self.timer.timeout.connect(self.update_frame)
    self.cap = None

def open_image(self):
    file_path, _ = QFileDialog.getOpenFileName(self, "打开图片", "", "Image Files (*.jpg *.png)")
    if file_path:
        image = cv2.imread(file_path)
        result_image = predict_image(model, image)
        self.show_image(result_image)

def open_video(self):
    file_path, _ = QFileDialog.getOpenFileName(self, "打开视频", "", "Video Files (*.mp4 *.avi)")
    if file_path:
        self.cap = cv2.VideoCapture(file_path)
        self.timer.start(30)

def open_camera(self):
    self.cap = cv2.VideoCapture(0)
    self.timer.start(30)

def update_frame(self):
    ret, frame = self.cap.read()
    if ret:
        result_frame = predict_image(model, frame)
        self.show_image(result_frame)
    else:
        self.timer.stop()
        self.cap.release()
        QMessageBox.warning(self, "警告", "视频结束或无法读取帧")

def show_image(self, image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, ch = image.shape
    bytes_per_line = ch * w
    q_image = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888)
    pixmap = QPixmap.fromImage(q_image)
    self.label_image.setPixmap(pixmap)
    self.label_image.setScaledContents(True)

def closeEvent(self, event):
    if self.cap is not None:
        self.cap.release()
    self.timer.stop()
    event.accept()

if name == “main”:
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
8.2 utils.py
这个文件包含了预测图像、视频和摄像头的函数。

python
深色版本
import cv2
import numpy as np
from ultralytics import YOLO

定义类别名称

class_names = [‘牛’, ‘羊’, ‘马’]

预测图像

def predict_image(model, image):
results = model(image)
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = box.conf[0].item()
cls = int(box.cls[0].item())
label = f’{class_names[cls]} {conf:.2f}’
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return image

预测视频

def predict_video(model, video_path):
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
result_frame = predict_image(model, frame)
cv2.imshow(‘Grassland Animal Detection’, result_frame)
if cv2.waitKey(1) & 0xFF == ord(‘q’):
break
cap.release()
cv2.destroyAllWindows()

预测摄像头

def predict_camera(model):
cap = cv2.VideoCapture(0)
while True:
ret, frame = cap.read()
if not ret:
break
result_frame = predict_image(model, frame)
cv2.imshow(‘Grassland Animal Detection’, result_frame)
if cv2.waitKey(1) & 0xFF == ord(‘q’):
break
cap.release()
cv2.destroyAllWindows()
9. 运行项目
确保你的数据集和训练好的权重已经放在相应的文件夹中。
在项目根目录下运行以下命令启动训练:
bash
深色版本
python src/train.py
训练完成后,运行以下命令进行评估:
bash
深色版本
python ultralytics/yolo/v8/detect/val.py --data grassland_animal_dataset.yaml --weights runs/train/exp/weights/best.pt --task val
python ultralytics/yolo/v8/detect/val.py --data grassland_animal_dataset.yaml --weights runs/train/exp/weights/best.pt --task test
运行以下命令启动UI界面:
bash
深色版本
python src/main.py
10. 功能说明
打开图片:点击“打开图片”按钮,选择一张图片进行草原牛羊马检测,结果显示在UI界面上。
打开视频:点击“打开视频”按钮,选择一个视频文件进行草原牛羊马检测,结果显示在UI界面上。
打开摄像头:点击“打开摄像头”按钮,使用摄像头进行实时草原牛羊马检测,结果显示在UI界面上。
11. 详细注释
train.py
导入库:导入subprocess库用于调用命令行命令。
定义训练函数:定义一个函数train_model,用于训练YOLOv8模型。
训练模型:在主程序中调用train_model函数,训练YOLOv8模型。
main.py
导入库:导入必要的库,包括PyQt5、OpenCV、NumPy、ultralytics。
加载模型:使用YOLO类加载训练好的权重。
初始化UI:定义MainWindow类,继承自QMainWindow和Ui_MainWindow,并在构造函数中调用setupUi方法初始化UI。
连接信号和槽:连接按钮的点击事件到相应的槽函数。
打开图片:定义open_image方法,选择图片文件并进行草原牛羊马检测,结果显示在UI界面上。
打开视频:定义open_video方法,选择视频文件并进行草原牛羊马检测,结果显示在UI界面上。
打开摄像头:定义open_camera方法,使用摄像头进行实时草原牛羊马检测,结果显示在UI界面上。
更新帧:定义update_frame方法,定时读取摄像头帧并进行草原牛羊马检测,结果显示在UI界面上。
显示图像:定义show_image方法,将检测结果图像显示在UI界面上。
关闭事件:定义closeEvent方法,释放摄像头资源并停止定时器。
utils.py
定义类别名称:定义类别名称。
预测图像:定义一个函数predict_image,用于对单张图像进行草原牛羊马检测,并在图像上绘制检测框和标签。
预测视频:定义一个函数predict_video,用于对视频文件进行草原牛羊马检测,并在OpenCV窗口中显示结果。
预测摄像头:定义一个函数predict_camera,用于使用摄像头进行实时草原牛羊马检测,并在OpenCV窗口中显示结果。

标签:py,训练,self,YOLO,cv2,YOLOv8,train,dataset,image
From: https://blog.csdn.net/2401_88441190/article/details/143654641

相关文章