深度学习目标检测中使用YOLOv8训练树冠检测数据集,从环境设置、数据准备、模型训练、推理和结果可视化
树冠检测数据集的训练及推理
1
使用YOLOv8训练树冠检测数据集,从环境设置、数据准备、模型训练、推理和结果可视化等方面进行详细介绍。
首先确保安装了Ultralytics提供的YOLOv8及其依赖项:
pip install ultralytics
假设您的数据集结构如下:
images/
文件夹包含所有图像。labels/
文件夹包含对应的YOLO格式的txt标签文件(每个对象一行,格式为:
)。创建一个YOLO配置文件data.yaml
来描述您的数据集:
train: ./path/to/train/images
val: ./path/to/val/images
test: ./path/to/test/images # 可选
nc: 1 # 类别数量,这里以树冠为例,可能有更多类别
names: ['canopy'] # 类别名称,根据实际情况调整
使用YOLOv8进行训练的代码示例如下:
from ultralytics import YOLO
# 加载预训练的YOLOv8模型或自定义模型
model = YOLO('yolov8n.yaml') # 或者 'yolov8n.pt' 来加载预训练权重
# 开始训练
results = model.train(data='data.yaml', epochs=100, imgsz=640, batch=16)
# 保存模型
model.save("best_model.pt")
完成训练后,您可以使用训练好的模型对新的图像进行预测,并可视化结果:
import cv2
import matplotlib.pyplot as plt
def visualize_predictions(model, image_path):
results = model(image_path)
boxes = results[0].boxes.xyxy.cpu().numpy()
confidences = results[0].boxes.conf.cpu().numpy()
class_ids = results[0].boxes.cls.cpu().numpy()
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for box, conf, cls in zip(boxes, confidences, class_ids):
x_min, y_min, x_max, y_max = map(int, box)
label = model.model.names[int(cls)]
cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
cv2.putText(img, f'{label}: {conf:.2f}', (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
plt.figure(figsize=(10, 10))
plt.imshow(img)
plt.axis('off')
plt.show()
# 示例:加载一张图像进行预测
image_path = './path/to/test/image.jpg'
visualize_predictions(YOLO('best_model.pt'), image_path)
为了提供更友好的用户体验,可以构建一个简单的PyQt5界面来选择图片并显示预测结果。
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QFileDialog
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
class CanopyDetectionApp(QMainWindow):
def __init__(self):
super().__init__()
self.initUI()
def initUI(self):
self.setWindowTitle('树冠检测系统')
self.setGeometry(100, 100, 800, 600)
self.image_label = QLabel(self)
self.layout = QVBoxLayout()
self.layout.addWidget(self.image_label)
self.button_open_image = QPushButton('打开图片', self)
self.button_open_image.clicked.connect(self.open_image)
self.layout.addWidget(self.button_open_image)
container = QWidget()
container.setLayout(self.layout)
self.setCentralWidget(container)
def open_image(self):
fname, _ = QFileDialog.getOpenFileName(self, '打开图片', '', "Image files (*.jpg *.png)")
if fname:
self.display_image(fname)
def display_image(self, fname):
model = YOLO('best_model.pt')
results = model(fname)
boxes = results[0].boxes.xyxy.cpu().numpy()
confidences = results[0].boxes.conf.cpu().numpy()
class_ids = results[0].boxes.cls.cpu().numpy()
img = cv2.imread(fname)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for box, conf, cls in zip(boxes, confidences, class_ids):
x_min, y_min, x_max, y_max = map(int, box)
label = model.model.names[int(cls)]
cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
cv2.putText(img, f'{label}: {conf:.2f}', (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
h, w, ch = img.shape
bytes_per_line = ch * w
q_img = QPixmap.fromImage(QImage(img.data, w, h, bytes_per_line, QImage.Format_RGB888))
self.image_label.setPixmap(q_img.scaled(self.image_label.size(), Qt.KeepAspectRatio))
if __name__ == '__main__':
app = QApplication(sys.argv)
ex = CanopyDetectionApp()
ex.show()
sys.exit(app.exec_())
使用YOLOv8高效地训练树冠检测数据集,并实现目标检测任务。