gy65896 commited on
Commit
ec17309
·
verified ·
1 Parent(s): 2940390

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +10 -264
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
- [![ArXiv](https://img.shields.io/badge/OneRestore-ArXiv-red.svg)](https://arxiv.org/abs/2407.04621)
6
- [![Paper](https://img.shields.io/badge/OneRestore-Paper-purple.svg)](https://arxiv.org/abs/2407.04621)
7
- [![Web](https://img.shields.io/badge/OneRestore-Web-blue.svg)](https://gy65896.github.io/projects/ECCV2024_OneRestore/index.html)
8
- [![Poster](https://img.shields.io/badge/OneRestore-Poster-green.svg)](https://github.com/gy65896/OneRestore/blob/main/poster/OneRestore_poster.png)
9
- [![Video](https://img.shields.io/badge/OneRestore-Video-orange.svg)](https://www.youtube.com/watch?v=AFr5tZdPlZ4)
10
-
11
- [![Hits](https://hits.seeyoufarm.com/api/count/incr/badge.svg?url=https%3A%2F%2Fgithub.com%2Fgy65896%2FOneRestore&count_bg=%2379C83D&title_bg=%23555555&icon=&icon_color=%23E7E7E7&title=hits&edge_flat=false)](https://hits.seeyoufarm.com)
12
- [![Closed Issues](https://img.shields.io/github/issues-closed/gy65896/OneRestore)](https://github.com/gy65896/OneRestore/issues?q=is%3Aissue+is%3Aclosed)
13
- [![Open Issues](https://img.shields.io/github/issues/gy65896/OneRestore)](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
- >**OneRestore: A Universal Restoration Framework for Composite Degradation**<br> [Yu Guo](https://scholar.google.com/citations?user=klYz-acAAAAJ&hl=zh-CN)<sup>† </sup>, [Yuan Gao](https://scholar.google.com.hk/citations?hl=zh-CN&user=4JpRnU4AAAAJ&view_op=list_works&sortby=pubdate)<sup>† </sup>, [Yuxu Lu](https://scholar.google.com.hk/citations?user=XXge2_0AAAAJ&hl=zh-CN), [Huilin Zhu](https://scholar.google.com.hk/citations?hl=zh-CN&user=fluPrxcAAAAJ), [Ryan Wen Liu](http://mipc.whut.edu.cn/index.html)<sup>* </sup>, [Shengfeng He](http://www.shengfenghe.com/)<sup>* </sup> <br>
22
- († Co-first Author, * Corresponding Author)<br>
23
- >European Conference on Computer Vision
24
-
25
- > **Abstract:** *In real-world scenarios, image impairments often manifest as composite degradations, presenting a complex interplay of elements such as low light, haze, rain, and snow. Despite this reality, existing restoration methods typically target isolated degradation types, thereby falling short in environments where multiple degrading factors coexist. To bridge this gap, our study proposes a versatile imaging model that consolidates four physical corruption paradigms to accurately represent complex, composite degradation scenarios. In this context, we propose OneRestore, a novel transformer-based framework designed for adaptive, controllable scene restoration. The proposed framework leverages a unique cross-attention mechanism, merging degraded scene descriptors with image features, allowing for nuanced restoration. Our model allows versatile input scene descriptors, ranging from manual text embeddings to automatic extractions based on visual attributes. Our methodology is further enhanced through a composite degradation restoration loss, using extra degraded images as negative samples to fortify model constraints. Comparative results on synthetic and real-world datasets demonstrate OneRestore as a superior solution, significantly advancing the state-of-the-art in addressing complex, composite degradations.*
 
 
 
 
26
  ---
27
 
28
- ## News 🚀
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