电网人员安全带识别检测数据集统计表
标签名称 图片数量 边框数量
未系安全带 (Belt_off) 3773 -
系安全带 (Belt_on) 2517 -
总计 6290 -




项目结构
深色版本
safety_belt_detection/
├── dataset/
│ ├── images/
│ │ └── *.jpg
│ ├── labels/
│ │ └── *.txt
├── models/
│ └── yolov8/
├── src/
│ ├── train.py
│ ├── predict.py
│ ├── utils.py
├── weights/
│ └── best_model.pt
├── requirements.txt
└── README.md

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

深色版本
torch
torchvision
numpy
pandas
matplotlib
tqdm
pyyaml
opencv-python
ultralytics
然后，使用以下命令安装依赖：

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

深色版本
dataset/
├── images/
│ └── *.jpg
├── labels/
│ └── *.txt
每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件都是.jpg格式，标签文件是YOLO格式的.txt文件。

3. 数据集配置
   创建一个数据集类，用于加载和预处理数据。

3.1 src/utils.py
python
深色版本
import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image

class SafetyBeltDetectionDataset(Dataset):
def init(self, image_dir, label_dir, transform=None):
self.image_dir = image_dir
self.label_dir = label_dir
self.transform = transform
self.image_files = os.listdir(image_dir)
self.label_files = os.listdir(label_dir)

def __len__(self):
    return len(self.image_files)

def __getitem__(self, index):
    img_path = os.path.join(self.image_dir, self.image_files[index])
    label_path = os.path.join(self.label_dir, self.image_files[index].replace('.jpg', '.txt'))

    image = Image.open(img_path).convert("RGB")

    if os.path.exists(label_path):
        with open(label_path, 'r') as f:
            labels = f.readlines()
        labels = [line.strip().split() for line in labels]
        labels = [[int(label[0])] + list(map(float, label[1:])) for label in labels]
    else:
        labels = []

    if self.transform:
        image = self.transform(image)

    return image, labels

def get_data_loaders(image_dir, label_dir, batch_size=16, num_workers=4):
transform = transforms.Compose([
transforms.Resize((640, 640)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

dataset = SafetyBeltDetectionDataset(image_dir, label_dir, transform=transform)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size

train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

return train_loader, val_loader

4. 模型定义
   使用YOLOv8模型进行检测任务。这里我们使用预训练的YOLOv8模型，并对其进行微调。

4.1 src/train.py
python
深色版本
import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from ultralytics import YOLO
from src.utils import get_data_loaders
import torch.nn.functional as F

def train_model(image_dir, label_dir, epochs=100, batch_size=16, learning_rate=1e-4):
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

model = YOLO('yolov8n.pt')  # Load a pretrained YOLOv8 model
model = model.to(device)

train_loader, val_loader = get_data_loaders(image_dir, label_dir, batch_size=batch_size)

optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.MSELoss()

writer = SummaryWriter()

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}"):
        images = images.to(device)
        labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]

        optimizer.zero_grad()
        outputs = model(images)
        loss = 0.0
        for i in range(len(labels)):
            if len(labels[i]) > 0:
                pred = outputs[i, :, :, :5 * len(labels[i])]
                pred = pred.permute(2, 0, 1).reshape(-1, 5)
                target = labels[i][:, 1:]
                loss += criterion(pred, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    train_loss = running_loss / len(train_loader)
    writer.add_scalar('Training Loss', train_loss, epoch)

    model.eval()
    running_val_loss = 0.0
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]

            outputs = model(images)
            loss = 0.0
            for i in range(len(labels)):
                if len(labels[i]) > 0:
                    pred = outputs[i, :, :, :5 * len(labels[i])]
                    pred = pred.permute(2, 0, 1).reshape(-1, 5)
                    target = labels[i][:, 1:]
                    loss += criterion(pred, target)

            running_val_loss += loss.item()

    val_loss = running_val_loss / len(val_loader)
    writer.add_scalar('Validation Loss', val_loss, epoch)

    print(f"Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

model.save("weights/best_model.pt")
writer.close()

if name == “main”:
image_dir = “dataset/images”
label_dir = “dataset/labels”
train_model(image_dir, label_dir)
5. 模型评估
训练完成后，可以通过验证集和测试集来评估模型的性能。示例如下：

5.1 src/predict.py
python
深色版本
import torch
import matplotlib.pyplot as plt
from ultralytics import YOLO
from src.utils import get_data_loaders
import numpy as np

def predict_and_plot(image_dir, label_dir, model_path, num_samples=5):
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

model = YOLO(model_path)
model = model.to(device)
model.eval()

_, val_loader = get_data_loaders(image_dir, label_dir)

fig, axes = plt.subplots(num_samples, 2, figsize=(10, 5 * num_samples))
with torch.no_grad():
    for i, (images, labels) in enumerate(val_loader):
        if i >= num_samples:
            break

        images = images.to(device)
        labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]

        outputs = model(images)
        outputs = outputs.cpu().numpy()

        images = images.cpu().numpy().transpose((0, 2, 3, 1))
        labels = [label.cpu().numpy() for label in labels]

        for j in range(len(images)):
            ax = axes[j] if num_samples > 1 else axes
            ax[0].imshow(images[j])
            ax[0].set_title(f"True: {labels[j]}")
            ax[0].axis('off')

            ax[1].imshow(images[j])
            ax[1].set_title(f"Predicted: {outputs[j]}")
            ax[1].axis('off')

plt.tight_layout()
plt.show()

if name == “main”:
image_dir = “dataset/images”
label_dir = “dataset/labels”
model_path = “weights/best_model.pt”
predict_and_plot(image_dir, label_dir, model_path)
6. 运行项目
确保你的数据集已经放在相应的文件夹中。
在项目根目录下运行以下命令启动训练：
bash
深色版本
python src/train.py
训练完成后，运行以下命令进行评估和可视化：
bash
深色版本
python src/predict.py
7. 功能说明
数据集类：SafetyBeltDetectionDataset类用于加载和预处理数据。
数据加载器：get_data_loaders函数用于创建训练和验证数据加载器。
训练模型：train.py脚本用于训练YOLOv8模型，使用均方误差损失函数和Adam优化器。
评估模型：predict.py脚本用于评估模型性能，并可视化输入图像、真实标签和预测结果。
8. 详细注释
utils.py
数据集类：定义了一个SafetyBeltDetectionDataset类，用于加载和预处理数据。
数据加载器：定义了一个get_data_loaders函数，用于创建训练和验证数据加载器。
train.py
训练函数：定义了一个train_model函数，用于训练YOLOv8模型。
训练过程：在每个epoch中，模型在训练集上进行前向传播和反向传播，并在验证集上进行评估。
predict.py
预测和可视化：定义了一个predict_and_plot函数，用于在验证集上进行预测，并可视化输入图像、真实标签和预测结果。