PyTorch+Mask R-CNN 建筑物识别：从训练到部署的完整实践

简介：这是一份基于Python语言与PyTorch框架的建筑物识别器源码项目，面向机器视觉初学者、遥感及地理信息分析人员，解决从卫星或航拍影像中自动检测、识别并框出建筑物的实际需求。资源包共16个文件，包含13个Python脚本，分别承担模型训练、测试、推理、评估与辅助工具等环节；另有Markdown说明文档、示例输出图片和字体文件各1个，整体压缩包大小仅416KB，内容紧凑、便于下载查阅。目前已有65人学习/下载，内容精炼，适合快速上手体验完整的检测流程。项目采用MaskRCNN模型，能够对建筑物进行像素级或区域级识别，并输出类别与边界框，可直接迁移到城市规划、灾害管理和环境监测等应用场景；源码结构清晰、注释与文档齐备，既可作为深度学习目标检测的教学范例，也便于开发者替换或增补数据进行二次开发与实验验证。

1. 建筑物识别任务为什么值得自己动手训练一个模型

卫星和航拍影像里的建筑物识别，一直是遥感、城市规划、灾害评估领域的刚需。公开的检测模型很多，但真正落到自己手里的遥感数据上，几乎都要重新训练或微调。这套基于 Python 和 PyTorch 的建筑物识别器，选用 Mask R-CNN 作为核心模型，把数据加载、训练、评估、推理、可视化全流程串成了一个可以直接跑的工程。对于想快速验证 maskrcnn 在自己数据集上效果的工程师来说，项目里 train.py、detect.py、eval.py、coco_eval.py 这些文件把训练循环、COCO 评估和推理逻辑都拆开了，便于按需改动。适合有 PyTorch 基础、准备用自定义卫星或航拍数据做目标检测与分割的开发者，也适合在搞 GIS 或遥感应用但对深度学习管道细节不够熟的人，顺着代码把整条链路走通。

2. Mask R-CNN 选型与数据预处理链路搭建

2.1 建筑物检测场景下为什么选 Mask R-CNN 而不是 YOLO 或 SSD

先明确一点：建筑物检测的目标不是单纯画一个框。城市遥感影像里建筑物往往挨得很近，边界轮廓又规则，单一的目标框很容易把两栋楼框成一个整体。Mask R-CNN 在 Faster R-CNN 的检测分支之外，额外引入了一个并行的 mask 分支，对每个 RoI 输出二值分割掩码，这样既拿到类别和边界框，又拿到像素级的轮廓。从项目文件里 building_detect.py 的存在也能看出，作者把检测和掩码生成绑定在同一个前向流程里，而不是像 YOLO 只输出框。

Mask R-CNN 在 PyTorch 生态里可以直接用 torchvision 提供的预训练权重maskrcnn_resnet50_fpn，它由 ResNet-50 提取特征、FPN 做多尺度融合、RPN 生成候选框，再经过 RoIAlign 对齐后分别送入分类/回归头和 mask 头。ResNet-50 作为骨干网络对建筑物这种纹理相对单一、结构边缘清晰的目标来说，特征提取能力已经够用，预训练权重在 COCO 上学习到的底层视觉特征可以迁移到遥感影像上。FPN 的多尺度特性对大小楼体混在一起的场景很关键，小房子和大楼宇在不同的金字塔层被感知，Recall 会明显高过单尺度模型。

相比之下，YOLO 速度快但实例分割能力弱，SSD 在密集小目标上表现不够稳。建筑物识别这种既要框又要轮廓的任务，Mask R-CNN 是最稳妥的选择。代价是显存占用更高、推理更慢，但在遥感场景里通常对实时性要求不高，优先保证精度更合理。

2.2 COCO 数据格式约束与 transforms.py 的增强策略

这套工程沿用了 torchvision 参考实现里的 COCO 数据流程，coco_utils.py 负责处理 COCO 标注格式的加载和转换。COCO 格式中图像对应一个标注列表，每个标注包含bbox、area、iscrowd、category_id，本例中的category_id只有两类：1 表示建筑物，0 为背景。数据加载时，代码会把 COCO 的 bbox 从[x, y, width, height]转为[x1, y1, x2, y2]，用于模型训练时的目标编码。

