README.md

---
license: mit
tags:
- image-segmentation
---

<h1 align="center">High-Precision Dichotomous Image Segmentation via Probing Diffusion Capacity</h1>


<div align="center" style="display: flex; justify-content: center; flex-wrap: wrap;">
  <a href='https://arxiv.org/pdf/2410.10105'><img src='https://img.shields.io/badge/arXiv-DiffDIS-B31B1B'></a>&ensp; 
  <a href='https://github.com/qianyu-dlut/DiffDIS'><img src='https://img.shields.io/badge/Github-DiffDIS-blue'></a>&ensp; 
  <a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a>&ensp; 
</div>


This repository contains the official implementation for the paper "[High-Precision Dichotomous Image Segmentation via Probing Diffusion Capacity](https://arxiv.org/pdf/2410.10105)" (ICLR 2025).

<p align="center">
    <img alt="DiffDIS teaser image" src="https://raw.githubusercontent.com/qianyu-dlut/DiffDIS/master/assets/image.png" width="900px">
</p>

## How to use 
>  For the complete training and inference process, please refer to our [GitHub Repository](https://github.com/qianyu-dlut/DiffDIS). This section specifically guides you on loading weights from Hugging Face.

### Install Packages:
```shell
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt
pip install -e diffusers-0.30.2/
```

### Load DiffDIS weights from Hugging Face:

```python
import torch
from diffusers import (
    DiffusionPipeline,
    DDPMScheduler,
    UNet2DConditionModel,
    AutoencoderKL,
)
from transformers import CLIPTextModel, CLIPTokenizer

hf_model_path = 'qianyu1217/diffdis'
vae = AutoencoderKL.from_pretrained(hf_model_path,subfolder='vae',trust_remote_code=True)
scheduler = DDPMScheduler.from_pretrained(hf_model_path,subfolder='scheduler')
text_encoder = CLIPTextModel.from_pretrained(hf_model_path,subfolder='text_encoder')
tokenizer = CLIPTokenizer.from_pretrained(hf_model_path,subfolder='tokenizer')
unet = UNet2DConditionModel_diffdis.from_pretrained(hf_model_path,subfolder="unet",
                                in_channels=8, sample_size=96,
                                low_cpu_mem_usage=False,
                                ignore_mismatched_sizes=False,
                                class_embed_type='projection',
                                projection_class_embeddings_input_dim=4,
                                mid_extra_cross=True,
                                mode = 'DBIA',
                                use_swci = True)
pipe = DiffDISPipeline(unet=unet,
                        vae=vae,
                        scheduler=scheduler,
                        text_encoder=text_encoder,
                        tokenizer=tokenizer)

```



## Citation

```
@article{DiffDIS,
  title={High-Precision Dichotomous Image Segmentation via Probing Diffusion Capacity},
  author={Yu, Qian and Jiang, Peng-Tao and Zhang, Hao and Chen, Jinwei and Li, Bo and Zhang, Lihe and Lu, Huchuan},
  journal={arXiv preprint arXiv:2410.10105},
  year={2024}
}
```