RobustSAM
1.0.0
Repositório oficial do RobustSAM: segmente qualquer coisa de maneira robusta em imagens degradadas
Página do Projeto | Papel | Vídeo | Conjunto de dados
Agosto de 2024: Você pode consultar os cartões e demonstração do modelo Hugging Face criados por @jadechoghari para facilitar o uso por meio deste link.
Julho de 2024: Código de treinamento, dados e pontos de verificação de modelo para diferentes backbones ViT são lançados!
Junho de 2024: O código de inferência foi lançado!
Fevereiro de 2024: RobustSAM foi aceito no CVPR 2024!
Segment Anything Model (SAM) surgiu como uma abordagem transformadora na segmentação de imagens, aclamada por seus robustos recursos de segmentação zero-shot e sistema de prompt flexível. No entanto, seu desempenho é desafiado por imagens com qualidade degradada. Abordando essa limitação, propomos o Robust Segment Anything Model (RobustSAM), que melhora o desempenho do SAM em imagens de baixa qualidade, preservando sua prontabilidade e generalização zero-shot.
Nosso método aproveita o modelo SAM pré-treinado apenas com incrementos marginais de parâmetros e requisitos computacionais. Os parâmetros adicionais do RobustSAM podem ser otimizados em 30 horas em oito GPUs, demonstrando sua viabilidade e praticidade para laboratórios de pesquisa típicos. Também apresentamos o conjunto de dados Robust-Seg, uma coleção de 688K pares de máscara de imagem com diferentes degradações projetadas para treinar e avaliar nosso modelo de forma otimizada. Extensos experimentos em várias tarefas de segmentação e conjuntos de dados confirmam o desempenho superior do RobustSAM, especialmente sob condições de disparo zero, ressaltando seu potencial para ampla aplicação no mundo real. Além disso, nosso método demonstrou melhorar efetivamente o desempenho de tarefas downstream baseadas em SAM, como desfocagem e desfoque de imagem única.
Crie um ambiente conda e ative-o.
conda create --name robustsam python=3.10 -y conda activate robustsam
Clone e entre no diretório repo.
git clone https://github.com/robustsam/RobustSAM cd RobustSAM
Use o comando abaixo para verificar sua versão CUDA.
nvidia-smi
Substitua a versão CUDA pela sua no comando abaixo.
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
Instale as dependências restantes
pip install -r requirements.txt
Baixe pontos de verificação RobustSAM pré-treinados de diferentes tamanhos e coloque-os no diretório atual.
Ponto de verificação ViT-B RobustSAM
Ponto de verificação ViT-L RobustSAM
Ponto de verificação ViT-H RobustSAM
Mude o diretório atual para o diretório "data".
cd data
Baixe trem, val, teste e conjunto de dados COCO e LVIS extra. (NOTA: As imagens no conjunto de dados train, val e test consistem em imagens de LVIS, MSRA10K, ThinObject-5k, NDD20, STREETS e FSS-1000)
bash download.sh
Existem apenas imagens nítidas baixadas na etapa anterior. Use o comando abaixo para gerar imagens degradadas correspondentes.
bash gen_data.sh
Se você quiser treinar do zero, use o comando abaixo.
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
Se você quiser treinar a partir do ponto de verificação pré-treinado, use o comando abaixo.
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
Preparamos algumas imagens na pasta demo_images para fins de demonstração. Além disso, dois modos de prompt estão disponíveis (prompts de caixa e prompts de ponto).
Para prompt de caixa:
python eval.py --bbox --model_size l
Para prompt de ponto:
python eval.py --model_size l
Por padrão, os resultados da demonstração serão salvos em demo_result/[$PROMPT_TYPE] .
 |  |
 |  |
Se você achar este trabalho útil, considere nos citar!
@inproceedings{chen2024robustsam, title={RobustSAM: Segmentar qualquer coisa robustamente em imagens degradadas}, autor={Chen, Wei-Ting e Vong, Yu-Jiet e Kuo, Sy-Yen e Ma, Sizhou e Wang, Jian}, diário= {CVPR}, ano={2024}}Agradecemos aos autores do SAM no qual nosso repositório se baseia.
