Unet遥感语义分割

### 基于 U-Net 网络的遥感图像语义分割实现 #### 背景介绍 U-Net 是一种经典的卷积神经网络架构，最初设计用于生物医学图像分割任务。它通过编码器-解码器结构以及跳跃连接实现了高分辨率特征图的重建[^2]。在遥感图像领域，U-Net 同样被广泛应用于地物分类、目标检测和语义分割等任务。 以下是基于 PyTorch 和 TensorFlow 实现 U-Net 进行遥感图像语义分割的具体方法： --- #### 使用 PyTorch 实现 U-Net 遥感图像语义分割 以下是一个简单的 PyTorch 版本的 U-Net 架构实现代码示例： ```python import torch import torch.nn as nn class DoubleConv(nn.Module): """(convolution => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = 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) ) def forward(self, x): return self.double_conv(x) class UNet(nn.Module): def __init__(self, n_channels, n_classes): super(UNet, self).__init__() self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) factor = 2 if torch.cuda.is_available() else 1 self.down4 = Down(512, 1024 // factor) self.up1 = Up(1024, 512 // factor) self.up2 = Up(512, 256 // factor) self.up3 = Up(256, 128 // factor) self.up4 = Up(128, 64) 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 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 = nn.functional.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) ``` 上述代码定义了一个标准的 U-Net 结构，适用于遥感图像语义分割任务。可以通过调整输入通道数 `n_channels` 和类别数 `n_classes` 来适配不同的数据集需求[^1]。 --- #### 使用 TensorFlow 实现 U-Net 遥感图像语义分割 对于 TensorFlow 用户，可以利用 Keras API 快速搭建 U-Net 模型： ```python from tensorflow.keras.models import Model from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Dropout, concatenate, UpSampling2D def unet(input_shape=(256, 256, 3), num_classes=1): inputs = Input(shape=input_shape) c1 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(inputs) c1 = Dropout(0.1)(c1) c1 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c1) p1 = MaxPooling2D((2, 2))(c1) c2 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p1) c2 = Dropout(0.1)(c2) c2 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c2) p2 = MaxPooling2D((2, 2))(c2) c3 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p2) c3 = Dropout(0.2)(c3) c3 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c3) p3 = MaxPooling2D((2, 2))(c3) c4 = Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p3) c4 = Dropout(0.2)(c4) c4 = Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c4) p4 = MaxPooling2D(pool_size=(2, 2))(c4) c5 = Conv2D(1024, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p4) c5 = Dropout(0.3)(c5) c5 = Conv2D(1024, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c5) u6 = Conv2DTranspose(512, (2, 2), strides=(2, 2), padding='same')(c5) u6 = concatenate([u6, c4]) c6 = Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u6) c6 = Dropout(0.2)(c6) c6 = Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c6) u7 = Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(c6) u7 = concatenate([u7, c3]) c7 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u7) c7 = Dropout(0.2)(c7) c7 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c7) u8 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c7) u8 = concatenate([u8, c2]) c8 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u8) c8 = Dropout(0.1)(c8) c8 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c8) u9 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c8) u9 = concatenate([u9, c1], axis=3) c9 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u9) c9 = Dropout(0.1)(c9) c9 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c9) outputs = Conv2D(num_classes, (1, 1), activation='sigmoid')(c9) model = Model(inputs=[inputs], outputs=[outputs]) return model ``` 此代码片段展示了如何使用 TensorFlow/Keras 定义一个完整的 U-Net 模型，并支持多类别的语义分割任务[^3]。 --- #### 数据预处理与训练流程 无论是 PyTorch 还是 TensorFlow，在实际应用中都需要完成以下几个步骤： 1. **数据准备**：加载并预处理遥感图像及其对应的标注掩膜。 2. **模型编译**：配置损失函数（如 Dice Loss 或 Cross Entropy）、优化器（Adam 或 SGD）和评估指标。 3. **模型训练**：划分训练集和验证集，设置合适的批量大小和迭代次数进行训练。 4. **性能评估**：测试模型在独立测试集上的表现，计算 IoU、Precision、Recall 等评价指标。 ---