/ControlNet4FSS

Official Implementation of DifFSS Generation Pipeline

Primary LanguagePythonApache License 2.0Apache-2.0

DifFSS: Diffusion Model for Few-Shot Semantic Segmentation

This is the implementation of the auxiliary support generator in paper "DifFSS: Diffusion Model for Few-Shot Semantic Segmentation".

For more information, Please refer to the the paper on [arXiv].

🛠️ Getting Started with ControlNet4FSS

1. Setup Environment

    conda env create -f environment.yaml
    conda activate control

2. Download the pre-trained models

All models and detectors can be downloaded from Hugging Face page. Make sure that SD models are put in "ControlNet/models" and detectors are put in "ControlNet/annotator/ckpts". Make sure that you download all necessary pretrained weights and detector models from that Hugging Face page.

For DifFSS, we utilized the seg, scribble, hed model of ControlNet.

3. Preparation of the dataset

Please download the following datasets:

4. Start generating

python FSSpregenerate.py --st 0 --end -1 --imgdir /data/user6/coco/ --maskdir /data/user6/coco/annotations/ --dstdir /data/user6/justtest/ --list ./list/coco_all.txt --dataset coco --guidance seg --save_control 0

Please refer to the source code for the functions of arguments.

🏆 Try DifFSS on FSS models

The DifFSS method can be easily applied to existing FSS models, please refer to one of our implementations on BAM here.