transforms.py 做的是输入图像的预处理和数据增强。最核心的是一个ToTensor转换，把 HWC 排列、取值 0 到 255 的 numpy 数组或 PIL 图像转成 CHW 排列、0.0 到 1.0 的浮点张量。另一个常用增强是随机水平翻转：

import torch import torchvision.transforms as T import torchvision.transforms.functional as F class RandomHorizontalFlip(T.RandomHorizontalFlip): def forward(self, image, target): if torch.rand(1) < self.p: image = F.hflip(image) if target is not None: # 边界框的水平翻转变换，w 为图像宽度 boxes = target["boxes"] boxes[:, [0, 2]] = image.shape[-1] - boxes[:, [2, 0]] target["boxes"] = boxes if "masks" in target: target["masks"] = target["masks"].flip(-1) return image, target

翻转增强的核心问题是：图像翻转后，目标框和掩码也必须同步变换，否则模型训练时看到的标注和图像内容对不上。代码里对 boxes 做了坐标重映射，xmin, xmax变成镜像位置，masks 沿宽度方向翻转。这种做法在建筑物数据集里尤其有用，因为建筑朝向没有固定语义，翻转不改变类别含义，等于白赚一倍训练样本。

2.3 helpers.py 与 utils.py 里的工具函数管线

工程里 helpers.py 和 utils.py 承担了零散的辅助逻辑。常见的 helpers 包括：将预测出的 tensor 形式的 boxes 转成列表、把 tensor 数据移到指定设备、在 tensor 上按阈值过滤低分检测结果。utils.py 则会提供平滑学习率相关的工具函数。在 torchvision 参考实现里，warmup_lr_scheduler会在训练前几个 epoch 从 1/3 初始学习率开始线性爬升，避免模型在开局阶段因学习率过大而震荡。这个 warmup 逻辑被用到自定义项目中是很常见的做法，尤其是用预训练权重微调时，骨干网络收敛状态和检测头不一致，一上来就用大学习率容易把 FPN 层参数带偏。

数据准备阶段的常见做法如下：

# 数据集目录结构建议 dataset/ train/ images/ # 卫星图或航拍图，支持 .png .jpg .tif annotations.json # COCO 格式标注 val/ images/ annotations.json

如果手上只有 TIF 格式的高分辨率影像，先通过 GDAL 或 rasterio 做切片预处理，切成 512x512 或 1024x1024 的图块再生成标注，否则整幅 TIF 直接读入显存基本会 OOM。transform 里还应当加入归一化，虽然 Mask R-CNN 的预训练权重要求输入像素范围落在 0 到 1 之间，但遥感影像经常是 16 位深度，需要先转成 8 位或做线性拉伸，否则对比度过低时模型很难学到有效的边缘特征。

3. 训练链路：engine.py、train.py 与 train_tif.py 的参数细节

3.1 优化器与学习率调度策略

train.py 是训练入口，它的主要职责是解析命令行参数、构建 datasets 和 dataloader、初始化模型，然后把训练循环交给 engine.py。项目默认优化器选择 SGD，动量设为 0.9，权重衰减 0.0005。对检测任务来说，SGD 配合动量在大多数情况下收敛行为比 Adam 更稳，尤其是使用了预训练权重时，SGD 不容易把已经学好的特征破坏掉。

学习率调度采用 step 式下降：初始学习率 0.005 或 0.01，在指定的 epoch 位置（通常为 8 和 11）缩放到原来的 0.1。一个典型的 12 epoch 训练计划如下表所示。

如果显存不够，batch_size 降到 1，学习率同时降到 0.0025 左右，否则同等学习率下梯度更新步长相对变大，容易出现 loss 爆炸。在遥感数据集上我一般会用 0.0005 到 0.001 的微调学习率，因为 COCO 预训练权重已经足够强，过大的学习率会把 FPN 的特征尺度打乱。

3.2 engine.py 训练循环与损失函数构成

engine.py 实现了一个 epoch 内的训练和验证流程。每个 batch 的图像和标注送入模型，模型前向时同时返回 loss 字典和预测结果。训练阶段只保留 loss：

def train_one_epoch(model, optimizer, data_loader, device, epoch): model.train() for images, targets in data_loader: images = [img.to(device) for img in images] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) optimizer.zero_grad() losses.backward() optimizer.step()

