Spaces:
Running
Running
Update README.md
Browse files
README.md
CHANGED
|
@@ -1,267 +1,13 @@
|
|
| 1 |
-
# <p align=center> [ECCV 2024] OneRestore: A Universal Restoration Framework for Composite Degradation</p>
|
| 2 |
-
|
| 3 |
-
<div align="center">
|
| 4 |
-
|
| 5 |
-
[](https://arxiv.org/abs/2407.04621)
|
| 6 |
-
[](https://arxiv.org/abs/2407.04621)
|
| 7 |
-
[](https://gy65896.github.io/projects/ECCV2024_OneRestore/index.html)
|
| 8 |
-
[](https://github.com/gy65896/OneRestore/blob/main/poster/OneRestore_poster.png)
|
| 9 |
-
[](https://www.youtube.com/watch?v=AFr5tZdPlZ4)
|
| 10 |
-
|
| 11 |
-
[](https://hits.seeyoufarm.com)
|
| 12 |
-
[](https://github.com/gy65896/OneRestore/issues?q=is%3Aissue+is%3Aclosed)
|
| 13 |
-
[](https://github.com/gy65896/OneRestore/issues)
|
| 14 |
-
|
| 15 |
-
</div>
|
| 16 |
-
<div align=center>
|
| 17 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/7e037f8e-8a8d-4953-8aa6-5142e64f2005" width="720">
|
| 18 |
-
</div>
|
| 19 |
-
|
| 20 |
---
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 26 |
---
|
| 27 |
|
| 28 |
-
|
| 29 |
-
* **2024.09.04**: Code for data synthesis is released.
|
| 30 |
-
* **2024.07.27**: Code for multiple GPUs training is released.
|
| 31 |
-
* **2024.07.20**: [New Website](https://gy65896.github.io/projects/ECCV2024_OneRestore) has been created.
|
| 32 |
-
* **2024.07.10**: [Paper](https://arxiv.org/abs/2407.04621) is released on ArXiv.
|
| 33 |
-
* **2024.07.07**: Code and Dataset are released.
|
| 34 |
-
* **2024.07.02**: OneRestore is accepted by ECCV2024.
|
| 35 |
-
|
| 36 |
-
## Network Architecture
|
| 37 |
-
|
| 38 |
-
</div>
|
| 39 |
-
<div align=center>
|
| 40 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/e26fcaae-3688-489f-8bb4-a698bae3e7fb" width="1080">
|
| 41 |
-
</div>
|
| 42 |
-
|
| 43 |
-
## Quick Start
|
| 44 |
-
|
| 45 |
-
### Install
|
| 46 |
-
|
| 47 |
-
- python 3.7
|
| 48 |
-
- cuda 11.7
|
| 49 |
-
|
| 50 |
-
```
|
| 51 |
-
# git clone this repository
|
| 52 |
-
git clone https://github.com/gy65896/OneRestore.git
|
| 53 |
-
cd OneRestore
|
| 54 |
-
|
| 55 |
-
# create new anaconda env
|
| 56 |
-
conda create -n onerestore python=3.7
|
| 57 |
-
conda activate onerestore
|
| 58 |
-
|
| 59 |
-
# download ckpts
|
| 60 |
-
put embedder_model.tar and onerestore_cdd-11.tar in ckpts folder
|
| 61 |
-
|
| 62 |
-
# install pytorch (Take cuda 11.7 as an example to install torch 1.13)
|
| 63 |
-
pip install torch==1.13.0+cu117 torchvision==0.14.0+cu117 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
|
| 64 |
-
|
| 65 |
-
# install other packages
|
| 66 |
-
pip install -r requirements.txt
|
| 67 |
-
pip install genism
|
| 68 |
-
```
|
| 69 |
-
|
| 70 |
-
### Pretrained Models
|
| 71 |
-
|
| 72 |
-
Please download our pre-trained models and put them in `./ckpts`.
|
| 73 |
-
|
| 74 |
-
| Model | Description
|
| 75 |
-
| :--- | :----------
|
| 76 |
-
|[embedder_model.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpnhSQFIoD9msXWOA?e=aUpHOT) | Text/Visual Embedder trained on our CDD-11.
|
| 77 |
-
|[onerestore_cdd-11.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpmWkGBku6oj33efg?e=7yUGfN) | OneRestore trained on our CDD-11.
|
| 78 |
-
|[onerestore_real.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpi-iJOyN6OSYqiaA?e=QFfMeL) | OneRestore trained on our CDD-11 for Real Scenes.
|
| 79 |
-
|[onerestore_lol.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpkSoVB1j-wYHFpHg?e=0gR9pn) | OneRestore trained on LOL (low light enhancement benchmark).
|
| 80 |
-
|[onerestore_reside_ots.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpjGh8KjfM_QIJzEw?e=zabGTw) | OneRestore trained on RESIDE-OTS (image dehazing benchmark).
|
| 81 |
-
|[onerestore_rain1200.tar](https://1drv.ms/u/s!As3rCDROnrbLgqplAFHv6B348jarGA?e=GuduMT) | OneRestore trained on Rain1200 (image deraining benchmark).
|
| 82 |
-
|[onerestore_snow100k.tar](https://1drv.ms/u/s!As3rCDROnrbLgqphsWWxLZN_7JFJDQ?e=pqezzo) | OneRestore trained on Snow100k-L (image desnowing benchmark).
|
| 83 |
-
|
| 84 |
-
### Inference
|
| 85 |
-
|
| 86 |
-
We provide two samples in `./image` for the quick inference:
|
| 87 |
-
|
| 88 |
-
```
|
| 89 |
-
python test.py --embedder-model-path ./ckpts/embedder_model.tar --restore-model-path ./ckpts/onerestore_cdd-11.tar --input ./image/ --output ./output/ --concat
|
| 90 |
-
```
|
| 91 |
-
|
| 92 |
-
You can also input the prompt to perform controllable restoration. For example:
|
| 93 |
-
|
| 94 |
-
```
|
| 95 |
-
python test.py --embedder-model-path ./ckpts/embedder_model.tar --restore-model-path ./ckpts/onerestore_cdd-11.tar --prompt low_haze --input ./image/ --output ./output/ --concat
|
| 96 |
-
```
|
| 97 |
-
|
| 98 |
-
## Training
|
| 99 |
-
|
| 100 |
-
### Prepare Dataset
|
| 101 |
-
|
| 102 |
-
We provide the download link of our Composite Degradation Dataset with 11 types of degradation ([CDD-11](https://1drv.ms/f/s!As3rCDROnrbLgqpezG4sao-u9ddDhw?e=A0REHx)).
|
| 103 |
-
|
| 104 |
-
Preparing the train and test datasets as follows:
|
| 105 |
-
|
| 106 |
-
```
|
| 107 |
-
./data/
|
| 108 |
-
|--train
|
| 109 |
-
| |--clear
|
| 110 |
-
| | |--000001.png
|
| 111 |
-
| | |--000002.png
|
| 112 |
-
| |--low
|
| 113 |
-
| |--haze
|
| 114 |
-
| |--rain
|
| 115 |
-
| |--snow
|
| 116 |
-
| |--low_haze
|
| 117 |
-
| |--low_rain
|
| 118 |
-
| |--low_snow
|
| 119 |
-
| |--haze_rain
|
| 120 |
-
| |--haze_snow
|
| 121 |
-
| |--low_haze_rain
|
| 122 |
-
| |--low_haze_snow
|
| 123 |
-
|--test
|
| 124 |
-
```
|
| 125 |
-
### Train Model
|
| 126 |
-
|
| 127 |
-
**1. Train Text/Visual Embedder by**
|
| 128 |
-
|
| 129 |
-
```
|
| 130 |
-
python train_Embedder.py --train-dir ./data/CDD-11_train --test-dir ./data/CDD-11_test --check-dir ./ckpts --batch 256 --num-workers 0 --epoch 200 --lr 1e-4 --lr-decay 50
|
| 131 |
-
```
|
| 132 |
-
|
| 133 |
-
**2. Remove the optimizer weights in the Embedder model file by**
|
| 134 |
-
|
| 135 |
-
```
|
| 136 |
-
python remove_optim.py --type Embedder --input-file ./ckpts/embedder_model.tar --output-file ./ckpts/embedder_model.tar
|
| 137 |
-
```
|
| 138 |
-
|
| 139 |
-
**3. Generate the `dataset.h5` file for training OneRestore by**
|
| 140 |
-
|
| 141 |
-
```
|
| 142 |
-
python makedataset.py --train-path ./data/CDD-11_train --data-name dataset.h5 --patch-size 256 --stride 200
|
| 143 |
-
```
|
| 144 |
-
|
| 145 |
-
**4. Train OneRestore model by**
|
| 146 |
-
|
| 147 |
-
- **Single GPU**
|
| 148 |
-
|
| 149 |
-
```
|
| 150 |
-
python train_OneRestore_single-gpu.py --embedder-model-path ./ckpts/embedder_model.tar --save-model-path ./ckpts --train-input ./dataset.h5 --test-input ./data/CDD-11_test --output ./result/ --epoch 120 --bs 4 --lr 1e-4 --adjust-lr 30 --num-works 4
|
| 151 |
-
```
|
| 152 |
-
|
| 153 |
-
- **Multiple GPUs**
|
| 154 |
-
|
| 155 |
-
Assuming you train the OneRestore model using 4 GPUs (e.g., 0, 1, 2, and 3), you can use the following command. Note that the number of nproc_per_node should equal the number of GPUs.
|
| 156 |
-
|
| 157 |
-
```
|
| 158 |
-
CUDA_VISIBLE_DEVICES=0, 1, 2, 3 torchrun --nproc_per_node=4 train_OneRestore_multi-gpu.py --embedder-model-path ./ckpts/embedder_model.tar --save-model-path ./ckpts --train-input ./dataset.h5 --test-input ./data/CDD-11_test --output ./result/ --epoch 120 --bs 4 --lr 1e-4 --adjust-lr 30 --num-works 4
|
| 159 |
-
```
|
| 160 |
-
|
| 161 |
-
**5. Remove the optimizer weights in the OneRestore model file by**
|
| 162 |
-
|
| 163 |
-
```
|
| 164 |
-
python remove_optim.py --type OneRestore --input-file ./ckpts/onerestore_model.tar --output-file ./ckpts/onerestore_model.tar
|
| 165 |
-
```
|
| 166 |
-
|
| 167 |
-
### Customize your own composite degradation dataset
|
| 168 |
-
|
| 169 |
-
**1. Prepare raw data**
|
| 170 |
-
|
| 171 |
-
- Collect your own clear images.
|
| 172 |
-
- Generate the depth map based on [MegaDepth](https://github.com/zhengqili/MegaDepth).
|
| 173 |
-
- Generate the light map based on [LIME](https://github.com/estija/LIME).
|
| 174 |
-
- Generate the rain mask database based on [RainStreakGen](https://github.com/liruoteng/RainStreakGen?tab=readme-ov-file).
|
| 175 |
-
- Download the snow mask database from [Snow100k](https://sites.google.com/view/yunfuliu/desnownet).
|
| 176 |
-
|
| 177 |
-
A generated example is as follows:
|
| 178 |
-
|
| 179 |
-
| Clear Image | Depth Map | Light Map | Rain Mask | Snow Mask
|
| 180 |
-
| :--- | :---| :---| :--- | :---
|
| 181 |
-
| <img src="https://github.com/user-attachments/assets/243a7018-48db-41af-ad80-1f3e2e67ccdb" width="200"> | <img src="https://github.com/user-attachments/assets/dea93b3a-88b9-41cc-9014-5c864d8755b3" width="200"> | <img src="https://github.com/user-attachments/assets/0860d4fc-eb13-4c50-88bf-739f3e3f9908" width="200"> | <img src="https://github.com/user-attachments/assets/0bb06b82-42ee-4b8c-807a-0f7cc728ec97" width="200"> | <img src="https://github.com/user-attachments/assets/d465721b-bc91-4671-817a-e7253be6c890" width="200">
|
| 182 |
-
|
| 183 |
-
(Note: The rain and snow masks do not require strict alignment with the image.)
|
| 184 |
-
|
| 185 |
-
- Prepare the dataset as follows:
|
| 186 |
-
|
| 187 |
-
```
|
| 188 |
-
./syn_data/
|
| 189 |
-
|--data
|
| 190 |
-
| |--clear
|
| 191 |
-
| | |--000001.png
|
| 192 |
-
| | |--000002.png
|
| 193 |
-
| |--depth_map
|
| 194 |
-
| | |--000001.png
|
| 195 |
-
| | |--000002.png
|
| 196 |
-
| |--light_map
|
| 197 |
-
| | |--000001.png
|
| 198 |
-
| | |--000002.png
|
| 199 |
-
| |--rain_mask
|
| 200 |
-
| | |--aaaaaa.png
|
| 201 |
-
| | |--bbbbbb.png
|
| 202 |
-
| |--snow_mask
|
| 203 |
-
| | |--cccccc.png
|
| 204 |
-
| | |--dddddd.png
|
| 205 |
-
|--out
|
| 206 |
-
```
|
| 207 |
-
|
| 208 |
-
**2. Generate composite degradation images**
|
| 209 |
-
|
| 210 |
-
- low+haze+rain
|
| 211 |
-
|
| 212 |
-
```
|
| 213 |
-
python syn_data.py --hq-file ./data/clear/ --light-file ./data/light_map/ --depth-file ./data/depth_map/ --rain-file ./data/rain_mask/ --snow-file ./data/snow_mask/ --out-file ./out/ --low --haze --rain
|
| 214 |
-
```
|
| 215 |
-
|
| 216 |
-
- low+haze+snow
|
| 217 |
-
|
| 218 |
-
```
|
| 219 |
-
python syn_data.py --hq-file ./data/clear/ --light-file ./data/light_map/ --depth-file ./data/depth_map/ --rain-file ./data/rain_mask/ --snow-file ./data/snow_mask/ --out-file ./out/ --low --haze --snow
|
| 220 |
-
```
|
| 221 |
-
(Note: The degradation types can be customized according to specific needs.)
|
| 222 |
-
|
| 223 |
-
| Clear Image | low+haze+rain | low+haze+snow
|
| 224 |
-
| :--- | :--- | :---
|
| 225 |
-
| <img src="https://github.com/user-attachments/assets/243a7018-48db-41af-ad80-1f3e2e67ccdb" width="200"> | <img src="https://github.com/user-attachments/assets/4e662af4-e731-449b-a5b2-c2e93d34be75" width="200"> | <img src="https://github.com/user-attachments/assets/108f08c6-d2e1-4a50-96fa-93615aaf5cb9" width="200">
|
| 226 |
-
|
| 227 |
-
## Performance
|
| 228 |
-
|
| 229 |
-
### CDD-11
|
| 230 |
-
|
| 231 |
-
</div>
|
| 232 |
-
<div align=center>
|
| 233 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/e8b5d6f6-b78c-43a8-9c21-4e78c166fecf" width="720">
|
| 234 |
-
</div>
|
| 235 |
-
|
| 236 |
-
</div>
|
| 237 |
-
<div align=center>
|
| 238 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/835edc0c-acfb-481c-9116-a23ce1929588" width="1080">
|
| 239 |
-
</div>
|
| 240 |
-
|
| 241 |
-
### Real Scene
|
| 242 |
-
|
| 243 |
-
</div>
|
| 244 |
-
<div align=center>
|
| 245 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/f9a4df1c-ad64-4339-8485-b76f29010bdd" width="1080">
|
| 246 |
-
</div>
|
| 247 |
-
|
| 248 |
-
### Controllability
|
| 249 |
-
|
| 250 |
-
</div>
|
| 251 |
-
<div align=center>
|
| 252 |
-
<img src="https://github.com/gy65896/OneRestore/assets/48637474/ed57114a-43a5-4221-bc3a-9bc7f3ac2dd5" width="410"><img src="https://github.com/gy65896/OneRestore/assets/48637474/fd8684f7-8494-4fba-8919-dc50e6acb26f" width="410">
|
| 253 |
-
</div>
|
| 254 |
-
|
| 255 |
-
|
| 256 |
-
## Citation
|
| 257 |
-
|
| 258 |
-
```
|
| 259 |
-
@inproceedings{guo2024onerestore,
|
| 260 |
-
title={OneRestore: A Universal Restoration Framework for Composite Degradation},
|
| 261 |
-
author={Guo, Yu and Gao, Yuan and Lu, Yuxu and Liu, Ryan Wen and He, Shengfeng},
|
| 262 |
-
booktitle={European Conference on Computer Vision},
|
| 263 |
-
year={2024}
|
| 264 |
-
}
|
| 265 |
-
```
|
| 266 |
-
|
| 267 |
-
#### If you have any questions, please get in touch with me (guoyu65896@gmail.com).
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: OneRestore
|
| 3 |
+
emoji: 🦀
|
| 4 |
+
colorFrom: purple
|
| 5 |
+
colorTo: blue
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 4.38.1
|
| 8 |
+
app_file: app.py
|
| 9 |
+
pinned: false
|
| 10 |
+
license: other
|
| 11 |
---
|
| 12 |
|
| 13 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|