基于PyTorch的建筑物识别：从语义分割到掩膜矢量化

简介：基于Python和PyTorch的建筑物识别器源码，面向计算机视觉入门者及城市规划、灾害管理等行业应用，采用MaskRCNN模型对卫星或航拍图像进行建筑物自动识别与定位。资源共16个文件，以13个Python脚本为主体，涵盖数据预处理、模型构建、训练、测试与评估等完整流程，另含字体文件、示例输出图像及中文说明文档，压缩包仅416KB，轻量便于快速部署。已有65人学习，适合作为深度学习目标检测与实例分割的实战参考。下载后可获得结构清晰的工程代码，包含独立的工具函数、模型文件、训练与测试入口，配合配置文件可自定义训练参数与数据集，直接运行即可得到带边界框和类别标签的预测结果及可视化输出。该源码既适合入门者学习MaskRCNN的工程实现，也能为遥感图像分析、城市建筑普查等实际需求提供基础框架，是值得收藏的参考示例。

1. 用 Python 和 PyTorch 识别建筑物，先分清检测与分割

把一张无人机或卫星影像丢给程序，让它把画面里的房子逐栋标出来，这个需求在 GIS 里叫建筑物识别，落到深度学习里通常不是目标检测而是语义分割。城市里房子贴着房子，检测框互相重叠，框给不了精确边界，而城市规划、地图更新要的是屋顶轮廓线。这类源码包在资源站上很常见，但下载后能直接跑通的很少，问题多出在数据加载、损失函数和推理切块策略。

下面按一个基于 Python 和 PyTorch 的建筑物识别器源码工程的标准结构，把四个环节讲透：为什么选分割模型而不是检测模型，最小可复现的 Dataset 和训练循环怎么写，训练参数和大影像推理有哪些坑，最后怎么把掩膜变成可交付的矢量轮廓。适合有 Python 基础、想跑通完整图像分割项目的人，也适合从别处下载过源码但训练一直不收敛的人对照排查。

2. 建筑物识别模型选型：PyTorch 里分割网络、数据与环境的搭配

2.1 为什么建筑物识别要做逐像素分割而不是目标检测

建筑物识别在遥感领域早期靠形态学、阴影分析和边缘检测，鲁棒性很差。深度学习成为主流后，工程实现基本分成两个方向：目标检测给出包围框，语义分割给出逐像素掩膜。航拍影像里建筑物排列密集、朝向任意，一排联排住宅的框会大面积重叠，NMS 之后两栋变一栋，这在检测方案的输出里几乎无解。栅格掩膜天然保留每个屋顶的完整形状，后续转成矢量轮廓只要提取边界即可，对 GIS 下游来说是唯一能直接用的形式。

另一个理由是输出形式。分割网络最后输出 H×W 的概率图，设一个阈值就是掩膜，不需要像检测那样设计锚框、做 NMS 调参。对源码工程来说，训练数据也更友好：在 GIS 里把屋顶多边形栅格化就能得到掩膜。如果需求里还要区分每栋楼并逐个计数，那得换成实例分割，比如 Mask R-CNN 那一路，但这类模型权重体积大、推理慢，对以二分类建筑物掩膜为目标的源码包来说并不划算。

2.2 U-Net、DeepLabV3+ 和 FCN 三个常用骨架怎么挑

FCN 是分割任务的起点，把分类网络的全连接层换成卷积，再用转置卷积把特征图恢复到原图尺寸。缺点是没有利用浅层细节，建筑物边缘容易糊成一团，现在一般只当基线用。U-Net 在编码器和解码器之间加了跳跃连接，浅层边缘信息可以直接传到解码器，这让它在小目标、单一类别、训练数据只有几千张的场景下表现突出。对建筑物识别来说，默认从 U-Net 起步是对的，PyTorch 里用 segmentation_models_pytorch 可以一行加载带预训练权重的版本：

import segmentation_models_pytorch as smp model = smp.Unet( encoder_name="resnet34", encoder_weights="imagenet", in_channels=3, classes=1 )

三者的取舍可以用一张表概括：

