Dépôt officiel de RobustSAM : segmentez n'importe quoi de manière robuste sur des images dégradées
Page du projet | Papier | Vidéo | Ensemble de données
Août 2024 : Vous pouvez vous référer aux cartes modèles et à la démo Hugging Face créées par @jadechoghari pour une utilisation plus facile via ce lien.
Juillet 2024 : le code de formation, les données et les points de contrôle du modèle pour différents backbones ViT sont publiés !
Juin 2024 : le code d'inférence a été publié !
Février 2024 : RobustSAM a été accepté dans le CVPR 2024 !
Segment Anything Model (SAM) est apparu comme une approche transformatrice de la segmentation d'images, acclamée pour ses robustes capacités de segmentation sans prise de vue et son système d'invite flexible. Néanmoins, ses performances sont mises à mal par des images de qualité dégradée. Pour remédier à cette limitation, nous proposons le modèle Robust Segment Anything (RobustSAM), qui améliore les performances de SAM sur les images de faible qualité tout en préservant sa rapidité et sa généralisation sans prise de vue.
Notre méthode exploite le modèle SAM pré-entraîné avec seulement des incréments de paramètres marginaux et des exigences de calcul. Les paramètres supplémentaires de RobustSAM peuvent être optimisés en 30 heures sur huit GPU, démontrant sa faisabilité et son caractère pratique pour les laboratoires de recherche typiques. Nous introduisons également l'ensemble de données Robust-Seg, une collection de 688 000 paires image-masque avec différentes dégradations conçues pour entraîner et évaluer notre modèle de manière optimale. Des expériences approfondies sur diverses tâches de segmentation et ensembles de données confirment les performances supérieures de RobustSAM, en particulier dans des conditions de tir nul, soulignant son potentiel pour des applications étendues dans le monde réel. De plus, il a été démontré que notre méthode améliore efficacement les performances des tâches en aval basées sur SAM, telles que le dévoilage et la suppression du flou d'une image unique.
Créez un environnement conda et activez-le.
conda create --name robustsam python=3.10 -y conda activate robustsam
Clonez et entrez dans le répertoire du dépôt.
git clone https://github.com/robustsam/RobustSAM cd RobustSAM
Utilisez la commande ci-dessous pour vérifier votre version CUDA.
nvidia-smi
Remplacez la version CUDA par la vôtre dans la commande ci-dessous.
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
Installer les dépendances restantes
pip install -r requirements.txt
Téléchargez des points de contrôle RobustSAM pré-entraînés de différentes tailles et placez-les dans le répertoire actuel.
Point de contrôle ViT-B RobustSAM
Point de contrôle ViT-L RobustSAM
Point de contrôle ViT-H RobustSAM
Changez le répertoire actuel en répertoire "data".
cd data
Téléchargez l'ensemble de données train, val, test et supplémentaire COCO & LVIS. (REMARQUE : les images des ensembles de données train, val et test sont constituées d'images de LVIS, MSRA10K, ThinObject-5k, NDD20, STREETS et FSS-1000)
bash download.sh
Il n'y a que des images claires téléchargées à l'étape précédente. Utilisez la commande ci-dessous pour générer les images dégradées correspondantes.
bash gen_data.sh
Si vous souhaitez vous entraîner à partir de zéro, utilisez la commande ci-dessous.
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
Si vous souhaitez vous entraîner à partir d'un point de contrôle pré-entraîné, utilisez la commande ci-dessous.
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
Nous avons préparé quelques images dans le dossier demo_images à des fins de démonstration. En outre, deux modes d'invite sont disponibles (invites de boîte et invites de point).
Pour l'invite de boîte :
python eval.py --bbox --model_size l
Pour l'invite de point :
python eval.py --model_size l
Par défaut, les résultats de la démonstration seront enregistrés dans demo_result/[$PROMPT_TYPE] .
Si vous trouvez ce travail utile, pensez à nous citer !
@inproceedings{chen2024robustsam, title={RobustSAM : segmentez n'importe quoi de manière robuste sur des images dégradées}, author={Chen, Wei-Ting et Vong, Yu-Jiet et Kuo, Sy-Yen et Ma, Sizhou et Wang, Jian}, journal= {CVPR}, année={2024}}
Nous remercions les auteurs de SAM sur lesquels notre dépôt est basé.