RobustSAM
1.0.0
RobustSAM 的官方存储库：在降级图像上稳健地分割任何内容
项目页面|纸|视频 |数据集
2024 年 8 月：您可以通过此链接参考 @jadechoghari 构建的 Hugging Face 模型卡和演示，以便更轻松地使用。
2024年7月：针对不同ViT主干的训练代码、数据和模型检查点发布！
2024年6月：推理代码已发布！
2024 年 2 月：RobustSAM 被 CVPR 2024 接收！
分段任意模型 (SAM) 已成为图像分割领域的一种变革性方法，因其强大的零样本分割功能和灵活的提示系统而备受赞誉。尽管如此，其性能仍受到图像质量下降的挑战。为了解决这一限制，我们提出了鲁棒分段任意模型 (RobustSAM)，它增强了 SAM 在低质量图像上的性能，同时保留了其及时性和零样本泛化能力。
我们的方法利用预训练的 SAM 模型，仅具有边际参数增量和计算要求。 RobustSAM 的附加参数可以在 8 个 GPU 上在 30 小时内优化，这证明了其对于典型研究实验室的可行性和实用性。我们还介绍了 Robust-Seg 数据集，这是一个具有不同退化的 688K 图像掩模对的集合，旨在最佳地训练和评估我们的模型。跨各种分割任务和数据集的大量实验证实了 RobustSAM 的卓越性能，尤其是在零样本条件下，强调了其在现实世界中广泛应用的潜力。此外，我们的方法已被证明可以有效提高基于 SAM 的下游任务（例如单图像去雾和去模糊）的性能。
创建conda环境并激活它。
conda create --name robustsam python=3.10 -y conda activate robustsam
克隆并进入 repo 目录。
git clone https://github.com/robustsam/RobustSAM cd RobustSAM
使用下面的命令检查您的 CUDA 版本。
nvidia-smi
在下面的命令中将 CUDA 版本替换为您的版本。
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu[$YOUR_CUDA_VERSION] # For example: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 # cu117 = CUDA_version_11.7
安装剩余的依赖项
pip install -r requirements.txt
下载不同大小的预训练 RobustSAM 检查点并将它们放入当前目录。
ViT-B RobustSAM 检查点
ViT-L RobustSAM 检查点
ViT-H RobustSAM 检查点
将当前目录更改为“data”目录。
cd data
下载训练、验证、测试和额外的 COCO 和 LVIS 数据集。 （注：训练、验证和测试数据集中的图像包含来自 LVIS、MSRA10K、ThinObject-5k、NDD20、STREETS 和 FSS-1000 的图像）
bash download.sh
上一步中仅下载了清晰的图像。使用下面的命令生成相应的降级图像。
bash gen_data.sh
如果您想从头开始训练，请使用下面的命令。
python -m torch.distributed.launch train_ddp.py --multiprocessing-distributed --exp_name [$YOUR_EXP_NAME] --model_size [$MODEL_SIZE] # Example usage: python -m torch.distributed.launch train_ddp.py --multiprocessing-distributed --exp_name test --model_size l
如果您想从预先训练的检查点进行训练，请使用下面的命令。
python -m torch.distributed.launch train_ddp.py --multiprocessing-distributed --exp_name [$YOUR_EXP_NAME] --model_size [$MODEL_SIZE] --load_model [$CHECKPOINT_PATH] # Example usage: python -m torch.distributed.launch train_ddp.py --multiprocessing-distributed --exp_name test --model_size l --load_model robustsam_checkpoint_l.pth
python gradio_app.py
我们在demo_images文件夹中准备了一些图像用于演示目的。此外，还有两种提示方式（框提示和点提示）。
对于框提示：
python eval.py --bbox --model_size l
对于点提示：
python eval.py --model_size l
默认情况下，演示结果将保存到demo_result/[$PROMPT_TYPE] 。
 |  |
 |  |
如果您发现这项工作有用，请考虑引用我们！
@inproceedings{chen2024robustsam，title={RobustSAM：在退化图像上鲁棒地分割任何东西}，作者={Chen，Wei-Ting and Vong，Yu-Jiet and Kuo，Sy-Yen and Ma，Sizhou and Wang，Jian}，journal= {CVPR}，年份={2024}}我们感谢我们的存储库所基于的 SAM 的作者。
