Personalize Segment Anything with 1 Shot in 10 Seconds
Official implementation of 'Personalize Segment Anything Model with One Shot'.
News
- TODO: Release the PerSAM-assisted Dreambooth for better fine-tuning Stable Diffusion
π . - We release the code of PerSAM and PerSAM-F
π₯ . Check our video here! - We release a new dataset for personalized segmentation, PerSeg
π₯ .
Introduction
How to customize SAM to automatically segment your pet dog in a photo album?
In this project, we propose a training-free Personalization approach for Segment Anything Model (SAM), termed as PerSAM. Given only a single image with a reference mask, PerSAM can segment specific visual concepts, e.g., your pet dog, within other images or videos without any training. For better performance, we further present an efficient one-shot fine-tuning variant, PerSAM-F. We freeze the entire SAM and introduce two learnable mask weights, which only trains 2 parameters within 10 seconds.
Besides, our approach can be utilized to assist DreamBooth in fine-tuning better Stable Diffusion for personalized image synthesis. We adopt PerSAM to segment the target object in the user-provided few-shot images, which eliminates the background disturbance and benefits the target representation learning.
Requirements
Installation
Clone the repo and create a conda environment:
git clone https://github.com/ZrrSkywalker/Personalize-SAM.git
cd Personalize-SAM
conda create -n persam python=3.8
conda activate persam
pip install -r requirements.txt
Similar to Segment Anything, our code requires pytorch>=1.7
and torchvision>=0.8
. Please follow the instructions here to install both PyTorch and TorchVision dependencies.
Preparation
Please download our constructed dataset PerSeg for personalized segmentation from Google Drive or Baidu Yun (code 222k
), and the pre-trained weights of SAM from here. Then, unzip the dataset file and organize them as
data/
|ββ Annotations/
|ββ Images/
sam_vit_h_4b8939.pth
Getting Started
Personalized Segmentation
For the training-free π§ PerSAM, just run:
python persam.py --outdir <output filename>
For 10-second fine-tuning of
python persam_f.py --outdir <output filename>
After running, the output masks and visualzations will be stored at outputs/<output filename>
.
Then, for mIoU evaluation, please run:
python eval_miou.py --pred_path <output filename>
Personalized Stable Diffusion
Our approach can enhance DreamBooth to better personalize Stable Diffusion for text-to-image generation.
Comming soon.
Citation
@article{zhang2023personalize,
title={Personalize Segment Anything Model with One Shot},
author={Zhang, Renrui and Jiang, Zhengkai and Guo, Ziyu and Yan, Shilin and Pan, Junting and Dong, Hao and Gao, Peng and Li, Hongsheng},
journal={arXiv preprint arXiv:2305.03048},
year={2023}
}
Acknowledgement
This repo benefits from Segment Anything and DreamBooth. Thanks for their wonderful works.
Contact
If you have any question about this project, please feel free to contact zhangrenrui@pjlab.org.cn.