RobustSAM
1.0.0
Официальный репозиторий RobustSAM: надежная сегментация любых изображений на поврежденных изображениях
Страница проекта | Бумага | Видео | Набор данных
Август 2024 г.: для упрощения использования вы можете обратиться к карточкам и демо-версии модели Hugging Face, созданным @jadechoghari, по этой ссылке.
Июль 2024 г.: выпущен обучающий код, данные и контрольные точки модели для различных магистральных сетей ViT!
Июнь 2024 г.: выпущен код вывода!
Февраль 2024 г.: RobustSAM принят в CVPR 2024!
Модель Segment Anything (SAM) возникла как революционный подход к сегментации изображений, получивший признание благодаря своим надежным возможностям сегментации с нулевым выстрелом и гибкой системе подсказок. Тем не менее, его производительность снижается из-за изображений с ухудшенным качеством. Устраняя это ограничение, мы предлагаем модель Robust Segment Anything Model (RobustSAM), которая повышает производительность SAM на изображениях низкого качества, сохраняя при этом ее оперативность и обобщение с нулевым выстрелом.
Наш метод использует предварительно обученную модель SAM только с незначительными приращениями параметров и вычислительными требованиями. Дополнительные параметры RobustSAM можно оптимизировать в течение 30 часов на восьми графических процессорах, демонстрируя его осуществимость и практичность для типичных исследовательских лабораторий. Мы также представляем набор данных Robust-Seg, коллекцию пар изображений-масок из 688 тысяч с различными ухудшениями, предназначенную для оптимального обучения и оценки нашей модели. Обширные эксперименты с различными задачами сегментации и наборами данных подтверждают превосходную производительность RobustSAM, особенно в условиях нулевого выстрела, подчеркивая его потенциал для широкого практического применения. Кроме того, было показано, что наш метод эффективно повышает производительность последующих задач на основе SAM, таких как устранение дымки и размытия одного изображения.
Создайте среду conda и активируйте ее.
conda create --name robustsam python=3.10 -y conda activate robustsam
Клонируйте и войдите в каталог репо.
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 разных размеров и поместите их в текущий каталог.
КПП ВИТ-Б РобустСАМ
КПП ВиТ-Л РобустСАМ
КПП ВиТ-Х РобустСАМ
Измените текущий каталог на каталог «данные».
cd data
Загрузите train, val, test и дополнительный набор данных 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
Мы подготовили несколько изображений в папке demo_images для демонстрационных целей. Кроме того, доступны два режима подсказок (подсказки в рамке и подсказки в виде точек).
Для подсказки коробки:
python eval.py --bbox --model_size l
Для подсказки точки:
python eval.py --model_size l
По умолчанию результаты демонстрации будут сохранены в demo_result/[$PROMPT_TYPE] .
 |  |
 |  |
Если вы считаете эту работу полезной, пожалуйста, цитируйте нас!
@inproceedings{chen2024robustsam, title={RobustSAM: надежная сегментация чего угодно на испорченных изображениях}, автор={Чен, Вэй-Тин и Вонг, Ю-Цзиет и Куо, Си-Йен и Ма, Сычжоу и Ван, Цзянь}, журнал= {CVPR}, год={2024}}Мы благодарим авторов SAM, на основе которого основан наш репозиторий.
