基于 PyTorch 实现训练滑坡语义分割数据集 UNet 语义分割模型训练滑坡分割数据集

DeepLearningYolo

277人浏览 · 2026-05-20 10:51:10

DeepLearningYolo · 2026-05-20 10:51:10 发布

文章目录

数据集描述：
**超大规模滑坡分割数据集，共34604张图像，全部处理为512×512尺度，涵盖航空，无人机，卫星多种传感器类型，涵盖0.1m-10m不同分辨率影像，3.6GB，标注质量很高，且全部为原始数据，无增广。**的训练及应用。在这里插入图片描述
1

1

滑坡语义分割数据集，共计 34604 张图像，每张大小为 512×512，包含多种传感器来源（航空、无人机、卫星），分辨率范围广泛（0.1m - 10m），总大小约 3.6GB。标注为原始数据，无增广处理，非常适合用于训练 UNet 语义分割模型。

🧱 项目目标

使用 UNet 模型 对滑坡进行像素级语义分割
基于 PyTorch 实现训练流程
构建一套 滑坡检测与可视化系统
支持单图推理、视频流或遥感影像批量推理

🛠️ 第一步：环境搭建

1. 安装 CUDA 驱动

确保你的显卡驱动已安装最新版本的 NVIDIA 驱动，并支持当前的 CUDA 版本（推荐 CUDA 11.8）：

nvidia-smi

2. 安装 Anaconda（如未安装）

从官网下载并安装 Anaconda

3. 创建虚拟环境

conda create --name unet-seg python=3.9
conda activate unet-seg

4. 安装依赖项

pip install torch torchvision torchaudio
pip install matplotlib numpy opencv-python tqdm scikit-learn
pip install albumentations  # 数据增强库（可选）
pip install tensorboard     # 可视化训练过程（可选）

📁 第二步：数据集结构整理（YOLO格式转UNet输入）

YOLO 格式通常如下：

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

但 UNet 要求的是图像 + mask 的配对形式，因此你需要将 YOLO 标注转换为灰度图 mask 图像。

✅ 将 YOLO 标注文件 `.txt` 转换为 mask 图像

你可以使用下面这个脚本将 .txt 文件转为对应尺寸的 mask 图像（每个像素值为 0 或 1）：

示例代码：`convert_yolo_to_mask.py`

import os
import cv2
import numpy as np

IMAGE_SIZE = (512, 512)

def yolo_to_mask(txt_path, img_size):
    h, w = img_size
    mask = np.zeros((h, w), dtype=np.uint8)
    with open(txt_path, 'r') as f:
        lines = f.readlines()
        for line in lines:
            data = list(map(float, line.strip().split()))
            class_id = int(data[0])
            coords = np.array(data[1:]).reshape(-1, 2)
            coords[:, 0] *= w
            coords[:, 1] *= h
            coords = np.round(coords).astype(np.int32)
            cv2.fillPoly(mask, [coords], color=1)  # 类别为1
    return mask

# 输入输出路径
image_dir = "dataset/images/train"
label_dir = "dataset/labels/train"
output_mask_dir = "dataset/masks/train"

os.makedirs(output_mask_dir, exist_ok=True)

for label_file in os.listdir(label_dir):
    base_name = os.path.splitext(label_file)[0]
    image_path = os.path.join(image_dir, base_name + ".jpg")
    label_path = os.path.join(label_dir, label_file)

    if not os.path.exists(image_path):
        continue

    mask = yolo_to_mask(label_path, IMAGE_SIZE)
    cv2.imwrite(os.path.join(output_mask_dir, base_name + ".png"), mask * 255)

执行后你会得到：

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

🧪 第三步：构建 UNet 数据集和模型

1. 自定义 Dataset 类

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

class LandslideDataset(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"))  # 灰度图

        # 二值化 mask（0 or 1）
        mask[mask > 0] = 1

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

        image = transforms.ToTensor()(image)
        mask = torch.from_numpy(mask).long()

        return image, mask

2. 定义 UNet 模型（PyTorch官方实现简化版）

import torch.nn as nn
import torch.nn.functional as F

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super(UNet, self).__init__()
        
        def CBR(in_channels, out_channels):
            return nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True)
            )

        self.enc1 = CBR(in_channels, 64)
        self.enc2 = CBR(64, 128)
        self.enc3 = CBR(128, 256)
        self.enc4 = CBR(256, 512)
        self.bottleneck = CBR(512, 1024)

        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.dec4 = CBR(1024, 512)
        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec3 = CBR(512, 256)
        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = CBR(256, 128)
        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = CBR(128, 64)

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

    def forward(self, x):
        enc1 = self.enc1(x)
        enc2 = self.enc2(F.max_pool2d(enc1, kernel_size=2, stride=2))
        enc3 = self.enc3(F.max_pool2d(enc2, kernel_size=2, stride=2))
        enc4 = self.enc4(F.max_pool2d(enc3, kernel_size=2, stride=2))

        bottleneck = self.bottleneck(F.max_pool2d(enc4, kernel_size=2, stride=2))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.dec4(dec4)

        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.dec3(dec3)

        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.dec2(dec2)

        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.dec1(dec1)

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

🏋️‍♂️ 第四步：训练模型

1. 训练脚本完整示例

from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm

# 设置参数
EPOCHS = 20
BATCH_SIZE = 8
LEARNING_RATE = 1e-4
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 加载数据集
train_dataset = LandslideDataset(
    image_dir="dataset/images/train",
    mask_dir="dataset/masks/train"
)
val_dataset = LandslideDataset(
    image_dir="dataset/images/val",
    mask_dir="dataset/masks/val"
)

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

# 初始化模型
model = UNet(in_channels=3, out_channels=1).to(DEVICE)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.BCELoss()  # 二分类交叉熵损失

# 开始训练
for epoch in range(EPOCHS):
    model.train()
    loop = tqdm(train_loader, total=len(train_loader))
    for images, masks in loop:
        images = images.to(DEVICE)
        masks = masks.float().unsqueeze(1).to(DEVICE)

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

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

        loop.set_postfix(loss=loss.item())

    # 验证阶段（可选）
    model.eval()
    with torch.no_grad():
        val_loss = 0
        for images, masks in val_loader:
            images = images.to(DEVICE)
            masks = masks.float().unsqueeze(1).to(DEVICE)
            preds = model(images)
            val_loss += criterion(preds, masks).item()
        print(f"Epoch {epoch+1} | Val Loss: {val_loss/len(val_loader):.4f}")

# 保存模型
torch.save(model.state_dict(), "unet_landslide.pth")

🧪 第五步：推理 & 构建滑坡检测系统

推理函数（单图测试）

def predict_image(model, image_path, device="cuda"):
    model.eval()
    image = Image.open(image_path).convert("RGB")
    image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(image_tensor)
        pred_mask = (output.squeeze().cpu().numpy() > 0.5).astype(np.uint8)

    # 显示结果
    import matplotlib.pyplot as plt
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title("Input Image")

    plt.subplot(1, 2, 2)
    plt.imshow(pred_mask, cmap='gray')
    plt.title("Predicted Mask")
    plt.show()

📦 第六步：部署 & 构建系统

你可以进一步构建以下功能：

功能	描述
Web 系统	使用 Flask / FastAPI 提供 API 接口
多图批处理	批量处理遥感图像，生成滑坡分布图
视频流检测	接入无人机视频流实时识别滑坡区域
地理坐标映射	结合 GIS 数据展示地理上的滑坡位置
报警机制	当检测到滑坡时自动报警