遥感图像语义分割CPU pytorch

### 使用 PyTorch 实现遥感图像语义分割 为了在 CPU 上使用 PyTorch 进行遥感图像的语义分割，可以按照以下方法构建完整的流程。这包括数据准备、模型定义、训练过程以及评估阶段。 #### 数据准备 首先需要准备好用于语义分割的数据集。通常情况下，遥感图像语义分割任务会涉及大量的标注数据。这些数据可以通过公开数据集获取或者自行制作[^2]。对于标签文件，它们通常是像素级别的分类图，其中不同的类别由不同数值表示[^3]。 以下是加载和预处理数据的一个简单例子： ```python import torch from torchvision import transforms from PIL import Image import os import numpy as np class RemoteSensingDataset(torch.utils.data.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 = Image.open(img_path).convert("RGB") mask = Image.open(mask_path) if self.transform is not None: augmentations = self.transform(image=np.array(image), mask=np.array(mask)) image = augmentations["image"] mask = augmentations["mask"] return image, mask ``` 上述代码展示了如何创建自定义数据集类 `RemoteSensingDataset` 来读取遥感图像及其对应的掩码。 #### 定义 U-Net 模型 U-Net 是一种广泛应用于医学影像分析和遥感图像分割的经典架构。下面是一个简单的 U-Net 结构实现： ```python import torch.nn as nn import torch.nn.functional as F class UNet(nn.Module): def __init__(self, num_classes=10): super(UNet, self).__init__() # Encoder layers (downsampling path) self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1) self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # Decoder layers (upsampling path) self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.conv3_1 = nn.Conv2d(192, 64, kernel_size=3, padding=1) self.conv3_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.final_conv = nn.Conv2d(64, num_classes, kernel_size=1) def forward(self, x): enc1 = F.relu(self.conv1_2(F.relu(self.conv1_1(x)))) pool1 = self.pool1(enc1) enc2 = F.relu(self.conv2_2(F.relu(self.conv2_1(pool1)))) pool2 = self.pool2(enc2) up1 = self.upconv1(pool2) concat1 = torch.cat([enc1, up1], dim=1) dec1 = F.relu(self.conv3_2(F.relu(self.conv3_1(concat1)))) output = self.final_conv(dec1) return output ``` 此部分实现了基础版本的 U-Net 架构，适用于多类别的语义分割任务。 #### 训练配置与优化器设置 由于目标是在 CPU 上运行，因此无需显式指定设备为 GPU。默认情况下，PyTorch 的张量操作会在 CPU 上执行。 ```python device = torch.device('cpu') # 默认使用 CPU model = UNet(num_classes=10).to(device) # 初始化模型并移动到 CPU criterion = nn.CrossEntropyLoss() # 处理多分类问题的标准损失函数 optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam 优化器 ``` #### 开始训练循环 下面是训练的主要逻辑框架： ```python def train_model(dataloader, model, criterion, optimizer, epochs=10): for epoch in range(epochs): running_loss = 0.0 for images, masks in dataloader: images = images.to(device) masks = masks.long().to(device) outputs = model(images) loss = criterion(outputs, masks.squeeze(1)) optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss:.4f}') ``` 通过调用该函数即可启动整个训练过程[^1]。 #### 可视化预测结果 完成训练之后，可以用测试样本验证模型性能，并绘制预测结果与真实标签对比图。 ```python import matplotlib.pyplot as plt def visualize_predictions(model, test_loader): with torch.no_grad(): for i, (images, masks) in enumerate(test_loader): images = images.to(device) predictions = model(images).argmax(dim=1).cpu().numpy()[0] fig, axarr = plt.subplots(1, 2) axarr[0].imshow(predictions, cmap='gray') axarr[0].set_title('Predicted Mask') true_mask = masks.numpy()[0] axarr[1].imshow(true_mask, cmap='gray') axarr[1].set_title('True Mask') plt.show() break ``` 以上脚本能够帮助展示单次推理的结果以便直观了解模型表现情况。 --- ###