File size: 3,052 Bytes
9cba551 4071595 9cba551 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
---
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> 
<a href='https://github.com/qianyu-dlut/DiffDIS'><img src='https://img.shields.io/badge/Github-DiffDIS-blue'></a> 
<a href='LICENSE'><img src='https://img.shields.io/badge/License-MIT-yellow'></a> 
</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}
}
``` |