model(images, targets)在训练模式下返回一个字典，包含四项：loss_classifier是 RPN 和检测头的分类损失，loss_box_reg是边界框回归的 SmoothL1 损失，loss_mask是掩码分支的 BCE 损失，loss_objectness和loss_rpn_box_reg来自 RPN 的候选框生成质量。总 loss 是这几项的加权和，torchvision 默认权重都是 1.0。

每个 epoch 结束后需要保存 checkpoint。常见做法是只保留验证集 mAP 最高的模型，或者保存最后一次权重，同时把优化器状态一并存下来，方便中断后恢复训练：

torch.save({ 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch, 'lr': optimizer.param_groups[0]['lr'] }, f'checkpoint_epoch_{epoch}.pth')

3.3 train_tif.py 对遥感 TIF 影像的适配

项目里单独拆出的 train_tif.py 是为了处理 TIF 格式遥感影像的训练场景。和普通 JPEG/PNG 相比，TIF 往往位深更高、尺寸极大，无法直接整图读入网络中。这个脚本会在读取阶段引入一个 tile 切分逻辑，先把大尺寸 TIF 按固定步长切成重叠的 patch，每个 patch 作为一个独立训练样本。

切分时需要注意几个点。patch 大小设成 512 或 768 比较合理，再大模型处理时 FPN 的内存消耗也会变大。步长可以小于 patch 大小，形成重叠区域，缓解建筑物被切在边缘的问题。被切断的建筑物标注面积如果太小，比如小于 20x20 像素，建议直接丢弃，否则标注噪声会干扰 RPN 训练。TIF 的通道数也可能不是 3，此时要看是 RGB 三波段还是 RGBN 四波段。四波段影像中近红外通道对建筑物和植被的区分很有用，但预训练 ResNet 接受的是三通道输入，要么取前三个波段，要么在工程中自行修改模型的输入 conv 层权重。

train_tif.py 中通常还要实现对 16 位 TIF 的归一化拉伸。简单做法是读取波段数据后计算 2% 和 98% 分位数，做线性映射到 0-255，再转成三通道图。这一步对模型稳定收敛的作用很直接：16 位影像的像素值范围过大，直接输入网络会让 BN 层的统计量偏差很大，预训练权重作废。

4. 推理、评估与可视化：detect.py、eval.py 与 coco_eval.py

4.1 detect.py 的预测流程与 score 阈值选择

detect.py 是推理入口。加载训练好的权重后，对单张或多张输入图像执行模型前向，拿到包含boxes、labels、scores、masks的字典。模型默认只会输出经过 NMS 后的前 100 个预测结果，但仍然需要按 score 过滤才能在业务里使用：

import torch import torchvision from PIL import Image from torchvision.models.detection import maskrcnn_resnet50_fpn device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = maskrcnn_resnet50_fpn(weights=None, num_classes=2) checkpoint = torch.load('checkpoint_final.pth', map_location=device) model.load_state_dict(checkpoint['model_state_dict']) model.to(device) model.eval() image = Image.open('aerial.png').convert('RGB') transform = torchvision.transforms.ToTensor() image_tensor = transform(image).to(device) with torch.no_grad(): pred = model([image_tensor]) boxes = pred[0]['boxes'].cpu().numpy() scores = pred[0]['scores'].cpu().numpy() masks = pred[0]['masks'].cpu().numpy().squeeze(1) score_threshold = 0.5 keep = scores >= score_threshold boxes = boxes[keep] masks = masks[keep] scores = scores[keep]

score 阈值控制在 0.5 是通用选择，但遥感场景我会分情况调整：如果只做粗略的建筑面积统计，阈值可以放到 0.3 以召回更多小楼；如果做精确到楼栋的地图更新，0.7 甚至 0.8 更合适，代价是召回变低。推理时model.eval()必须调用，它关闭了 BatchNorm 的统计更新和 Dropout，否则同一张图每次预测结果都会不一样。

4.2 eval.py 与 coco_eval.py 的指标计算逻辑

eval.py 和 coco_eval.py 组合起来做的是 COCO 标准评估。coco_eval.py是 torchvision 参考实现里对 COCO 官方评估代码的封装，计算 mAP 和 mAR 等指标。COCO 评估的核心是 AP，它把 IoU 阈值从 0.5 到 0.95 每 0.05 取一个档位，在每个档位上计算 precision-recall 曲线下的面积，再求平均。

