0.08米超高清无人机图像建筑物分割，共18396张，尺寸为4096×4096，或处理为1177344张尺寸为512×512图像
超高清无人机图像建筑物分割数据集，提供使用深度学习模型进行训练的代码。

数据集介绍

数据集概述

• 数据集名称：Ultra-High Resolution Drone Image Building Segmentation Dataset (UHDRDIBSD)
• 数据类型：超高清无人机航拍图像
• 目标任务：建筑物分割
• 样本数量：18396张图像，每张图像尺寸为4096×4096
• 处理后图像数量：1177344张图像，每张图像尺寸为512×512
• 标注格式：像素级分割标签

数据集目录结构

UHDRDIBSD/
├── images/
│   ├── train/
│   └── val/
├── masks/
│   ├── train/
│   └── val/
└── data.yaml

数据集配置文件

创建一个data.yaml文件，配置数据集的路径和类别信息：

path: ./UHDRDIBSD  # 数据集路径
train: images/train  # 训练集图像路径
val: images/val  # 验证集图像路径

nc: 1  # 类别数（二值分割）
names: ['building']  # 类别名称

图像预处理

由于原始图像尺寸较大（4096×4096），我们需要将其切分为更小的图像块（512×512）。

切分脚本

import os
import cv2
import numpy as np

def split_image(image, mask, patch_size=512, stride=512):
    h, w = image.shape[:2]
    patches = []
    mask_patches = []

    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            patch = image[i:i+patch_size, j:j+patch_size]
            mask_patch = mask[i:i+patch_size, j:j+patch_size]
            patches.append(patch)
            mask_patches.append(mask_patch)

    return patches, mask_patches

def save_patches(patches, mask_patches, output_dir, prefix):
    os.makedirs(output_dir, exist_ok=True)
    for idx, (patch, mask_patch) in enumerate(zip(patches, mask_patches)):
        cv2.imwrite(os.path.join(output_dir, f"{prefix}_{idx}.png"), patch)
        cv2.imwrite(os.path.join(output_dir, f"{prefix}_mask_{idx}.png"), mask_patch)

def process_images(input_dir, output_dir, patch_size=512, stride=512):
    os.makedirs(output_dir, exist_ok=True)
    for filename in os.listdir(input_dir):
        if filename.endswith('.png'):
            image_path = os.path.join(input_dir, filename)
            mask_path = os.path.join(input_dir.replace('images', 'masks'), filename)
            image = cv2.imread(image_path)
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            patches, mask_patches = split_image(image, mask, patch_size, stride)
            save_patches(patches, mask_patches, output_dir, filename.split('.')[0])

if __name__ == "__main__":
    input_train_dir = 'UHDRDIBSD/images/train'
    output_train_dir = 'UHDRDIBSD/images/train_512'
    process_images(input_train_dir, output_train_dir)

    input_val_dir = 'UHDRDIBSD/images/val'
    output_val_dir = 'UHDRDIBSD/images/val_512'
    process_images(input_val_dir, output_val_dir)

模型训练

我们将使用PyTorch和Segmentation Models PyTorch (SMP)库进行训练。

1. 安装依赖项：

   pip install torch torchvision segmentation-models-pytorch
   

2. 数据加载器：

   import os
   import cv2
   import torch
   from torch.utils.data import Dataset, DataLoader
   from albumentations import Compose, Resize, Normalize
   from albumentations.pytorch 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 = [f for f in os.listdir(image_dir) if f.endswith('.png')]
   
       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('.png', '_mask.png'))
           image = cv2.imread(img_path)
           mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
   
           if self.transform:
               augmented = self.transform(image=image, mask=mask)
               image = augmented['image']
               mask = augmented['mask']
   
           return image, mask
   
   def get_data_loaders(image_dir, mask_dir, batch_size=16, num_workers=4):
       transform = Compose([
           Resize(512, 512),
           Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
           ToTensorV2()
       ])
   
       train_dataset = BuildingSegmentationDataset(image_dir, mask_dir, transform=transform)
       train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
   
       val_dataset = BuildingSegmentationDataset(image_dir.replace('train', 'val'), mask_dir.replace('train', 'val'), transform=transform)
       val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
   
       return train_loader, val_loader
   

