使用深度学习框架（如PyTorch/YOLO）训练指针式仪表数据集的目标检测和语义分割模型，并建立推理代码及可视化界面，训练指针式仪表分割数据集

使用深度学习框架（如PyTorch/YOLO）训练指针式仪表的目标检测和语义分割模型，并建立推理代码及可视化界面，训练指针式仪表数据集 检测和分割数据集

3000多张指针式仪表，目标检测，语义分割，数据集和标注文件

1

1

使用深度学习框架（如PyTorch或YOLO）训练指针式仪表的目标检测和语义分割模型，并建立推理代码及可视化界面。

---

任务概述

1. 目标：
   • 目标检测：识别仪表盘上的指针位置。
   • 语义分割：对仪表盘的各个部件进行像素级分类。
2. 数据集：
   • 3000多张图像，已标注。
   • 标注文件可能是边界框（目标检测）或像素掩码（语义分割）。
3. 工具：
   • 目标检测：YOLOv5/v8 或 PyTorch Faster R-CNN。
   • 语义分割：PyTorch DeepLabV3 或 MMSegmentation。
4. 可视化界面：使用Tkinter或Streamlit构建图形化界面。

---

步骤 1：环境准备

确保安装了必要的库。

# 安装依赖
pip install torch torchvision matplotlib opencv-python ultralytics streamlit

---

步骤 2：数据集处理

目标检测数据集

• 格式：YOLO需要每个图像对应一个txt文件，包含目标的类别和归一化的边界框坐标。
• 目录结构：

  dataset/
  ├── images/
  │   ├── train/
  │   ├── val/
  │   └── test/
  └── labels/
      ├── train/
      ├── val/
      └── test/
  

语义分割数据集

• 格式：每个图像对应一个掩码图像，掩码图像的像素值表示类别索引。
• 目录结构：

  dataset/
  ├── images/
  │   ├── train/
  │   ├── val/
  │   └── test/
  └── masks/
      ├── train/
      ├── val/
      └── test/
  

---

步骤 3：目标检测训练代码

YOLOv8 训练代码

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8n.pt')  # 使用YOLOv8 Nano模型

# 配置数据集
data_yaml = 'dataset.yaml'

# 开始训练
results = model.train(
    data=data_yaml,
    epochs=50,       # 训练轮数
    imgsz=640,       # 输入图像尺寸
    batch=16,        # 批量大小
    device='cuda'    # 使用GPU
)

# 验证模型
metrics = model.val()

# 保存模型
model.export(format='onnx')  # 导出为ONNX格式

自定义目标检测模型（PyTorch Faster R-CNN）

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from PIL import Image
import os

# 自定义数据集类
class PointerDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms=None):
        self.root = root
        self.transforms = transforms
        self.images = list(sorted(os.listdir(os.path.join(root, "images"))))
        self.labels = list(sorted(os.listdir(os.path.join(root, "labels"))))

    def __getitem__(self, idx):
        img_path = os.path.join(self.root, "images", self.images[idx])
        label_path = os.path.join(self.root, "labels", self.labels[idx])

        image = Image.open(img_path).convert("RGB")
        with open(label_path, 'r') as f:
            boxes = []
            for line in f.readlines():
                class_id, x_center, y_center, width, height = map(float, line.strip().split())
                x_min = (x_center - width / 2) * image.width
                x_max = (x_center + width / 2) * image.width
                y_min = (y_center - height / 2) * image.height
                y_max = (y_center + height / 2) * image.height
                boxes.append([x_min, y_min, x_max, y_max])

        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((len(boxes),), dtype=torch.int64)  # 假设所有类别为1
        target = {"boxes": boxes, "labels": labels}

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

        return image, target

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

# 创建模型
def get_model(num_classes):
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

# 数据加载器
dataset = PointerDataset(root="dataset/train", transforms=ToTensor())
data_loader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=4)

# 训练模型
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = get_model(num_classes=5)  # 包括背景类
model.to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)

for epoch in range(10):  # 训练10轮
    model.train()
    for images, targets in data_loader:
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {losses.item()}")

---

步骤 4：语义分割训练代码

DeepLabV3 训练代码

import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from PIL import Image
import os

# 自定义数据集类
class SegmentationDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = os.listdir(image_dir)

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.images[idx].replace(".jpg", ".png"))

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

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

        return image, mask

# 数据加载器
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])
dataset = SegmentationDataset(image_dir="dataset/images/train", mask_dir="dataset/masks/train", transform=transform)
data_loader = DataLoader(dataset, batch_size=8, shuffle=True)

# 创建模型
model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True)
model.classifier[-1] = nn.Conv2d(256, 5, kernel_size=(1, 1))  # 修改输出类别数
model.to(torch.device('cuda'))

# 训练模型
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    model.train()
    for images, masks in data_loader:
        images, masks = images.to('cuda'), masks.to('cuda')

        outputs = model(images)['out']
        loss = criterion(outputs, masks.long())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

---

步骤 5：推理代码与可视化界面

推理代码

import cv2
import torch
from ultralytics import YOLO

# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')

# 推理
image = cv2.imread('test_image.jpg')
results = model(image)

# 绘制结果
for box in results[0].boxes:
    x1, y1, x2, y2 = map(int, box.xyxy[0])
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

cv2.imshow("Result", image)
cv2.waitKey(0)

可视化界面

使用Streamlit创建简单的GUI：

import streamlit as st
from PIL import Image
import torch
from ultralytics import YOLO

# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')

# Streamlit应用
st.title("指针式仪表检测")
uploaded_file = st.file_uploader("上传一张图片", type=["jpg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    results = model(image)

    # 显示结果
    st.image(results[0].plot(), caption="检测结果", use_column_width=True)

---

总结

以上代码展示了如何使用YOLOv8和PyTorch完成目标检测和语义分割任务，并通过Streamlit构建可视化界面。您可以根据实际需求调整模型、数据集和界面设计。

仅供参考，