当前位置：首页 > news >正文

建筑物检测数据集深度学习中的语义分割方法来识别图像中的建筑物区域

news 2026/8/9 13:07:43

以U-Net作为模型采用深度学习中的语义分割方法来识别图像中的建筑物区域_识别检测并训练spacenet2建筑物检测数据集

spacenet2建筑物检测数据集

SpaceNet 2 数据集主要用于建筑物检测任务，该数据集包含高分辨率卫星图像及其对应的建筑物掩码。为了使用这套数据集进行训练和评估，我们可以采用深度学习中的语义分割方法来识别图像中的建筑物区域。这里我们将以U-Net作为示例模型，并提供从数据准备到模型训练的完整流程。
1

SpaceNet 2 数据集主要用于建筑物检测任务，该数据集包含高分辨率卫星图像及其对应的建筑物掩码。使用这套数据集进行训练和评估，采用深度学习中的语义分割方法来识别图像中的建筑物区域。以U-Net作为模型，从数据准备到模型训练的完整流程。

1. 环境配置

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

pipinstalltorch torchvision albumentations opencv-python scikit-image

2. 数据准备

SpaceNet 2 数据集通常由卫星图像和相应的建筑物掩码组成。我们需要将这些数据转换为适合训练的形式。假设您的数据结构如下：

/path/to/spacenet2/ ├── train_images/ │ ├── img1.tif │ ├── img2.tif │ └── ... ├── train_masks/ │ ├── mask1.png │ ├── mask2.png │ └── ...

编写一个自定义的数据集类来加载这些图像和掩码。

importcv2importtorchfromtorch.utils.dataimportDatasetimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2importosfromskimageimportioclassSpaceNetDataset(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):returnlen(self.images)def__getitem__(self,index):img_path=os.path.join(self.image_dir,self.images[index])mask_path=os.path.join(self.mask_dir,self.images[index].replace('.tif','.png'))image=io.imread(img_path)mask=cv2.imread(mask_path,cv2.IMREAD_GRAYSCALE)mask[mask>0]=1# 将所有非零值设置为1，表示建筑物区域ifself.transformisnotNone:augmentations=self.transform(image=image,mask=mask)image=augmentations["image"]mask=augmentations["mask"]returnimage,mask

定义一些基本的数据增强操作：

transform=A.Compose([A.Resize(height=512,width=512),A.Normalize(mean=(0.485,),std=(0.229,)),# 根据实际情况调整ToTensorV2(),])

3. 模型定义

使用U-Net模型进行语义分割：

importtorch.nnasnnimporttorchclassUNet(nn.Module):def__init__(self,n_channels,n_classes):super(UNet,self).__init__()defconv_block(in_channels,out_channels):returnnn.Sequential(nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1),nn.ReLU(inplace=True),nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1),nn.ReLU(inplace=True),)self.encoder1=conv_block(n_channels,64)self.pool1=nn.MaxPool2d(kernel_size=2,stride=2)self.encoder2=conv_block(64,128)self.pool2=nn.MaxPool2d(kernel_size=2,stride=2)self.encoder3=conv_block(128,256)self.pool3=nn.MaxPool2d(kernel_size=2,stride=2)self.encoder4=conv_block(256,512)self.pool4=nn.MaxPool2d(kernel_size=2,stride=2)self.bottleneck=conv_block(512,1024)self.upconv4=nn.ConvTranspose2d(1024,512,kernel_size=2,stride=2)self.decoder4=conv_block(1024,512)self.upconv3=nn.ConvTranspose2d(512,256,kernel_size=2,stride=2)self.decoder3=conv_block(512,256)self.upconv2=nn.ConvTranspose2d(256,128,kernel_size=2,stride=2)self.decoder2=conv_block(256,128)self.upconv1=nn.ConvTranspose2d(128,64,kernel_size=2,stride=2)self.decoder1=conv_block(128,64)self.out_conv=nn.Conv2d(64,n_classes,kernel_size=1)defforward(self,x):enc1=self.encoder1(x)enc2=self.encoder2(self.pool1(enc1))enc3=self.encoder3(self.pool2(enc2))enc4=self.encoder4(self.pool3(enc3))bottleneck=self.bottleneck(self.pool4(enc4))dec4=self.upconv4(bottleneck)dec4=torch.cat((dec4,enc4),dim=1)dec4=self.decoder4(dec4)dec3=self.upconv3(dec4)dec3=torch.cat((dec3,enc3),dim=1)dec3=self.decoder3(dec3)dec2=self.upconv2(dec3)dec2=torch.cat((dec2,enc2),dim=1)dec2=self.decoder2(dec2)dec1=self.upconv1(dec2)dec1=torch.cat((dec1,enc1),dim=1)dec1=self.decoder1(dec1)returntorch.sigmoid(self.out_conv(dec1))model=UNet(n_channels=3,n_classes=1).cuda()# 假设输入图像是RGB三通道

4. 训练过程

编写训练循环：

fromtorch.utils.dataimportDataLoaderimporttorch.optimasoptim dataset=SpaceNetDataset('/path/to/train_images','/path/to/train_masks',transform=transform)dataloader=DataLoader(dataset,batch_size=4,shuffle=True)criterion=nn.BCELoss()optimizer=optim.Adam(model.parameters(),lr=0.001)deftrain_model(model,criterion,optimizer,dataloader,num_epochs=25):forepochinrange(num_epochs):print(f'Epoch{epoch}/{num_epochs-1}')print('-'*10)model.train()running_loss=0.0forinputs,masksindataloader:inputs=inputs.float().cuda()masks=masks.unsqueeze(1).float().cuda()optimizer.zero_grad()withtorch.set_grad_enabled(True):outputs=model(inputs)loss=criterion(outputs,masks)loss.backward()optimizer.step()running_loss+=loss.item()*inputs.size(0)epoch_loss=running_loss/len(dataloader.dataset)print(f'Train Loss:{epoch_loss:.4f}')returnmodel trained_model=train_model(model,criterion,optimizer,dataloader,num_epochs=25)

5. 模型保存与加载

训练完成后，保存模型权重以便后续使用：

torch.save(model.state_dict(),'/path/to/save/your/model.pth')

加载已保存的模型：

model.load_state_dict(torch.load('/path/to/save/your/model.pth'))

6. 模型评估

在验证集或测试集上评估模型性能：

model.eval()withtorch.no_grad():forinputs,masksinval_dataloader:inputs=inputs.float().cuda()masks=masks.unsqueeze(1).float().cuda()outputs=model(inputs)# 这里可以根据需要计算IoU等指标