3. 模型训练：

   import torch
   import torch.nn as nn
   import torch.optim as optim
   from torch.utils.tensorboard import SummaryWriter
   from segmentation_models_pytorch import Unet
   
   def train_model(train_loader, val_loader, model, criterion, optimizer, num_epochs=100, device='cuda'):
       writer = SummaryWriter()
       best_val_loss = float('inf')
   
       for epoch in range(num_epochs):
           model.train()
           running_loss = 0.0
           for images, masks in train_loader:
               images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
               optimizer.zero_grad()
               outputs = model(images)
               loss = criterion(outputs, masks)
               loss.backward()
               optimizer.step()
               running_loss += loss.item()
   
           train_loss = running_loss / len(train_loader)
           writer.add_scalar('Loss/train', train_loss, epoch)
   
           model.eval()
           val_loss = 0.0
           with torch.no_grad():
               for images, masks in val_loader:
                   images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
                   outputs = model(images)
                   loss = criterion(outputs, masks)
                   val_loss += loss.item()
   
           val_loss /= len(val_loader)
           writer.add_scalar('Loss/val', val_loss, epoch)
   
           if val_loss < best_val_loss:
               best_val_loss = val_loss
               torch.save(model.state_dict(), 'best_model.pth')
   
           print(f"Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
   
       writer.close()
   
   if __name__ == "__main__":
       image_dir = 'UHDRDIBSD/images/train_512'
       mask_dir = 'UHDRDIBSD/masks/train_512'
       batch_size = 16
       num_workers = 4
       num_epochs = 100
       device = 'cuda' if torch.cuda.is_available() else 'cpu'
   
       train_loader, val_loader = get_data_loaders(image_dir, mask_dir, batch_size, num_workers)
   
       model = Unet(encoder_name='resnet34', in_channels=3, classes=1).to(device)
       criterion = nn.BCEWithLogitsLoss()
       optimizer = optim.Adam(model.parameters(), lr=1e-4)
   
       train_model(train_loader, val_loader, model, criterion, optimizer, num_epochs, device)
   

详细解释

1. 安装依赖项：

   • 安装PyTorch、TorchVision和Segmentation Models PyTorch库。

2. 数据加载器：

   • 创建一个自定义的BuildingSegmentationDataset类，用于加载图像和对应的分割标签。
   • 使用Albumentations库进行数据增强和标准化。
   • 创建数据加载器，用于批量加载数据。

3. 模型训练：

   • 使用Unet模型进行建筑物分割。
   • 定义损失函数和优化器。
   • 训练模型并在每个epoch结束时记录训练和验证损失。
   • 保存最佳模型。

运行训练脚本

将上述脚本保存为一个Python文件（例如train_building_segmentation.py），然后运行它。

python train_building_segmentation.py

评估模型

1. 评估模型：

   import torch
   from torch.utils.data import DataLoader
   from segmentation_models_pytorch import Unet
   from sklearn.metrics import jaccard_score
   import numpy as np
   
   def evaluate_model(val_loader, model, device='cuda'):
       model.eval()
       jaccard_scores = []
   
       with torch.no_grad():
           for images, masks in val_loader:
               images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
               outputs = model(images)
               preds = (torch.sigmoid(outputs) > 0.5).float()
               jaccard_scores.extend([jaccard_score(masks.cpu().numpy().flatten(), preds.cpu().numpy().flatten())])
   
       mean_jaccard_score = np.mean(jaccard_scores)
       print(f"Mean Jaccard Score: {mean_jaccard_score:.4f}")
   
   if __name__ == "__main__":
       image_dir = 'UHDRDIBSD/images/val_512'
       mask_dir = 'UHDRDIBSD/masks/val_512'
       batch_size = 16
       num_workers = 4
       device = 'cuda' if torch.cuda.is_available() else 'cpu'
   
       _, val_loader = get_data_loaders(image_dir, mask_dir, batch_size, num_workers)
   
       model = Unet(encoder_name='resnet34', in_channels=3, classes=1).to(device)
       model.load_state_dict(torch.load('best_model.pth'))
   
       evaluate_model(val_loader, model, device)
   

详细解释

1. 评估模型：
   • 加载训练好的模型权重。
   • 在验证集上进行评估，计算Jaccard分数（IoU）。
   • 打印平均Jaccard分数。

运行评估脚本

将上述脚本保存为一个Python文件（例如evaluate_building_segmentation.py），然后运行它。

python evaluate_building_segmentation.py

一键运行脚本

为了实现一键运行，可以将图像预处理、训练和评估脚本合并到一个主脚本中，并添加命令行参数来控制运行模式。

import argparse
import os
import cv2
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from segmentation_models_pytorch import Unet
from albumentations import Compose, Resize, Normalize
from albumentations.pytorch import ToTensorV2
from sklearn.metrics import jaccard_score
import numpy as np

def split_image(image, mask, patch_size=512, stride=512):
    h, w = image.shape[:2]
    patches = []
    mask_patches = []

    for i in range(0, h - patch_size + 1, stride):
        for j in range(0, w - patch_size + 1, stride):
            patch = image[i:i+patch_size, j:j+patch_size]
            mask_patch = mask[i:i+patch_size, j:j+patch_size]
            patches.append(patch)
            mask_patches.append(mask_patch)

    return patches, mask_patches

def save_patches(patches, mask_patches, output_dir, prefix):
    os.makedirs(output_dir, exist_ok=True)
    for idx, (patch, mask_patch) in enumerate(zip(patches, mask_patches)):
        cv2.imwrite(os.path.join(output_dir, f"{prefix}_{idx}.png"), patch)
        cv2.imwrite(os.path.join(output_dir, f"{prefix}_mask_{idx}.png"), mask_patch)

def process_images(input_dir, output_dir, patch_size=512, stride=512):
    os.makedirs(output_dir, exist_ok=True)
    for filename in os.listdir(input_dir):
        if filename.endswith('.png'):
            image_path = os.path.join(input_dir, filename)
            mask_path = os.path.join(input_dir.replace('images', 'masks'), filename)
            image = cv2.imread(image_path)
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            patches, mask_patches = split_image(image, mask, patch_size, stride)
            save_patches(patches, mask_patches, output_dir, filename.split('.')[0])

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 = [f for f in os.listdir(image_dir) if f.endswith('.png')]

    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('.png', '_mask.png'))
        image = cv2.imread(img_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

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

        return image, mask

def get_data_loaders(image_dir, mask_dir, batch_size=16, num_workers=4):
    transform = Compose([
        Resize(512, 512),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2()
    ])

    train_dataset = BuildingSegmentationDataset(image_dir, mask_dir, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

    val_dataset = BuildingSegmentationDataset(image_dir.replace('train', 'val'), mask_dir.replace('train', 'val'), transform=transform)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    return train_loader, val_loader

def train_model(train_loader, val_loader, model, criterion, optimizer, num_epochs=100, device='cuda'):
    writer = SummaryWriter()
    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for images, masks in train_loader:
            images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        train_loss = running_loss / len(train_loader)
        writer.add_scalar('Loss/train', train_loss, epoch)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for images, masks in val_loader:
                images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
                outputs = model(images)
                loss = criterion(outputs, masks)
                val_loss += loss.item()

        val_loss /= len(val_loader)
        writer.add_scalar('Loss/val', val_loss, epoch)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pth')

        print(f"Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

    writer.close()

def evaluate_model(val_loader, model, device='cuda'):
    model.eval()
    jaccard_scores = []

    with torch.no_grad():
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device).unsqueeze(1).float()
            outputs = model(images)
            preds = (torch.sigmoid(outputs) > 0.5).float()
            jaccard_scores.extend([jaccard_score(masks.cpu().numpy().flatten(), preds.cpu().numpy().flatten())])

    mean_jaccard_score = np.mean(jaccard_scores)
    print(f"Mean Jaccard Score: {mean_jaccard_score:.4f}")

def main(mode):
    if mode == 'preprocess':
        input_train_dir = 'UHDRDIBSD/images/train'
        output_train_dir = 'UHDRDIBSD/images/train_512'
        process_images(input_train_dir, output_train_dir)

        input_val_dir = 'UHDRDIBSD/images/val'
        output_val_dir = 'UHDRDIBSD/images/val_512'
        process_images(input_val_dir, output_val_dir)
    elif mode == 'train':
        image_dir = 'UHDRDIBSD/images/train_512'
        mask_dir = 'UHDRDIBSD/masks/train_512'
        batch_size = 16
        num_workers = 4
        num_epochs = 100
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        train_loader, val_loader = get_data_loaders(image_dir, mask_dir, batch_size, num_workers)

        model = Unet(encoder_name='resnet34', in_channels=3, classes=1).to(device)
        criterion = nn.BCEWithLogitsLoss()
        optimizer = optim.Adam(model.parameters(), lr=1e-4)

        train_model(train_loader, val_loader, model, criterion, optimizer, num_epochs, device)
    elif mode == 'eval':
        image_dir = 'UHDRDIBSD/images/val_512'
        mask_dir = 'UHDRDIBSD/masks/val_512'
        batch_size = 16
        num_workers = 4
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        _, val_loader = get_data_loaders(image_dir, mask_dir, batch_size, num_workers)

        model = Unet(encoder_name='resnet34', in_channels=3, classes=1).to(device)
        model.load_state_dict(torch.load('best_model.pth'))

        evaluate_model(val_loader, model, device)
    else:
        print("Invalid mode. Use 'preprocess', 'train', or 'eval'.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Preprocess, train, or evaluate building segmentation model.")
    parser.add_argument('mode', type=str, choices=['preprocess', 'train', 'eval'], help="Mode: 'preprocess', 'train', or 'eval'")
    args = parser.parse_args()

    main(args.mode)

详细解释

1. 命令行参数：

   • 使用argparse库添加命令行参数，控制脚本的运行模式（预处理、训练或评估）。

2. 主函数：

   • 根据传入的模式参数，调用相应的预处理、训练或评估函数。

运行主脚本

将上述脚本保存为一个Python文件（例如main_building_segmentation.py），然后运行它。

预处理图像

python main_building_segmentation.py preprocess

训练模型

python main_building_segmentation.py train

评估模型

python main_building_segmentation.py eval

总结

通过以上步骤，你可以准备好超高清无人机图像建筑物分割数据集，并使用深度学习模型进行训练和评估。希望这些信息对你有帮助！如果你有任何其他问题或需要进一步的帮助，请随时告诉我。