精度和召回率的定义对检测结果质量影响最大，也是最容易被忽略的部分。同一个模型，用 AP@0.5 衡量可能表现极好，用 AP@0.75 衡量就掉得厉害，说明它的定位精度不够，框的位置偏移较大。对建筑物识别，框的偏移直接影响面积估算，所以我倾向于把 AP@0.75 和 AP@0.5:0.95 一起关注。

运行评估的方式是把验证集预测结果转成 COCO 格式的 JSON，然后通过 COCO API 计算指标：

python eval.py --weights checkpoint_final.pth --data val --output pred_results.json

eval.py 会遍历验证集每一张图像，保存预测框、分数、掩码，然后调用 coco_eval.py 里的COCOeval类完成指标计算。输出里每一行对应一个小类的 AP 值，汇总行mAP是各类平均。

4.3 从 output.png 观察模型预测行为

项目根目录的 output.png 是模型在某一批测试图像上的可视化结果。生成这种可视化通常是把预测的 boxes 画在原始图上，同时把 mask 以半透明方式叠加。对输出图的观察有一个容易被忽略的技巧：不要只看检测对不对，还要看 mask 边缘是否贴合建筑物轮廓。如果 mask 锯齿严重或大面积溢到周边空地，说明训练数据里的掩码标注质量不够，或者模型在 mask 分支的 loss 权重偏低，可以尝试在总 loss 中加大loss_mask的权重系数。

绘制边界框和 mask 时可以使用 OpenCV 或 matplotlib，但注意 mask 的输出是一个概率值矩阵，需要用阈值二值化再叠加：

import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as patches mask = masks[0] # 建议保留 5 个边界检查 binary_mask = mask > 0.5 plt.imshow(image_np) plt.imshow(binary_mask, alpha=0.4, cmap='Reds') rect = patches.Rectangle((boxes[0][0], boxes[0][1]), boxes[0][2] - boxes[0][0], boxes[0][3] - boxes[0][1], linewidth=2, edgecolor='lime', facecolor='none') plt.gca().add_patch(rect) plt.savefig('output.png', dpi=150, bbox_inches='tight')

5. 用 diff.py 做模型回归验证与误检控制

diff.py 在项目里承担的是一个常被作者们忽略但实际很实用的功能：对比两个模型在相同输入上的输出差异。典型使用场景是训练了一个新 checkpoint 后，不想全量跑 eval.py，先挑几组代表性图像，把新旧模型的预测结果做逐框对比，找出新版本相比旧版本哪些目标被漏掉、哪些被新增检测出来，哪些框的位置发生了明显偏移。

具体做法是把两个模型分别跑一遍前向，按类别和位置做 IoU 配对，如果一个框在新模型中出现但旧模型没有，标记为new_fp；反过来标记为new_fn；IoU 低于 0.5 但面积相近的标记为shifted。这样能快速判断一次参数调整的方向对不对。比如减小 score 阈值后，new_fp数量激增，说明模型在低分区间产生了大量噪声；而如果new_fn减少明显，说明召回确实提升，这时候再去权衡阈值，就有了量化依据。

diff.py 里常见的核心判断逻辑大致是这样的：

def match_boxes(gt_boxes, dt_boxes, iou_threshold=0.5): matched_dt = set() matched_gt = set() for i, gt in enumerate(gt_boxes): for j, dt in enumerate(dt_boxes): iou_val = compute_iou(gt, dt) if iou_val >= iou_threshold and j not in matched_dt: matched_dt.add(j) matched_gt.add(i) break return matched_gt, matched_dt

误检的地类通常集中在阴影、停车场、混凝土空地这些与建筑物颜色纹理相近的区域。想减少这类 false positive，首选方法是在训练集中主动加入不含建筑物的负样本 patch，让 RPN 学会抑制这些区域的 proposal。另一个做法是在后处理阶段加一个面积过滤，删除明显小于真实最小建筑面积的预测框——这个阈值从验证集的标注分布里统计出来即可，不必拍脑袋。

在做模型上线前，diff.py 配合 COCO 评估是个高效的组合：先跑 diff 做定性分析，再跑 eval.py 拿定量指标。最终选择哪个 checkpoint 上线，我一般会看综合指标而不是 mAP 单值，把mAP@0.5:0.95、mAP@0.75、AR@100三个值一起拉出来权衡，同时回看 diff 输出确认没有明显的新增误检。建筑物识别落在遥感业务里通常对接的是面积统计或变化检测，这种回归验证比单纯追求指标数字更有参考价值。

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