RobustSAM의 공식 저장소: 저하된 이미지에서 무엇이든 강력하게 분할
프로젝트 페이지 | 종이 | 비디오 | 데이터세트
2024년 8월: 이 링크를 통해 더 쉽게 사용할 수 있도록 @jadechoghari가 제작한 Hugging Face 모델 카드 및 데모를 참조할 수 있습니다.
2024년 7월: 다양한 ViT 백본에 대한 학습 코드, 데이터 및 모델 체크포인트가 출시되었습니다!
2024년 6월: 추론 코드가 출시되었습니다!
2024년 2월: RobustSAM이 CVPR 2024에 승인되었습니다!
SAM(Segment Anything Model)은 이미지 분할의 혁신적인 접근 방식으로 등장했으며, 강력한 제로샷 분할 기능과 유연한 프롬프트 시스템으로 호평을 받았습니다. 그럼에도 불구하고 품질이 저하된 이미지로 인해 성능이 저하됩니다. 이러한 한계를 해결하기 위해 우리는 SAM의 신속성과 제로샷 일반화를 유지하면서 저품질 이미지에 대한 SAM의 성능을 향상시키는 RobustSegment Anything Model(RobustSAM)을 제안합니다.
우리의 방법은 한계 매개변수 증가와 계산 요구 사항만으로 사전 훈련된 SAM 모델을 활용합니다. RobustSAM의 추가 매개변수는 8개의 GPU에서 30시간 이내에 최적화될 수 있어 일반적인 연구 실험실에 대한 타당성과 실용성을 입증합니다. 또한 모델을 최적으로 훈련하고 평가하기 위해 설계된 다양한 성능 저하가 있는 688K 이미지-마스크 쌍 모음인 Robust-Seg 데이터 세트를 소개합니다. 다양한 세분화 작업과 데이터 세트에 대한 광범위한 실험을 통해 특히 제로샷 조건에서 RobustSAM의 탁월한 성능이 확인되었으며 광범위한 실제 적용 가능성이 강조되었습니다. 또한, 우리의 방법은 단일 이미지 디헤이징 및 디블러링과 같은 SAM 기반 다운스트림 작업의 성능을 효과적으로 향상시키는 것으로 나타났습니다.
콘다 환경을 생성하고 활성화합니다.
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 데이터세트를 다운로드하세요. (참고: train, val 및 test 데이터 세트의 이미지는 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
데모용으로 몇 가지 이미지를 데모_이미지 폴더에 준비했습니다. 또한 두 가지 프롬프트 모드(박스 프롬프트와 포인트 프롬프트)를 사용할 수 있습니다.
상자 프롬프트의 경우:
python eval.py --bbox --model_size l
포인트 프롬프트의 경우:
python eval.py --model_size l
기본적으로 데모 결과는 deco_result/[$PROMPT_TYPE] 에 저장됩니다.
이 작업이 유용하다고 생각되면 우리를 인용해 보세요!
@inproceedings{chen2024robustsam, title={RobustSAM: 저하된 이미지에서 무엇이든 강력하게 분할}, 작성자={Chen, Wei-Ting 및 Vong, Yu-Jiet 및 Kuo, Sy-Yen 및 Ma, Sizhou 및 Wang, Jian}, 저널= {CVPR}, 연도={2024}}
우리 저장소의 기반이 되는 SAM의 작성자에게 감사드립니다.