Urban_seg: 基于遥感图像的高效语义分割项目

Urban_seg:让遥感图像语义分割变得简单高效

遥感图像语义分割是一项具有挑战性的任务,通常需要大量带标注的训练数据和复杂的模型才能取得良好效果。然而,近日GitHub上开源的Urban_seg项目为我们展示了一种简单高效的解决方案,仅需4张遥感图像就能训练出性能优异的分割模型。本文将深入介绍这个新手友好的项目,包括其背景、特点、使用方法等内容。

项目背景与亮点

Urban_seg项目由GitHub用户anxiangsir开发并开源,旨在为遥感图像语义分割任务提供一个简单易用的解决方案。该项目的核心亮点包括:

基于预训练模型:利用在4亿张图像上预训练的unicom模型作为基础,大大降低了对训练数据的需求。
极少数据即可训练:仅需4张带标注的遥感图像就能获得优秀的分割效果,极大降低了数据标注成本。
新手友好:代码简洁易懂,仅200行左右即可实现单GPU训练。
性能出色:尽管使用极少数据,但分割效果仍然非常优秀,展现了强大的泛化能力。
灵活可扩展:提供了单GPU和多GPU两种训练方式,可根据需求选择。

这些特点使得Urban_seg成为一个非常适合新手入门和快速验证想法的遥感图像分割项目。

预测效果图

安装与环境配置

要开始使用Urban_seg,首先需要克隆项目并安装所需依赖:

git clone https://github.com/anxiangsir/urban_seg.git
cd urban_seg
pip install -r requirements.txt

项目依赖的主要环境包括Python 3、PyTorch等常用深度学习库。

数据准备

Urban_seg使用的数据集来自CCF卫星影像的AI分类与识别比赛,包含5张带标注的遥感图像。您可以从百度网盘(密码:3ih2)下载数据集。

下载完成后,将数据集解压到项目根目录,目录结构如下:

dataset
├── origin  # 5张带标注的遥感图像
├── test    # 3张无标注的测试图像(本项目未使用)
└── train   # 通过预处理脚本生成的训练数据
    ├── images       
    └── labels

除了图像数据,您还需要下载预训练的ViT模型权重文件,包括:

FP16-ViT-B-32.pt
FP16-ViT-B-16.pt
FP16-ViT-L-14.pt
FP16-ViT-L-14-336px.pt

这些文件同样应放置在项目根目录下。

数据预处理

在开始训练之前,需要对原始数据进行预处理。Urban_seg提供了一个简单的预处理脚本preprocess.py,运行它即可自动完成数据处理:

python preprocess.py

这个脚本会从原始图像中随机采样生成训练集,并将结果保存在dataset/train目录下。

模型训练

Urban_seg提供了两种训练方式:单GPU训练和多GPU训练。

单GPU训练

对于大多数用户来说,单GPU训练已经足够。运行以下命令即可开始训练:

python train_one_gpu.py

train_one_gpu.py是一个仅200行左右的简洁脚本,非常适合新手学习和理解整个训练流程。

多GPU训练

如果您有多个GPU资源并希望加速训练,可以使用多GPU训练脚本:

torchrun --nproc_per_node 8 train_multi_gpus.py

这里的8表示使用8个GPU,您可以根据实际情况调整。注意,多GPU训练可能需要额外的配置,请仔细阅读代码中的注释说明。

训练过程可视化

Urban_seg在训练过程中会输出丰富的日志信息,包括loss、学习率、各类指标等。您可以通过这些信息实时监控训练进展。

此外,项目还提供了训练结果的可视化展示。下面是一个预测效果的动态展示:

预测效果动图

模型评估与应用

训练完成后,您可以使用项目提供的评估脚本对模型性能进行评估。评估指标包括像素级准确率、IoU等常用分割指标。

对于实际应用,您可以参考项目中的推理代码,将训练好的模型应用于新的遥感图像。下面是一个分割效果示例:

分割效果示例

进阶优化建议

尽管Urban_seg使用极少的数据就能取得不错的效果,但如果您希望进一步提升模型性能,可以考虑以下优化方向:

增加训练数据:虽然4张图像就能工作,但更多的高质量训练数据通常能带来更好的效果。
数据增强:尝试更多的数据增强技术,如旋转、翻转、色彩抖动等,提高模型泛化能力。
模型优化:尝试其他backbone网络或分割头,如DeepLabV3+、UNet++等。
超参数调优:使用网格搜索或贝叶斯优化等方法寻找最佳超参数组合。
集成学习:训练多个模型并集成预测结果,可能会带来额外的性能提升。

社区交流与支持

Urban_seg项目维护了一个活跃的QQ群(679897018),您可以在这里与其他用户和开发者交流经验、反馈问题。此外,GitHub仓库的Issue区也是寻求帮助和讨论的好地方。

总结与展望

Urban_seg项目为遥感图像语义分割任务提供了一个简单高效的解决方案。它的成功证明了预训练模型的强大潜力,以及在特定任务上微调的有效性。这种方法不仅大大降低了对标注数据的需求,还简化了模型训练过程,使得即使是新手也能快速上手并获得不错的效果。

展望未来,我们可以期待看到更多类似Urban_seg这样的项目,它们将先进的AI技术与特定领域任务相结合,为各行各业带来创新和效率提升。对于遥感图像处理领域来说,这意味着更快速、更精确的地理信息提取,这将在城市规划、环境监测、灾害预警等众多领域发挥重要作用。

如果您对遥感图像分割感兴趣,不妨尝试使用Urban_seg项目,相信它会为您的研究或应用带来新的灵感和可能性。

📝 参考文献:

An, X., Deng, J., Yang, K., Li, J., Feng, Z., Guo, J., Yang, J., & Liu, T. (2023). Unicom: Universal and Compact Representation Learning for Image Retrieval. ICLR.
Urban_seg GitHub仓库
CCF卫星影像的AI分类与识别竞赛