![在这里插入图片描述](https://developer.qcloudimg.com/http-save/yehe-12490569/f4a4fefc24ad1b9d4f4e963fa7a2a8c1.png)
钢材表面缺陷语义分割数据集 运行代码
总共3630张数据（jpg）和掩码（png）
包含三种缺陷类型（像素标签）
0为背景 1为夹杂物 2为补丁 3为划痕
可通过脚本修改标签映射关系![在这里插入图片描述](https://developer.qcloudimg.com/http-save/yehe-12490569/df14052db0f43404e8fd37ef50166dbe.png)
![在这里插入图片描述](https://developer.qcloudimg.com/http-save/yehe-12490569/5183390eb41ac512ebcb343ffd5f27c2.png)
假设你有一个钢材表面缺陷语义分割数据集，包含3630张图片和对应的掩码，共有三种缺陷类型（夹杂物、补丁、划痕），并且可以通过脚本修改标签映射关系。我们将使用PyTorch和U-Net架构来实现这个任务。
以下是完整的项目代码，包括数据预处理、模型训练、评估和推理部分。
### 完整项目代码

```
import os
import sys
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QLabel, QPushButton
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QTimer
import cv2

# 数据集类
class SteelDefectDataset(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 = np.array(Image.open(img_path).convert("RGB"))
        mask = np.array(Image.open(mask_path).convert("L"), dtype=np.int64)
        
        if self.transform:
            augmentations = self.transform(image=image, mask=mask)
            image = augmentations["image"]
            mask = augmentations["mask"]
        
        return image, mask

# U-Net 模型
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

# 数据增强
transform = A.Compose([
    A.Resize(height=256, width=256),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.Rotate(limit=30, p=0.5),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

# 训练脚本
def train_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    image_dir = 'datasets/steel_defects/images'
    mask_dir = 'datasets/steel_defects/masks'

    images = os.listdir(image_dir)
    masks = os.listdir(mask_dir)

    train_images, val_images = train_test_split(images, test_size=0.2, random_state=42)
    val_images, test_images = train_test_split(val_images, test_size=0.5, random_state=42)

    train_dataset = SteelDefectDataset(image_dir, mask_dir, transform=transform)
    val_dataset = SteelDefectDataset(image_dir, mask_dir, transform=transform)
    test_dataset = SteelDefectDataset(image_dir, mask_dir, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

    model = UNet(n_channels=3, n_classes=4)
    model.to(device=device)

    optimizer = optim.Adam(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss()

    num_epochs = 50

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for i, (images, masks) in enumerate(train_loader):
            images = images.to(device=device)
            masks = masks.to(device=device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 10 == 9:
                print(f'[Epoch {epoch+1}, Batch {i+1}] loss: {running_loss / 10:.4f}')
                running_loss = 0.0

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for images, masks in val_loader:
                images = images.to(device=device)
                masks = masks.to(device=device)

                outputs = model(images)
                loss = criterion(outputs, masks)
                val_loss += loss.item()

        print(f'Validation Loss after Epoch {epoch+1}: {val_loss / len(val_loader):.4f}')

    torch.save(model.state_dict(), 'best_steel_defect_model.pth')

# 评估脚本
def evaluate_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    image_dir = 'datasets/steel_defects/images'
    mask_dir = 'datasets/steel_defects/masks'

    images = os.listdir(image_dir)
    masks = os.listdir(mask_dir)

    _, val_images = train_test_split(images, test_size=0.2, random_state=42)
    val_images, _ = train_test_split(val_images, test_size=0.5, random_state=42)

    val_dataset = SteelDefectDataset(image_dir, mask_dir, transform=transform)
    val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)

    model = UNet(n_channels=3, n_classes=4)
    model.load_state_dict(torch.load('best_steel_defect_model.pth'))
    model.to(device=device)
    model.eval()

    criterion = nn.CrossEntropyLoss()
    val_loss = 0.0

    with torch.no_grad():
        for images, masks in val_loader:
            images = images.to(device=device)
            masks = masks.to(device=device)

            outputs = model(images)
            loss = criterion(outputs, masks)
            val_loss += loss.item()

    print(f'Evaluation Loss: {val_loss / len(val_loader):.4f}')

# 推理脚本
class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("Steel Surface Defect Detection")
        self.setGeometry(100, 100, 800, 600)
        
        # 初始化模型
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = UNet(n_channels=3, n_classes=4)
        self.model.load_state_dict(torch.load('best_steel_defect_model.pth'))
        self.model.to(device=self.device)
        self.model.eval()
        
        # 设置类别名称
        self.class_names = ['背景', '夹杂物', '补丁', '划痕']
        
        # 创建界面元素
        self.label_display = QLabel(self)
        self.label_display.setGeometry(10, 10, 780, 400)
        
        self.button_select_image = QPushButton("Select Image", self)
        self.button_select_image.setGeometry(10, 420, 150, 30)
        self.button_select_image.clicked.connect(self.select_image)
        
        self.button_select_video = QPushButton("Select Video", self)
        self.button_select_video.setGeometry(170, 420, 150, 30)
        self.button_select_video.clicked.connect(self.select_video)
        
        self.button_start_camera = QPushButton("Start Camera", self)
        self.button_start_camera.setGeometry(330, 420, 150, 30)
        self.button_start_camera.clicked.connect(self.start_camera)
        
        self.button_stop_camera = QPushButton("Stop Camera", self)
        self.button_stop_camera.setGeometry(490, 420, 150, 30)
        self.button_stop_camera.clicked.connect(self.stop_camera)
        
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        
        self.cap = None
        self.results = []

    def select_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.jpg *.jpeg *.png);;All Files (*)", options=options)
        if file_path:
            self.process_image(file_path)

    def process_image(self, image_path):
        frame = cv2.imread(image_path)
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        transformed = transform(image=frame_rgb, mask=np.zeros_like(frame_rgb[:,:,0]))
        input_tensor = transformed['image'].unsqueeze(0).to(device=self.device)
        output = self.model(input_tensor)
        pred_mask = output.argmax(dim=1).squeeze().cpu().numpy()
        annotated_frame = self.draw_annotations(frame, pred_mask)
        self.display_image(annotated_frame)
        self.results.append((image_path, annotated_frame))

    def select_video(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(self, "选择视频", "", "Videos (*.mp4 *.avi);;All Files (*)", options=options)
        if file_path:
            self.process_video(file_path)

    def process_video(self, video_path):
        self.cap = cv2.VideoCapture(video_path)
        while self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            transformed = transform(image=frame_rgb, mask=np.zeros_like(frame_rgb[:,:,0]))
            input_tensor = transformed['image'].unsqueeze(0).to(device=self.device)
            output = self.model(input_tensor)
            pred_mask = output.argmax(dim=1).squeeze().cpu().numpy()
            annotated_frame = self.draw_annotations(frame, pred_mask)
            self.display_image(annotated_frame)
            self.results.append((video_path, annotated_frame))
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        self.cap.release()

    def start_camera(self):
        self.cap = cv2.VideoCapture(0)
        self.timer.start(30)

    def stop_camera(self):
        self.timer.stop()
        if self.cap is not None:
            self.cap.release()
            self.label_display.clear()

    def update_frame(self):
        ret, frame = self.cap.read()
        if not ret:
            return
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        transformed = transform(image=frame_rgb, mask=np.zeros_like(frame_rgb[:,:,0]))
        input_tensor = transformed['image'].unsqueeze(0).to(device=self.device)
        output = self.model(input_tensor)
        pred_mask = output.argmax(dim=1).squeeze().cpu().numpy()
        annotated_frame = self.draw_annotations(frame, pred_mask)
        self.display_image(annotated_frame)
        self.results.append(('camera', annotated_frame))

    def draw_annotations(self, frame, mask):
        colors = [(0, 0, 0), (255, 0, 0), (0, 255, 0), (0, 0, 255)]  # 背景, 夹杂物, 补丁, 划痕
        h, w = mask.shape
        overlay = np.zeros((h, w, 3), dtype=np.uint8)
        for cls_id in range(1, 4):
            overlay[mask == cls_id] = colors[cls_id]
        opacity = 0.5
        cv2.addWeighted(overlay, opacity, frame, 1 - opacity, 0, frame)
        return frame

    def display_image(self, frame):
        rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qt_image)
        self.label_display.setPixmap(pixmap.scaled(self.label_display.width(), self.label_display.height()))

if __name__ == "__main__":
    # 运行模式选项
    mode = input("Enter mode (train, evaluate, infer): ")

    if mode == "train":
        train_model()
    
    elif mode == "evaluate":
        evaluate_model()
    
    elif mode == "infer":
        app = QApplication(sys.argv)
        window = MainWindow()
        window.show()
        sys.exit(app.exec_())
    
    else:
        print("Invalid mode. Please choose from 'train', 'evaluate', 'infer'.")
```
![在这里插入图片描述](https://developer.qcloudimg.com/http-save/yehe-12490569/166053291b12613825bed06c566480b9.png)
### 文件结构

```
steel_surface_defect_detection/
├── main.py
├── datasets/
│   └── steel_defects/
│       ├── images/
│       │   ├── image1.jpg
│       │   ├── image2.jpg
│       │   └── ...
│       └── masks/
│           ├── image1.png
│           ├── image2.png
│           └── ...
├── requirements.txt
└── data.yaml
```
### 安装依赖项

首先，确保你已经安装了所有必要的依赖项。你可以通过以下命令安装：
```
pip install -r requirements.txt
```
`requirements.txt` 内容如下：
```
opencv-python
torch==1.9
albumentations
matplotlib
PyQt5
```
### 运行步骤总结

1. **克隆项目仓库**（如果有的话）：
	```
	git clone https://github.com/yourusername/steel_surface_defect_detection.git
	cd steel_surface_defect_detection
	```
2. **安装依赖项**：
	```
	pip install -r requirements.txt
	```
3. **训练模型**：
	```
	python main.py train
	```
4. **评估模型**：
	```
	python main.py evaluate
	```
5. **运行推理界面**：
	```
	python main.py infer
	```
6. **操作界面**：
	- 选择图片进行检测。
	- 选择视频进行检测。
	- 使用摄像头进行实时检测。

### 总结

以上是完整的基于U-Net的钢材表面缺陷语义分割系统的项目介绍和代码实现。该项目支持图片检测、视频检测和调用摄像头进行实时检测，并提供了详细的注释和使用说明。希望这些详细的信息和代码能够帮助你顺利实施和优化你的钢材表面缺陷检测系统。

钢材表面缺陷语义分割数据集 运行代码总共3630张数据（jpg）和掩码（png）包含三种缺陷类型（像素标签）0为背景 1为夹杂物 2为补丁 3为划痕可通过脚本修改标签映射关系假设你有一个钢材表面缺陷语义分割数据集，包含3630张图片和对应的掩码，共有三种缺陷类型（夹杂物、补丁、划痕），并且可以通过脚本修改标签映射关系。我们将使用PyTorch和U-Net架构来实现这个任务。以下是完整的项目代码，包括数据预处理、模型训练、评估和推理部分。

使用PyTorch和U-Net架构来实现钢材表面缺陷语义分割检测 3630张图片和对应掩码，三种缺陷类型（夹杂物、补丁、划痕）

钢材表面缺陷语义分割数据集 运行代码总共3630张数据（jpg）和掩码（png）包含三种缺陷类型（像素标签）0为背景 1为夹杂物 2为补丁 3为划痕可通过脚本修改标签映射关系假设你有一个钢材表面缺陷语义分割数据集，包含3630张图片和对应的掩码，共有三种缺陷类型（夹杂物、补丁、划痕），并且可以通过脚本修改标签映射关系。我们将使用PyTorch和U-Net架构来实现这个任务。以下是完整的项目代码，包

人工智能

钢材表面缺陷检测系统采用U-Net深度学习模型，支持夹杂物、补丁和划痕三种缺陷的语义分割。包含3630张标注数据，提供训练、评估和推理功能，支持图片、视频和实时摄像头检测。完整项目代码包含数据预处理、模型训练和可视化界面。

数据预处理

代码解释

2026上云采购 | AI焕新·智启新局

datainlong

文章

问答

视频

教程

学习中心

腾讯云实验室

直播

竞赛

腾讯云代码分析专区

腾讯iOA零信任安全管理系统专区

腾讯云架构师技术同盟交流圈

腾讯云数据库专区

腾讯云智能顾问专区

腾讯云原生专区

腾讯混元专区

腾讯云TCE专区

腾讯云Lighthouse专区

腾讯云HAI专区

腾讯云Edgeone专区

腾讯云存储专区

腾讯云智能专区

腾讯轻联专区 

腾讯云开发专区

TAPD专区

腾讯轻量云游戏服专区

EdgeOne AI 安全实战专区

腾讯云最具价值专家

腾讯云架构师技术同盟

腾讯云创作之星

腾讯云开发者先锋

腾讯云代码助手

云原生构建

TAPD 敏捷项目管理

Cloud Studio

SDK中心

API中心

命令行工具

涵盖代码开发、场景应用、自动测试全流程，助你从零构建专属AI助手

一站式MCP教程库，解锁AI应用新玩法

聚焦“写作效率、视觉美观与运行性能”三方面进行全面升级，为您提供更高效、稳定的创作环境

社区富文本&Markdown编辑器全新改版上线，欢迎大家体验!

诚挚邀请您参与本次调研，分享您的真实使用感受与建议。您的反馈至关重要，感谢您的支持与参与！

社区新版编辑器体验调研

使用PyTorch和U-Net架构来实现钢材表面缺陷语义分割检测 3630张图片和对应掩码，三种缺陷类型（夹杂物、补丁、划痕）-腾讯云开发者社区-腾讯云

使用PyTorch和U-Net架构来实现钢材表面缺陷语义分割检测 3630张图片和对应掩码，三种缺陷类型（夹杂物、补丁、划痕）

使用PyTorch和U-Net架构来实现钢材表面缺陷语义分割检测 3630张图片和对应掩码，三种缺陷类型（夹杂物、补丁、划痕）

社区

活动

圈层

关于

腾讯云开发者

热门产品

热门推荐

更多推荐