DeepLabV3+ 用空洞卷积加 ASPP 模块扩大感受野，对整幅大影像、目标跨度大的场景表现更好，但显存占用高一个量级，训练速度也明显更慢。

提示：如果机器显存只有 6G 左右，先把 encoder 换成 mobilenet，或者直接用只有四五层卷积的轻量 U-Net，识别精度下降有限，训练速度提升明显。

2.3 公开数据集、掩膜格式与 PyTorch 环境检查

公开数据方面，INRIA Aerial Image Labeling 是 0.3m 分辨率的航空影像，标注就是逐像素的建筑、非建筑二分类；Massachusetts Buildings 是 1m 分辨率、覆盖 137 平方公里；WHU 数据集有航空和卫星两个子集。这些数据集的掩膜都是单通道 PNG，像素值为 0 或 255，下载后先确认编码和通道数：

from PIL import Image import numpy as np mask = np.array(Image.open("mask.png")) print(mask.shape, np.unique(mask)) # 期望输出 (512, 512) [ 0 255]

这个检查很必要。很多网上下载的源码包里，掩膜是 RGB 三通道或者带调色板的伪彩色 PNG，直接转 tensor 会和标签对不上，训练时 loss 一路飘红。

环境方面，PyTorch 的安装版本依赖 CUDA 版本，官网有对应表。我一般先在 Python 里验证当前组合能不能调用 GPU：

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

返回 True 再继续。PyTorch 基础框架层只需要 torch 和 torchvision 两个包，torchvision 里的 v2 版本还自带了针对分割任务的增强工具。网上流传较多的 python 3.10 与 pytorch 2.x、CUDA 12.x 组合能用，但换机器前先用 nvidia-smi 确认驱动支持的 CUDA 版本不低于 torch 要求的版本。用 anaconda 配置 pytorch 环境是最省事的路径，conda 会把 cudatoolkit 和 cudnn 的匹配关系一并处理掉；如果用 VSCode 写这套源码，记得 interpreter 选 conda 环境里的 python，否则跑起来用的不是 torch 所在的环境。

3. 基于 PyTorch 的建筑物识别器源码：数据读取、损失函数与训练循环

3.1 自定义 Dataset：把影像和掩膜成对读进来

分割任务的数据加载核心是保证影像和掩膜逐像素对齐，任何一步 resize 或归一化不对都会直接毁掉训练。最小可复现的 Dataset 写法如下：

import os import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class BuildingDataset(Dataset): def __init__(self, image_dir, mask_dir, tile_size=256): self.image_dir = image_dir self.mask_dir = mask_dir self.image_names = sorted(os.listdir(image_dir)) self.tile_size = tile_size def __len__(self): return len(self.image_names) def __getitem__(self, idx): img = Image.open(os.path.join(self.image_dir, self.image_names[idx])).convert("RGB") mask_name = self.image_names[idx].replace(".jpg", ".png") mask = Image.open(os.path.join(self.mask_dir, mask_name)) img = img.resize((self.tile_size, self.tile_size), Image.BILINEAR) mask = mask.resize((self.tile_size, self.tile_size), Image.NEAREST) # 关键 img = np.array(img).transpose(2, 0, 1) / 255.0 mask = (np.array(mask) > 127).astype(np.float32) return torch.tensor(img, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32)

mask 的 resize 必须用 NEAREST 最近邻插值。用 BILINEAR 会给掩膜边缘插入中间灰度值，比如 0 和 255 之间插出 128，再拿去算 loss 会出现模糊梯度，训练出来的掩膜边缘全是渐变带。影像归一化放到 0 到 1，掩膜二值化用大于 127 判断而不是严格等于 255，这样能同时兼容 0/1 和 0/255 两种标签编码。

数据增强一般用随机裁剪、翻转和 90 度旋转。注意 torchvision 的 RandomResizedCrop 对影像和掩膜要传同一组随机参数，否则两张图对不上，这是源码工程里常见的隐性 bug，跑出来的 loss 不降先查这里。

3.2 BCE 加 Dice 的组合损失，解决前景像素偏少的问题

