PyTorch 实现 深度学习语义分割模型unet训练 道路裂缝分割数据集，实现深度学习道路裂缝分割检测任务 语义分割。

道路的裂缝数据集，适合做深度学习语义分割。标签，掩膜，图像。250张

1

你提供的是一个用于语义分割的小型道路裂缝数据集，具有以下特点：

🧾 数据集概述

• 图像数量: 250张
• 图像类型: RGB 图像（如 .jpg `）
• 标签格式: PNG 格式的二值 mask 掩膜（单通道），表示裂缝区域（像素值为1）和非裂缝区域（像素值为0）
• 任务类型: 语义分割
• 用途: 道路裂缝检测、自动驾驶辅助系统、路面维护监控等

—

📦 下载 U-Net 示例代码（PyTorch 实现）

1. 自定义 Dataset 类

import os
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset

class CrackDataset(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])

        image = np.array(Image.open(img_path).convert("RGB"))
        mask = np.array(Image.open(mask_path).convert("L"), dtype=np.float32)
        mask[mask == 255.0] = 1.0  # 二值mask，转为0和1

        if self.transform:
            augmentations = self.transform(image=image, mask=mask)
            image = augmentations["image"]
            mask = augmentations["mask"]

        return image, mask.unsqueeze(0)  # shape: [C, H, W]

2. 构建 U-Net 模型（简化版）

import torch
import torch.nn as nn

def double_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, 3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

class UNet(nn.Module):
    def __init__(self, in_channels=3):
        super(UNet, self).__init__()
        
        self.down1 = double_conv(in_channels, 64)
        self.down2 = double_conv(64, 128)
        self.down3 = double_conv(128, 256)
        self.down4 = double_conv(256, 512)

        self.maxpool = nn.MaxPool2d(2)

        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.up_conv3 = double_conv(512, 256)

        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.up_conv2 = double_conv(256, 128)

        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.up_conv1 = double_conv(128, 64)

        self.final_conv = nn.Conv2d(64, 1, kernel_size=1)

    def forward(self, x):
        # 下采样路径
        conv1 = self.down1(x)
        x = self.maxpool(conv1)

        conv2 = self.down2(x)
        x = self.maxpool(conv2)

        conv3 = self.down3(x)
        x = self.maxpool(conv3)

        x = self.down4(x)

        # 上采样路径
        x = self.up3(x)
        x = torch.cat([x, conv3], dim=1)
        x = self.up_conv3(x)

        x = self.up2(x)
        x = torch.cat([x, conv2], dim=1)
        x = self.up_conv2(x)

        x = self.up1(x)
        x = torch.cat([x, conv1], dim=1)
        x = self.up_conv1(x)

        return torch.sigmoid(self.final_conv(x))

3. 训练脚本

from torch.utils.data import DataLoader
from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2

# 数据增强与加载
transform = A.Compose([
    A.Resize(height=256, width=256),
    A.ToFloat(max_value=255),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2()
])

train_dataset = CrackDataset("images/train", "masks/train", transform=transform)
val_dataset = CrackDataset("images/val", "masks/val", transform=transform)

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

# 初始化模型、损失函数、优化器
device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet().to(device)
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 训练循环
for epoch in range(50):  # 训练50轮
    model.train()
    for images, masks in train_loader:
        images = images.to(device)
        masks = masks.to(device)

        outputs = model(images)
        loss = criterion(outputs, masks)

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

    print(f"Epoch [{epoch+1}/50], Loss: {loss.item():.4f}")

    # 验证阶段（略）

4. 推理 & 可视化

import matplotlib.pyplot as plt

model.eval()
with torch.no_grad():
    images, masks = next(iter(val_loader))
    outputs = model(images.to(device)).cpu().numpy()

    for i in range(len(outputs)):
        plt.figure(figsize=(10, 5))
        plt.subplot(1, 3, 1)
        plt.title("Image")
        plt.imshow(images[i].permute(1, 2, 0).cpu().numpy())

        plt.subplot(1, 3, 2)
        plt.title("Mask")
        plt.imshow(masks[i][0].cpu().numpy(), cmap='gray')

        plt.subplot(1, 3, 3)
        plt.title("Prediction")
        plt.imshow(outputs[i][0] > 0.5, cmap='gray')

        plt.show()

📊 模型评估指标（语义分割）

from sklearn.metrics import jaccard_score, f1_score

def calculate_metrics(preds, targets):
    preds = preds.flatten()
    targets = targets.flatten()
    iou = jaccard_score(targets, preds)
    dice = f1_score(targets, preds)
    return {"IoU": iou, "Dice": dice}

📦 模型导出（ONNX / TorchScript）

dummy_input = torch.randn(1, 3, 256, 256).to(device)
torch.onnx.export(model, dummy_input, "unet_crack.onnx", export_params=True)

以上文字及代码仅供参考学习。