项目介绍
urban_seg 是一个专为初学者设计的遥感图像语义分割项目。其核心是利用在4亿张图像上预训练的unicom模型,这一模型在处理遥感分割任务方面表现卓越。令人兴奋的是,该项目在训练时仅使用了4张遥感图像,却取得了出色的效果。
快速开始
对于希望快速上手的用户,可以使用简洁的 train_one_gpu.py
脚本开始训练,该脚本仅有200行代码。而对于追求更好性能的用户,可以使用支持多GPU训练的 train_multi_gpus.py
脚本。请注意,在使用 train_multi_gpus.py
时,可能需要进行一些额外的配置和设置,因此务必仔细阅读代码中的说明和文档,确保正确配置。
安装指南
首先,用户需要克隆项目并安装必要的依赖包:
git clone https://github.com/anxiangsir/urban_seg.git
pip install -r requirements.txt
数据和预训练模型
项目使用了CCF卫星影像AI分类与识别提供的五张卫星遥感图像数据集。数据集初赛复赛训练集可从指定的百度云盘下载。
数据结构如下:
dataset
├── origin //五张带标签的遥感图片
├── test //三张无标签的遥感图片(本项目中未使用)
└── train //为空,通过`python preprocess.py`生成
├── images
└── labels
FP16-ViT-B-32.pt
FP16-ViT-B-16.pt
FP16-ViT-L-14.pt
FP16-ViT-L-14-336px.pt
一张GPU训练流程
-
下载数据集到当前目录。
-
预处理数据:
python preprocess.py
-
开始训练:
python train_one_gpu.py
多GPU训练流程
-
下载数据集到当前目录。
-
预处理数据:
python preprocess.py
-
启动训练:
torchrun --nproc_per_node 8 train_multi_gpus.py
讨论与反馈
如果用户对项目有任何疑问或意见,欢迎加入QQ群参与讨论:679897018。
引用
如果该项目对您的研究有所帮助,欢迎引用我们的论文:
@inproceedings{anxiang_2023_unicom,
title={Unicom: Universal and Compact Representation Learning for Image Retrieval},
author={An, Xiang and Deng, Jiankang and Yang, Kaicheng and Li, Jiawei and Feng, Ziyong and Guo, Jia and Yang, Jing and Liu, Tongliang},
booktitle={ICLR},
year={2023}
}
通过以上内容,用户可以轻松了解 urban_seg 项目的特点和使用方法,并快速开始进行遥感图像的语义分割研究。