建筑物和背景的像素比例往往达到 1:10 甚至 1:50。单独用 BCEWithLogitsLoss 训练时，模型会把概率压得很低来减少背景误报，结果是掩膜里零星几点亮点，建筑物连不成片。分割任务里默认做法是把 BCE 和 Dice Loss 相加：

def dice_loss(pred_logits, target, smooth=1.0): pred = torch.sigmoid(pred_logits) intersection = (pred * target).sum() return 1.0 - (2.0 * intersection + smooth) / (pred.sum() + target.sum() + smooth)

Dice Loss 直接优化交并比，对前景占比不敏感；smooth 参数避免分子分母同时为 0，取 1.0 就行，不需要额外调。组合时 Dice 的梯度占比会比 BCE 高，这正好把训练重心拉回到建筑物轮廓上。

3.3 训练循环、优化器参数与 IoU 计算

下面是一段可以直接跑的完整训练循环：

from torch.utils.data import DataLoader model = smp.Unet(encoder_name="resnet34", encoder_weights="imagenet", in_channels=3, classes=1).cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50) bce = torch.nn.BCEWithLogitsLoss() for epoch in range(50): model.train() total_loss = 0.0 for imgs, masks in DataLoader(train_set, batch_size=8, shuffle=True, num_workers=4, drop_last=True): imgs, masks = imgs.cuda(), masks.cuda() logits = model(imgs) loss = bce(logits, masks) + dice_loss(logits, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() scheduler.step() torch.save(model.state_dict(), f"building_unet_{epoch}.pth") print(f"epoch {epoch} loss {total_loss / len(train_set):.4f}")

参数说明：AdamW 的 weight_decay 对分割网络取 1e-5 到 1e-4 之间，取大了编码器权重被压得过死，迁移学习的效果会打折。CosineAnnealingLR 的 T_max 与总 epoch 一致，让学习率在最后一个 epoch 降到接近 0；如果验证集指标波动大，改成 ReduceLROnPlateau 按验证 IoU 触发更可靠。drop_last=True 防止最后一批不足 batch_size 时 BN 统计量抖动。

评估阶段最核心的指标是 IoU：

def compute_iou(logits, mask, threshold=0.5): pred = (torch.sigmoid(logits) > threshold).bool() mask = mask.bool() inter = (pred & mask).sum().item() union = (pred | mask).sum().item() return inter / (union + 1e-6)

阈值默认 0.5，但实际交付时可以在验证集上扫 0.3 到 0.7 的区间，挑出 IoU 最高的那个阈值，这一步通常能白捡一两个点。保存权重时建议把验证 IoU 最高的 checkpoint 单独存一份，不要用最后一个 epoch 的权重，分割网络在训练后期往往有过拟合倾向。

4. 建筑物识别参数调优与整幅影像的滑窗推理

4.1 学习率、批大小与切块尺寸的联动关系

遥感影像原图动辄几千像素，直接整图进网络不现实，一般切成 256 到 1024 的小块训练。切块尺寸、批大小和初始学习率是联动的，参考配置如下：

切块越大，单张图能看到的上下文越多，U-Net 对连片小区、大型公建的整体把握更好，但批大小被迫变小，BN 统计量抖动变大，学习率也要跟着降。显存不够时先降 batch size，再考虑冻结 encoder 前几层的梯度；PyTorch 2.x 用 AMP 半精度训练可以把 batch 翻倍，torch.amp 接口包一层 GradScaler 就行，代码改动半小时内能完成。实际训练中先跑 10 个 epoch 看 loss 曲线，初始学习率合理时 loss 应该在前两三个 epoch 明显下降，纹丝不动就调大一个量级，剧烈震荡就调小。

4.2 建筑物占比太低时的损失函数补偿

建筑物在影像里的像素占比通常不到 10%，训练集里如果出现大片农田、水面，模型很容易把所有像素判成背景，因为这样已经能拿到很高的准确率。除了加 Dice Loss，还可以给 BCE 加 pos_weight，把背景类梯度降权：

pos_weight = torch.tensor([background_count / building_count]) bce = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)

