利用U-Net或TransUNet架构创建基于PyTorch框架构建针对不同城市建筑物精准提取遥感图像语义分割系统

以下文字及代码仅供参考。

遥感图像语义分割，基于Pytorch框架训练遥感建筑物数据集，针对不同城市建筑物精准提取。
模型可U-Net、TransU-Net等

1

1

构建基于PyTorch框架的遥感图像语义分割模型，专门用于不同城市的建筑物精准提取，可采用U-Net、TransUNet等架构。

如何从数据准备到模型训练及推理的完整流程？

1. 环境设置

首先确保安装了必要的库：

pip install torch torchvision albumentations opencv-python matplotlib timm

对于TransUNet，还需要额外安装pytorch-pretrained-ViT或类似库来支持Transformer模块：

pip install pytorch-pretrained-vit

2. 数据准备

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

• images/ 文件夹包含所有图像。
• masks/ 文件夹包含对应的掩码图像（标签）。

创建一个自定义的数据加载器：

import torch
from torch.utils.data import Dataset, DataLoader
import cv2
import os
import numpy as np
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

class BuildingSegmentationDataset(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(".tif", "_mask.tif"))  # 根据实际情况调整文件名格式
        image = cv2.imread(img_path)
        mask = cv2.imread(mask_path, 0)  # Load grayscale mask

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

        return image, mask.long()

transform = A.Compose(
    [
        A.Resize(512, 512),
        A.Normalize(),
        ToTensorV2(),
    ],
)

dataset = BuildingSegmentationDataset(image_dir="path/to/images", mask_dir="path/to/masks", transform=transform)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

3. 模型定义

U-Net

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

class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        # Define your U-Net architecture here
        # Example: Encoder and Decoder layers
        pass

    def forward(self, x):
        # Implement the forward pass
        pass

unet_model = UNet().cuda()

TransUNet

TransUNet结合了卷积神经网络和Vision Transformer的优点。这里提供一个简化的实现示例：

from vit_pytorch import ViT

class TransUNet(nn.Module):
    def __init__(self):
        super(TransUNet, self).__init__()
        self.vit = ViT(image_size=512, patch_size=32, num_classes=1, dim=1024, depth=6, heads=16, mlp_dim=2048)
        # Add additional layers for segmentation task if necessary

    def forward(self, x):
        x = self.vit(x)
        return x

transunet_model = TransUNet().cuda()

4. 模型训练

定义损失函数和优化器，并开始训练：

import torch.optim as optim

model = unet_model  # 或者 transunet_model

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, masks in dataloader:
        images = images.cuda()
        masks = masks.cuda()

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    
    avg_loss = running_loss / len(dataloader)
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss}")

5. 推理与结果可视化

完成训练后，使用训练好的模型进行预测并可视化结果：

import matplotlib.pyplot as plt

def visualize_predictions(model, image_tensor, original_image):
    model.eval()
    with torch.no_grad():
        prediction = model(image_tensor.unsqueeze(0).cuda())
        prediction = torch.argmax(prediction.squeeze(), dim=0).detach().cpu().numpy()

    plt.figure(figsize=(10,7))
    plt.subplot(1,2,1)
    plt.title("Original Image")
    plt.imshow(original_image)
    plt.subplot(1,2,2)
    plt.title("Predicted Mask")
    plt.imshow(prediction, cmap='gray')
    plt.show()

# 示例：加载一张图像进行预测
img_path = "path/to/test/image.tif"
original_image = cv2.imread(img_path)
image = cv2.resize(original_image, (512, 512))  # 需要与训练时尺寸一致
image_transform = transform(image=image)
image_tensor = image_transform['image'].cuda()

visualize_predictions(model, image_tensor, original_image)

总结

基于PyTorch框架构建针对不同城市建筑物精准提取的遥感图像语义分割系统。该系统利用U-Net或TransUNet架构，从数据准备、模型训练到推理和结果可视化的全流程。