使用深度学习模型SMP对全球尺度下基于哨兵一号（Sentinel-1）与哨兵二号（Sentinel-2）卫星的水体分割数据集进行语义分割

本文链接：https://blog.csdn.net/2401_86822270/article/details/146094764

全球尺度下基于哨兵一号与哨兵二号卫星的水体分割数据集

文章目录

使用深度学习模型对全球尺度下基于哨兵一号（Sentinel-1）与哨兵二号（Sentinel-2）卫星的水体分割数据集进行语义分割，Segmentation Models for PyTorch (SMP)。这个框架支持多种流行的架构（如UNet、DeepLabV3等），并允许使用预训练权重和自定义损失函数。

1. 数据集准备

假设你的数据集结构如下：

global_water_segmentation/
├── sentinel1_images/
│   ├── train/
│   │   ├── img1.tif
│   │   └── ...
│   ├── val/
│   │   ├── img1.tif
│   │   └── ...
│   └── test/
│       ├── img1.tif
│       └── ...
├── sentinel2_images/
│   ├── train/
│   │   ├── img1.tif
│   │   └── ...
│   ├── val/
│   │   ├── img1.tif
│   │   └── ...
│   └── test/
│       ├── img1.tif
│       └── ...
└── masks/
    ├── train/
    │   ├── img1.png
    │   └── ...
    ├── val/
    │   ├── img1.png
    │   └── ...
    └── test/
        ├── img1.png
        └── ...
data_water_segmentation.yaml

data_water_segmentation.yaml 文件内容示例：

train_s1: ./global_water_segmentation/sentinel1_images/train/
val_s1: ./global_water_segmentation/sentinel1_images/val/

train_s2: ./global_water_segmentation/sentinel2_images/train/
val_s2: ./global_water_segmentation/sentinel2_images/val/

mask_train: ./global_water_segmentation/masks/train/
mask_val: ./global_water_segmentation/masks/val/

nc: 2  # 类别数量：Water 和 Non-Water
names: ['Non-Water', 'Water']

确保每个掩码图像（mask）是单通道的PNG图像，其中像素值代表类别ID（例如0代表非水域，1代表水域）。

2. 安装依赖库

确保安装了必要的库：

pip install segmentation-models-pytorch rasterio albumentations opencv-python-headless torch torchvision

3. 自定义数据集类

编写一个自定义的数据集类来读取来自哨兵一号和哨兵二号的多光谱图像以及对应的掩码。

import os
import cv2
import numpy as np
import rasterio
from torch.utils.data import Dataset

class WaterSegmentationDataset(Dataset):
    def __init__(self, s1_dir, s2_dir, mask_dir, transform=None):
        self.s1_fps = sorted([os.path.join(s1_dir, fname) for fname in os.listdir(s1_dir)])
        self.s2_fps = sorted([os.path.join(s2_dir, fname) for fname in os.listdir(s2_dir)])
        self.mask_fps = sorted([os.path.join(mask_dir, fname).replace('.tif', '.png') for fname in os.listdir(s1_dir)])
        self.transform = transform

    def __getitem__(self, i):
        with rasterio.open(self.s1_fps[i]) as src1:
            s1_img = src1.read().transpose((1, 2, 0))  # Sentinel-1 images
        with rasterio.open(self.s2_fps[i]) as src2:
            s2_img = src2.read().transpose((1, 2, 0))  # Sentinel-2 images
        image = np.concatenate((s1_img, s2_img), axis=2)  # Concatenate both images
        mask = cv2.imread(self.mask_fps[i], 0)

        if self.transform:
            augmented = self.transform(image=image, mask=mask)
            image = augmented['image']
            mask = augmented['mask']

        return image, mask

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

4. 模型定义与训练

选择一个合适的模型架构（例如UNet或DeepLabV3），并配置优化器、损失函数和评估指标。

训练脚本

import torch
import segmentation_models_pytorch as smp
from torch.utils.data import DataLoader

# 加载数据集
train_dataset = WaterSegmentationDataset('./global_water_segmentation/sentinel1_images/train/', './global_water_segmentation/sentinel2_images/train/', './global_water_segmentation/masks/train/')
valid_dataset = WaterSegmentationDataset('./global_water_segmentation/sentinel1_images/val/', './global_water_segmentation/sentinel2_images/val/', './global_water_segmentation/masks/val/')

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)
valid_loader = DataLoader(valid_dataset, batch_size=4, shuffle=False, num_workers=4)

# 定义模型
ENCODER = 'resnet34'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['Non-Water', 'Water']
ACTIVATION = 'sigmoid'  # 二分类问题可以使用sigmoid激活函数

model = smp.Unet(
    encoder_name=ENCODER, 
    encoder_weights=ENCODER_WEIGHTS, 
    classes=len(CLASSES), 
    activation=ACTIVATION,
)

preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)

# 设置损失函数和优化器
loss = smp.utils.losses.CrossEntropyLoss()
metrics = [
    smp.utils.metrics.IoU(threshold=0.5),
]

optimizer = torch.optim.Adam([ 
    dict(params=model.parameters(), lr=0.0001),
])

train_epoch = smp.utils.train.TrainEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    optimizer=optimizer,
    device='cuda',
    verbose=True,
)

valid_epoch = smp.utils.train.ValidEpoch(
    model, 
    loss=loss, 
    metrics=metrics, 
    device='cuda',
    verbose=True,
)

max_score = 0

for i in range(0, 40):  # 训练周期数
    print('\nEpoch: {}'.format(i))
    train_logs = train_epoch.run(train_loader)
    valid_logs = valid_epoch.run(valid_loader)
    
    if max_score < valid_logs['iou_score']:
        max_score = valid_logs['iou_score']
        torch.save(model, './best_model.pth')
        print('Model saved!')

    if i == 25:
        optimizer.param_groups[0]['lr'] /= 10
        print('Decrease decoder learning rate to 1e-5!')

5. 推理与结果可视化

训练完成后，我们可以利用训练好的模型对新图片进行预测，并将结果可视化。

推理脚本

import matplotlib.pyplot as plt

best_model = torch.load('./best_model.pth')

def visualize(image, gt_mask, pr_mask):
    figure, ax = plt.subplots(1, 3, figsize=(10, 10))

    ax[0].imshow(image[..., :3])  # 显示RGB部分
    ax[0].set_title("Image")
    ax[1].imshow(gt_mask, cmap='gray')
    ax[1].set_title("Ground Truth Mask")
    ax[2].imshow(pr_mask, cmap='gray')
    ax[2].set_title("Predicted Mask")

    plt.show()

test_dataset = WaterSegmentationDataset('./global_water_segmentation/sentinel1_images/test/', './global_water_segmentation/sentinel2_images/test/', './global_water_segmentation/masks/test/')
for i in range(5):  # 可视化前5个测试样本
    image, gt_mask = test_dataset[i]
    x_tensor = torch.from_numpy(image).to('cuda').permute(2, 0, 1).unsqueeze(0).float() / 255.
    pr_mask = best_model.predict(x_tensor)
    pr_mask = pr_mask.squeeze().cpu().numpy()
    visualize(image, gt_mask, pr_mask)