使用PyTorch和YOLOv8的官方实现来进行训练—粉尘识别数据集。工地/矿下粉尘数据识别，数据集已划分，YOLO格式-有权重，相关指数，map相当高 2000多张

使用PyTorch和YOLOv8的官方实现来进行训练，并介绍数据集的准备和使用方法。

一、数据集介绍

数据集结构
总共有2000多张图片和对应的txt标签文件。
标签类别：假设只有一类“dust”（粉尘）。
目录结构
深色版本
DustDetection/
├── data/
│ ├── images/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
├── models/
│ └── yolov8/
├── utils/
│ ├── data_loader.py
│ ├── metrics.py
│ └── plot.py
├── main.py
├── train.py
├── infer.py
└── README.md
二、数据集准备

1. 数据配置文件
   创建一个数据配置文件data/dust.yaml，指定数据集路径和类别信息：

train: …/data/images/train/
val: …/data/images/val/
test: …/data/images/test/

nc: 1 # 类别数量
names:

• dust
  三、安装YOLOv8
  首先，克隆YOLOv8仓库并安装必要的依赖项：

bash
深色版本
cd models/
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip install -r requirements.txt
四、训练模型
编写一个训练脚本train.py：

python
深色版本
import os
from ultralytics import YOLO

def main():
# 设置训练参数
data_yaml = ‘…/data/dust.yaml’
model_yaml = ‘yolov8n.yaml’ # 选择合适的YOLOv8模型
initial_weights = ‘path/to/initial/weights.pt’ # 初始权重文件
epochs = 100
batch_size = 16
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

# 加载模型
model = YOLO(model_yaml)
model.load(initial_weights)  # 加载初始权重

# 开始训练
results = model.train(
    data=data_yaml,
    epochs=epochs,
    batch=batch_size,
    imgsz=img_size,
    device=device
)

if name == ‘main’:
main()
运行训练脚本：

bash
深色版本
python train.py
五、推理和可视化
编写一个推理脚本infer.py，用于加载训练好的模型并对新图像进行预测：

python
深色版本
import os
import torch
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.yolo.utils.general import non_max_suppression, scale_coords
from ultralytics.yolo.utils.plots import Annotator

def load_model(weights_path, device):
model = YOLO(weights_path)
model.to(device)
return model

def infer_image(model, img_path, img_size, device):
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (img_size, img_size))
img = img.transpose(2, 0, 1)
img = img.astype(np.float32) / 255.0
img = torch.from_numpy(img).unsqueeze(0).to(device)

with torch.no_grad():
    pred = model(img)[0]

pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

return pred, img

def draw_boxes(img, pred, names):
annotator = Annotator(img, line_width=2, example=str(names))
for i, det in enumerate(pred):
if len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[2:]).round()
for *xyxy, conf, cls in reversed(det):
c = int(cls)
label = f’{names[c]} {conf:.2f}’
annotator.box_label(xyxy, label, color=(0, 255, 0))
return annotator.result()

if name == ‘main’:
weights_path = ‘runs/train/exp/weights/best.pt’
img_path = ‘…/data/images/test/000001.jpg’
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

model = load_model(weights_path, device)
pred, img = infer_image(model, img_path, img_size, device)

# 读取原始图像
original_img = cv2.imread(img_path)

# 绘制预测框
annotated_img = draw_boxes(original_img, pred, names=['dust'])

# 显示结果
cv2.imshow('Result', annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

六、评价指标
编写一个脚本evaluate.py来计算模型的评价指标（如IoU、mAP等）：

python
深色版本
import os
import torch
from ultralytics import YOLO
from ultralytics.yolo.utils.general import check_img_size
from ultralytics.yolo.utils.datasets import create_dataloader

def main():
data_yaml = ‘…/data/dust.yaml’
weights_path = ‘runs/train/exp/weights/best.pt’
img_size = 640
batch_size = 16
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’

# 加载数据
dataloader = create_dataloader('../data/images/val/', img_size, batch_size, 32, single_cls=False, pad=0.5, rect=True)[0]

# 加载模型
model = YOLO(weights_path)
model.to(device)

# 验证模型
results = model.val(
    data=data_yaml,
    batch=batch_size,
    imgsz=img_size,
    device=device,
    dataloader=dataloader
)

if name == ‘main’:
main()
运行评价脚本：

bash
深色版本
python evaluate.py
七、运行项目
确保数据集和标签文件已经准备好，并放在相应的目录中。
运行训练脚本：
bash
深色版本
python train.py
运行推理脚本：
bash
深色版本
python infer.py
运行评价脚本：
bash
深色版本
python evaluate.py
八、代码说明
数据配置文件：data/dust.yaml指定了数据集的路径和类别信息。
训练模型：train.py负责加载数据、初始化模型、损失函数和优化器，并进行训练。
推理和可视化：infer.py用于加载训练好的模型并对新图像进行预测，显示预测结果。
评价指标：evaluate.py计算模型的评价指标，如IoU和mAP。
希望这些代码和说明能帮助你完成基于YOLOv8的工地粉尘识别项目。