pos_weight 这样设置等价于让模型更关注稀有类别。Focal Loss 是另一种做法，alpha 取 0.75、gamma 取 2，对难分样本的梯度放大更明显，但多出两个要调的参数。对应的验证指标也要改：类别不均衡时准确率没有参考意义，全预测背景都可能到 90% 以上，所以只看 IoU 和 F1，这两个指标都对前景类别敏感。

4.3 整幅遥感影像的滑窗推理与重叠融合

推理阶段要把训练好的模型搬回整幅大影像。直接切块推断再接回去，块与块之间会有明显的条带痕迹，原因是边缘像素看到的上下文比中心少，置信度偏低。常见做法是滑窗加重叠平均：

import numpy as np import torch def sliding_window_infer(model, big_img, tile_size=512, stride=256, device="cuda"): h, w = big_img.shape[:2] prob = np.zeros((h, w), dtype=np.float32) weight = np.zeros((h, w), dtype=np.float32) model.eval() with torch.no_grad(): for y in range(0, h - tile_size + 1, stride): for x in range(0, w - tile_size + 1, stride): tile = big_img[y:y + tile_size, x:x + tile_size] t = torch.tensor(tile.transpose(2, 0, 1) / 255.0, dtype=torch.float32).unsqueeze(0).to(device) p = torch.sigmoid(model(t)).squeeze().cpu().numpy() prob[y:y + tile_size, x:x + tile_size] += p weight[y:y + tile_size, x:x + tile_size] += 1 return prob / np.clip(weight, 1.0, None)

stride 小于 tile_size 时相邻图块有重叠，每个像素被多次预测后取平均，能明显压掉边缘衔接处的条带。stride 取 tile_size 的一半是速度和质量的折中写法。图像尺寸不能被 tile_size 整除时，先按对称 padding 把边缘补齐，推理完再裁回来，避免最后一行一列出现黑边。推理时务必带着 model.eval() 和 torch.no_grad()，这两个少一个，显存和延迟都会明显上涨。如果输入是 geotiff，推理结果直接写成单波段 tiff，中间别存 png，避免 8bit 量化丢掉概率信息。

5. 结果交付前的后处理：连通域过滤与掩膜矢量化

5.1 用连通域面积过滤孤立误检

模型输出的概率图阈值化之后，经常带一批零星的小噪声，比如树冠阴影、汽车顶棚被误判成建筑。形态学开运算能去掉一部分，但更可控的做法是按连通域面积过滤：

from scipy import ndimage pred_mask = (prob > threshold).astype(np.uint8) labeled, num = ndimage.label(pred_mask) sizes = ndimage.sum(pred_mask, labeled, range(1, num + 1)) for label_id, size in enumerate(sizes, start=1): if size < min_area: pred_mask[labeled == label_id] = 0

min_area 的取值要结合影像分辨率换算。0.5 米分辨率下，一个像素对应 0.25 平方米，10 平方米的杂物间约 40 像素，正常住宅楼至少几百像素，取 200 能把大多数噪声滤掉。1 米分辨率时同一栋楼像素数变成四分之一，min_area 相应除以 4。

5.2 从掩膜提取建筑物轮廓并输出 GeoJSON

掩膜变成矢量轮廓用 OpenCV 的 findContours 一行完成：

import cv2 mask_u8 = (pred_mask * 255).astype(np.uint8) contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) polygons = [cv2.approxPolyDP(c, epsilon=2.0, closed=True) for c in contours]

RETR_EXTERNAL 只取最外层轮廓，屋顶中间的庭院、天窗不会被拆成多个多边形；approxPolyDP 的 epsilon 控制简化程度，2.0 像素可以把锯齿边压平，又不会丢失直角结构。要交付 GeoJSON 时，把每个轮廓的坐标按 [x, y] 顺序写入 Polygon 的坐标数组，再套上地理参考信息就能直接进 QGIS。到这一步，一个基于 Python 和 PyTorch 的建筑物识别器就从原始影像走到了带轮廓的交付物，后续无论是统计屋顶面积还是做变化检测，拿到的都是干净可用的边界数据。

本文还有配套的精品资源，点击获取