当前位置：首页 > news >正文

使用 SAM + YOLO + ResNet 检测工业开关状态：从零到部署

news 来源：原创 2025/6/8 9:29:28

使用 SAM + YOLO + ResNet 检测工业开关状态：从零到部署

在工业自动化和智能制造领域，实时监控设备状态是提高生产效率和安全性的重要环节。其中，工业开关的状态检测是一个常见的需求。传统的人工检查不仅耗时耗力，还容易出错。因此，利用计算机视觉技术实现自动化检测显得尤为重要。

本文将介绍如何结合 SAM（Spatial Attention Module）、YOLO（You Only Look Once） 和 ResNet（Residual Network），实现一个高效、准确的工业开关状态检测系统。我们将从数据准备、模型训练到最终的实时检测部署，逐步展开。

一、技术选型

（一）YOLO：目标检测

YOLO 是一种流行的实时目标检测算法，以其速度快、精度高而闻名。它能够快速定位图像中的目标物体，并返回其位置信息。在我们的场景中，YOLO 将用于检测工业开关在图像中的位置。

（二）ResNet：状态分类

ResNet 是一种深度卷积神经网络，广泛应用于图像分类任务。它能够有效提取图像的特征，并对开关的状态进行分类（开闸、合闸、分闸、中间状态）。

（三）SAM：空间注意力机制

SAM（Spatial Attention Module）是一种空间注意力机制模块，能够增强模型对目标的空间特征提取能力。通过引入 SAM，我们可以进一步提升模型对开关状态的识别精度。

二、数据准备

（一）数据采集

首先，我们需要采集工业开关的图像数据。可以使用工业摄像头，从不同角度和距离拍摄开关的图像。确保采集的图像涵盖了开关的四种状态：开闸、合闸、分闸和中间状态。

（二）数据标注

使用标注工具（如 LabelImg）对采集到的图像进行标注。标注时，需要框出开关的位置，并标记其状态（开闸、合闸、分闸、中间状态）。标注完成后，将标注数据转换为 YOLO 格式的 TXT 文件。

import xml.etree.ElementTree as ET
import osdef convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]return (x * dw, y * dh, w * dw, h * dh)def convert_annotation(xml_file, txt_file):tree = ET.parse(xml_file)root = tree.getroot()size = root.find('size')width = int(size.find('width').text)height = int(size.find('height').text)with open(txt_file, 'w') as out_file:for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))bb = convert((width, height), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')classes = ["open", "close", "middle"]
xml_files = os.listdir('annotations')
for xml_file in xml_files:txt_file = xml_file.replace('.xml', '.txt')convert_annotation(os.path.join('annotations', xml_file), os.path.join('labels', txt_file))

三、模型训练

（一）YOLO 模型训练

我们将使用 YOLOv8 进行目标检测。YOLOv8 是目前最新的版本，具有更高的检测精度和速度。

准备数据集：将标注好的数据集分为训练集和验证集。
修改配置文件：根据我们的数据集修改 YOLO 的配置文件，包括类别数和路径。
训练模型：使用 YOLO 的训练脚本进行训练。

from ultralytics import YOLO# 加载预训练模型
model = YOLO('yolov8n.yaml')
model = YOLO('yolov8n.pt')# 修改模型配置文件
with open('yolov8n.yaml', 'r') as f:config = f.read()config = config.replace('nc: 80', 'nc: 4')  # 修改类别数为 4
with open('yolov8n.yaml', 'w') as f:f.write(config)# 训练模型
results = model.train(data='data.yaml', epochs=100, imgsz=640)

（二）ResNet 模型训练

我们将使用 ResNet50 进行状态分类。ResNet50 是一个经典的深度卷积神经网络，适合图像分类任务。

准备数据集：将标注好的图像数据分为训练集和验证集。
数据预处理：对图像进行缩放、归一化等预处理操作。
训练模型：使用 PyTorch 的训练脚本进行训练。

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import osclass CustomDataset(Dataset):def __init__(self, root_dir, transform=None):self.root_dir = root_dirself.transform = transformself.images = os.listdir(root_dir)def __len__(self):return len(self.images)def __getitem__(self, idx):img_name = os.path.join(self.root_dir, self.images[idx])image = Image.open(img_name)label = int(self.images[idx].split('_')[0])  # 假设文件名格式为 "label_image.jpg"if self.transform:image = self.transform(image)return image, labeltransform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])dataset = CustomDataset(root_dir='images', transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 4)  # 修改分类层为 4 类criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)for epoch in range(100):  # 训练 100 个 epochmodel.train()running_loss = 0.0for inputs, labels in dataloader:inputs, labels = inputs.to(device), labels.to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item() * inputs.size(0)epoch_loss = running_loss / len(dataset)print(f'Epoch {epoch+1}/100, Loss: {epoch_loss:.4f}')

四、模型部署

（一）实时检测

我们将训练好的 YOLO 模型和 ResNet 模型部署到摄像头设备上，实现实时检测开关状态。

import cv2
import numpy as np
import torch
import torchvision.transforms as transforms# 加载 YOLO 模型
yolo_model = YOLO('best.pt')# 加载 ResNet 模型
resnet_model = torch.load('resnet.pth')
resnet_model.eval()# 摄像头捕获
cap = cv2.VideoCapture(0)while True:ret, frame = cap.read()if not ret:break# YOLO 检测results = yolo_model.predict(frame)for result in results.xyxy[0].numpy():x1, y1, x2, y2, conf, cls = resultif conf > 0.5:x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)# ResNet 分类img = frame[y1:y2, x1:x2]img = cv2.resize(img, (224, 224))img = transforms.ToTensor()(img)img = img.unsqueeze(0).to(device)output = resnet_model(img)_, predicted = torch.max(output, 1)label = classes[predicted.item()]cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)cv2.imshow('frame', frame)if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()
cv2.destroyAllWindows()