Upload 36 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+image/low_haze_rain_00469_01_lq.png filter=lfs diff=lfs merge=lfs -text
+image/low_haze_snow_00337_01_lq.png filter=lfs diff=lfs merge=lfs -text
+img_file/OneRestore_poster.png filter=lfs diff=lfs merge=lfs -text
+output/low_haze_rain_00469_01_lq.png filter=lfs diff=lfs merge=lfs -text
+output/low_haze_snow_00337_01_lq.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,267 @@
+ # <p align=center> [ECCV 2024] OneRestore: A Universal Restoration Framework for Composite Degradation</p>
+
+<div align="center">
+ 
+[![ArXiv](https://img.shields.io/badge/OneRestore-ArXiv-red.svg)](https://arxiv.org/abs/2407.04621)
+[![Paper](https://img.shields.io/badge/OneRestore-Paper-purple.svg)](https://arxiv.org/abs/2407.04621)
+[![Web](https://img.shields.io/badge/OneRestore-Web-blue.svg)](https://gy65896.github.io/projects/ECCV2024_OneRestore/index.html)
+[![Poster](https://img.shields.io/badge/OneRestore-Poster-green.svg)](https://github.com/gy65896/OneRestore/blob/main/poster/OneRestore_poster.png)
+[![Video](https://img.shields.io/badge/OneRestore-Video-orange.svg)](https://www.youtube.com/watch?v=AFr5tZdPlZ4)
+
+[![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)
+[![Closed Issues](https://img.shields.io/github/issues-closed/gy65896/OneRestore)](https://github.com/gy65896/OneRestore/issues?q=is%3Aissue+is%3Aclosed)
+[![Open Issues](https://img.shields.io/github/issues/gy65896/OneRestore)](https://github.com/gy65896/OneRestore/issues)
+
+</div>
+<div align=center>
+<img src="https://github.com/gy65896/OneRestore/assets/48637474/7e037f8e-8a8d-4953-8aa6-5142e64f2005" width="720">
+</div>
+
 ---
-title: OneRestore
-emoji: 👀
-colorFrom: purple
-colorTo: yellow
-sdk: gradio
-sdk_version: 4.43.0
-app_file: app.py
-pinned: false
+>**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>
+(† Co-first Author, * Corresponding Author)<br>
+>European Conference on Computer Vision
+
+> **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.*
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## News 🚀
+* **2024.09.04**: Code for data synthesis is released.
+* **2024.07.27**: Code for multiple GPUs training is released.
+* **2024.07.20**: [New Website](https://gy65896.github.io/projects/ECCV2024_OneRestore) has been created.
+* **2024.07.10**: [Paper](https://arxiv.org/abs/2407.04621) is released on ArXiv.
+* **2024.07.07**: Code and Dataset are released.
+* **2024.07.02**: OneRestore is accepted by ECCV2024.
+
+## Network Architecture
+
+</div>
+<div align=center>
+<img src="https://github.com/gy65896/OneRestore/assets/48637474/e26fcaae-3688-489f-8bb4-a698bae3e7fb" width="1080">
+</div>
+
+## Quick Start
+
+### Install
+
+- python 3.7
+- cuda 11.7
+
+```
+# git clone this repository
+git clone https://github.com/gy65896/OneRestore.git
+cd OneRestore
+
+# create new anaconda env
+conda create -n onerestore python=3.7
+conda activate onerestore 
+
+# download ckpts
+put embedder_model.tar and onerestore_cdd-11.tar in ckpts folder
+
+# install pytorch (Take cuda 11.7 as an example to install torch 1.13)
+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
+
+# install other packages
+pip install -r requirements.txt
+pip install genism
+```
+
+### Pretrained Models
+
+Please download our pre-trained models and put them in  `./ckpts`.
+
+| Model | Description
+| :--- | :----------
+|[embedder_model.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpnhSQFIoD9msXWOA?e=aUpHOT)  | Text/Visual Embedder trained on our CDD-11.
+|[onerestore_cdd-11.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpmWkGBku6oj33efg?e=7yUGfN)  | OneRestore trained on our CDD-11.
+|[onerestore_real.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpi-iJOyN6OSYqiaA?e=QFfMeL)  | OneRestore trained on our CDD-11 for Real Scenes.
+|[onerestore_lol.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpkSoVB1j-wYHFpHg?e=0gR9pn)  | OneRestore trained on LOL (low light enhancement benchmark).
+|[onerestore_reside_ots.tar](https://1drv.ms/u/s!As3rCDROnrbLgqpjGh8KjfM_QIJzEw?e=zabGTw)  | OneRestore trained on RESIDE-OTS (image dehazing benchmark).
+|[onerestore_rain1200.tar](https://1drv.ms/u/s!As3rCDROnrbLgqplAFHv6B348jarGA?e=GuduMT)  | OneRestore trained on Rain1200 (image deraining benchmark).
+|[onerestore_snow100k.tar](https://1drv.ms/u/s!As3rCDROnrbLgqphsWWxLZN_7JFJDQ?e=pqezzo)  | OneRestore trained on Snow100k-L (image desnowing benchmark).
+
+### Inference
+
+We provide two samples in `./image` for the quick inference:
+
+```
+python test.py --embedder-model-path ./ckpts/embedder_model.tar --restore-model-path ./ckpts/onerestore_cdd-11.tar --input ./image/ --output ./output/ --concat
+```
+
+You can also input the prompt to perform controllable restoration. For example:
+
+```
+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
+```
+
+## Training
+
+### Prepare Dataset
+
+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)).
+
+Preparing the train and test datasets as follows:
+
+```
+./data/
+|--train
+|  |--clear
+|  |  |--000001.png
+|  |  |--000002.png
+|  |--low
+|  |--haze
+|  |--rain
+|  |--snow
+|  |--low_haze
+|  |--low_rain
+|  |--low_snow
+|  |--haze_rain
+|  |--haze_snow
+|  |--low_haze_rain
+|  |--low_haze_snow
+|--test
+```
+### Train Model
+
+**1. Train Text/Visual Embedder by**
+
+```
+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
+```
+
+**2. Remove the optimizer weights in the Embedder model file by**
+
+```
+python remove_optim.py --type Embedder --input-file ./ckpts/embedder_model.tar --output-file ./ckpts/embedder_model.tar
+```
+
+**3. Generate the `dataset.h5` file for training OneRestore by**
+
+```
+python makedataset.py --train-path ./data/CDD-11_train --data-name dataset.h5 --patch-size 256 --stride 200
+```
+
+**4. Train OneRestore model by**
+
+- **Single GPU**
+
+```
+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
+```
+
+- **Multiple GPUs**
+
+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.
+
+```
+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
+```
+
+**5. Remove the optimizer weights in the OneRestore model file by**
+
+```
+python remove_optim.py --type OneRestore --input-file ./ckpts/onerestore_model.tar --output-file ./ckpts/onerestore_model.tar
+```
+
+### Customize your own composite degradation dataset
+
+**1. Prepare raw data**
+
+ - Collect your own clear images.
+ - Generate the depth map based on [MegaDepth](https://github.com/zhengqili/MegaDepth).
+ - Generate the light map based on [LIME](https://github.com/estija/LIME).
+ - Generate the rain mask database based on [RainStreakGen](https://github.com/liruoteng/RainStreakGen?tab=readme-ov-file).
+ - Download the snow mask database from [Snow100k](https://sites.google.com/view/yunfuliu/desnownet).
+
+A generated example is as follows:
+
+| Clear Image | Depth Map | Light Map | Rain Mask | Snow Mask
+| :--- | :---| :---| :--- | :---
+| <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">
+
+(Note: The rain and snow masks do not require strict alignment with the image.)
+
+ - Prepare the dataset as follows:
+
+```
+./syn_data/
+|--data
+|  |--clear
+|  |  |--000001.png
+|  |  |--000002.png
+|  |--depth_map
+|  |  |--000001.png
+|  |  |--000002.png
+|  |--light_map
+|  |  |--000001.png
+|  |  |--000002.png
+|  |--rain_mask
+|  |  |--aaaaaa.png
+|  |  |--bbbbbb.png
+|  |--snow_mask
+|  |  |--cccccc.png
+|  |  |--dddddd.png
+|--out
+```
+
+**2. Generate composite degradation images**
+
+ - low+haze+rain
+
+```
+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
+```
+
+ - low+haze+snow
+
+```
+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
+```
+(Note: The degradation types can be customized according to specific needs.)
+
+| Clear Image | low+haze+rain | low+haze+snow
+| :--- | :--- | :---
+| <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">
+
+## Performance
+
+### CDD-11
+
+</div>
+<div align=center>
+<img src="https://github.com/gy65896/OneRestore/assets/48637474/e8b5d6f6-b78c-43a8-9c21-4e78c166fecf" width="720">
+</div>
+
+</div>
+<div align=center>
+<img src="https://github.com/gy65896/OneRestore/assets/48637474/835edc0c-acfb-481c-9116-a23ce1929588" width="1080">
+</div>
+
+### Real Scene
+
+</div>
+<div align=center>
+<img src="https://github.com/gy65896/OneRestore/assets/48637474/f9a4df1c-ad64-4339-8485-b76f29010bdd" width="1080">
+</div>
+
+### Controllability
+
+</div>
+<div align=center>
+<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">
+</div>
+
+
+## Citation
+
+```
+@inproceedings{guo2024onerestore,
+  title={OneRestore: A Universal Restoration Framework for Composite Degradation},
+  author={Guo, Yu and Gao, Yuan and Lu, Yuxu and Liu, Ryan Wen and He, Shengfeng},
+  booktitle={European Conference on Computer Vision},
+  year={2024}
+}
+```
+
+#### If you have any questions, please get in touch with me (guoyu65896@gmail.com).

app.py

@@ -0,0 +1,85 @@
+
+import torch
+import gradio as gr
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+from utils.utils import load_restore_ckpt, load_embedder_ckpt
+import os
+
+# Enforce CPU usage
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+embedder_model_path = "ckpts/embedder_model.tar"  # Update with actual path to embedder checkpoint
+restorer_model_path = "ckpts/onerestore_cdd-11.tar"   # Update with actual path to restorer checkpoint
+
+# Load models on CPU only
+embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=embedder_model_path)
+restorer = load_restore_ckpt(device, freeze_model=True, ckpt_name=restorer_model_path)
+
+# Define image preprocessing and postprocessing
+transform_resize = transforms.Compose([
+        transforms.Resize([224,224]),
+        transforms.ToTensor()
+        ]) 
+
+
+def postprocess_image(tensor):
+    image = tensor.squeeze(0).cpu().detach().numpy()
+    image = (image) * 255  # Assuming output in [-1, 1], rescale to [0, 255]
+    image = np.clip(image, 0, 255).astype("uint8")  # Clip values to [0, 255]
+    return Image.fromarray(image.transpose(1, 2, 0))  # Reorder to (H, W, C)
+
+# Define the enhancement function
+def enhance_image(image, degradation_type=None):
+    # Preprocess the image
+    input_tensor = torch.Tensor((np.array(image)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+    lq_em = transform_resize(image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+    lq_em = transform_resize(image).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Generate embedding
+    if degradation_type == "auto" or degradation_type is None:
+        text_embedding, _, [text] = embedder(lq_em, 'image_encoder')
+    else:
+        text_embedding, _, [text] = embedder([degradation_type], 'text_encoder')
+    
+    # Model inference
+    with torch.no_grad():
+        enhanced_tensor = restorer(input_tensor, text_embedding)
+    
+    # Postprocess the output
+    return postprocess_image(enhanced_tensor), text
+
+# Define the Gradio interface
+def inference(image, degradation_type=None):
+    return enhance_image(image, degradation_type)
+
+#### Image,Prompts examples
+examples = [
+            ['image/low_haze_rain_00469_01_lq.png'],
+            ['image/low_haze_snow_00337_01_lq.png'],
+            ]
+
+
+
+# Create the Gradio app interface using updated API
+interface = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.Image(type="pil", value="image/low_haze_rain_00469_01_lq.png"),  # Image input
+        gr.Dropdown(['auto', 'low', 'haze', 'rain', 'snow',\
+                                            'low_haze', 'low_rain', 'low_snow', 'haze_rain',\
+                                                    'haze_snow', 'low_haze_rain', 'low_haze_snow'], label="Degradation Type", value="auto")  # Manual or auto degradation
+    ],
+    outputs=[
+        gr.Image(type="pil"),  # Enhanced image output
+        gr.Textbox(label="Degradation Type")  # Display the estimated degradation type
+    ],
+    title="Image Restoration with OneRestore",
+    description="Upload an image and enhance it using OneRestore model. You can choose to let the model automatically estimate the degradation type or set it manually.",
+    examples=examples,
+)
+
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()

ckpts/embedder_model.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01c1ed1fe4fd06a73c78a7d1fac1b4092a098f3ea08bc71bcab1e4fdad7f6ec2
+size 48638120

ckpts/onerestore_cdd-11.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e02e2d87ce56740a9bedb5cffe8b129d3b206590ae3b97050c9b53aa549c7c9f
+size 23993607

makedataset.py

@@ -0,0 +1,157 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Feb 12 20:00:46 2020
+
+@author: Administrator
+"""
+
+import os
+import os.path
+import random
+import numpy as np
+import cv2
+import h5py
+import torch
+import torch.utils.data as udata
+import argparse
+from PIL import Image
+class Dataset(udata.Dataset):
+	r"""Implements torch.utils.data.Dataset
+	"""
+	def __init__(self, file, trainrgb=True,trainsyn = True, shuffle=False):
+		super(Dataset, self).__init__()
+		self.trainrgb = trainrgb
+		self.trainsyn = trainsyn
+		self.train_haze	 = file
+		
+		h5f = h5py.File(self.train_haze, 'r')
+		  
+		self.keys = list(h5f.keys())
+		if shuffle:
+			random.shuffle(self.keys)
+		h5f.close()
+
+	def __len__(self):
+		return len(self.keys)
+
+	def __getitem__(self, index):
+
+		h5f = h5py.File(self.train_haze, 'r')
+		  
+		key = self.keys[index]
+		data = np.array(h5f[key])
+		h5f.close()
+		return torch.Tensor(data)
+
+def data_augmentation(clear, mode):
+    r"""Performs dat augmentation of the input image
+
+    Args:
+        image: a cv2 (OpenCV) image
+        mode: int. Choice of transformation to apply to the image
+            0 - no transformation
+            1 - flip up and down
+            2 - rotate counterwise 90 degree
+            3 - rotate 90 degree and flip up and down
+            4 - rotate 180 degree
+            5 - rotate 180 degree and flip
+            6 - rotate 270 degree
+            7 - rotate 270 degree and flip
+    """
+    clear = np.transpose(clear, (2, 3, 0, 1))
+    if mode == 0:
+        # original
+        clear = clear
+    elif mode == 1:
+        # flip up and down
+        clear = np.flipud(clear)
+    elif mode == 2:
+        # rotate counterwise 90 degree
+        clear = np.rot90(clear)
+    elif mode == 3:
+        # rotate 90 degree and flip up and down
+        clear = np.rot90(clear)
+        clear = np.flipud(clear)
+    elif mode == 4:
+        # rotate 180 degree
+        clear = np.rot90(clear, k=2)
+    elif mode == 5:
+        # rotate 180 degree and flip
+        clear = np.rot90(clear, k=2)
+        clear = np.flipud(clear)
+    elif mode == 6:
+        # rotate 270 degree
+        clear = np.rot90(clear, k=3)
+    elif mode == 7:
+        # rotate 270 degree and flip
+        clear = np.rot90(clear, k=3)
+        clear = np.flipud(clear)
+    else:
+        raise Exception('Invalid choice of image transformation')
+    return np.transpose(clear, (2, 3, 0, 1))
+
+def img_to_patches(img,win,stride,Syn=True):
+    typ, chl, raw, col = img.shape
+    chl = int(chl)
+    num_raw = np.ceil((raw-win)/stride+1).astype(np.uint8)
+    num_col = np.ceil((col-win)/stride+1).astype(np.uint8)
+    count = 0
+    total_process = int(num_col)*int(num_raw)
+    img_patches = np.zeros([typ, chl, win, win, total_process])
+    if Syn:
+        for i in range(num_raw):
+            for j in range(num_col):
+                if stride * i + win <= raw and stride * j + win <=col:
+                    img_patches[:,:,:,:,count] = img[:, :, stride*i : stride*i + win, stride*j : stride*j + win]
+                elif stride * i + win > raw and stride * j + win<=col:
+                    img_patches[:,:,:,:,count] = img[:, :,raw-win : raw,stride * j : stride * j + win]
+                elif stride * i + win <= raw and stride*j + win>col:
+                    img_patches[:,:,:,:,count] = img[:, :,stride*i : stride*i + win, col-win : col]
+                else:
+                    img_patches[:,:,:,:,count] = img[:, :,raw-win : raw,col-win : col]
+                img_patches[:,:,:,:,count] = data_augmentation(img_patches[:, :, :, :, count], np.random.randint(0, 7))
+                count +=1
+    return img_patches
+
+def read_img(img):
+    return np.array(Image.open(img))/255.
+
+def Train_data(args):
+    file_list = os.listdir(f'{args.train_path}/{args.gt_name}')
+
+    with h5py.File(args.data_name, 'w') as h5f:
+        count = 0
+        for i in range(len(file_list)):
+            print(file_list[i])
+            img_list = []
+            
+            img_list.append(read_img(f'{args.train_path}/{args.gt_name}/{file_list[i]}'))
+            for j in args.degradation_name:
+                img_list.append(read_img(f'{args.train_path}/{j}/{file_list[i]}'))
+            
+            img = np.stack(img_list,0)
+            img = img_to_patches(img.transpose(0, 3, 1, 2), args.patch_size, args.stride)
+            
+            for nx in range(img.shape[4]):
+                data = img[:,:,:,:,nx]
+                print(count, data.shape)
+                h5f.create_dataset(str(count), data=data)
+                count += 1
+        h5f.close()
+
+if __name__ == "__main__":
+        
+    parser = argparse.ArgumentParser(description = "Building the training patch database")
+    parser.add_argument("--patch-size", type = int, default=256, help="Patch size")
+    parser.add_argument("--stride", type = int, default=200, help="Size of stride")
+
+    parser.add_argument("--train-path", type = str, default='./data/CDD-11_train', help="Train path")
+    parser.add_argument("--data-name", type = str, default='dataset.h5', help="Data name")
+
+    parser.add_argument("--gt-name", type = str, default='clear', help="HQ name")
+    parser.add_argument("--degradation-name", type = list, default=['low','haze','rain','snow',\
+    'low_haze','low_rain','low_snow','haze_rain','haze_snow','low_haze_rain','low_haze_snow'], help="LQ name")
+
+    args = parser.parse_args()
+    
+    Train_data(args)

model/Embedder.py

@@ -0,0 +1,238 @@
+import numpy as np
+import torch, torchvision
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from utils.utils_word_embedding import initialize_wordembedding_matrix
+
+class Backbone(nn.Module):
+    def __init__(self, backbone='resnet18'):
+        super(Backbone, self).__init__()
+
+        if backbone == 'resnet18':
+            resnet = torchvision.models.resnet.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
+        elif backbone == 'resnet50':
+            resnet = torchvision.models.resnet.resnet50(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
+        elif backbone == 'resnet101':
+            resnet = torchvision.models.resnet.resnet101(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
+
+        self.block0 = nn.Sequential(
+            resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
+        )
+        self.block1 = resnet.layer1
+        self.block2 = resnet.layer2
+        self.block3 = resnet.layer3
+        self.block4 = resnet.layer4
+
+    def forward(self, x, returned=[4]):
+        blocks = [self.block0(x)]
+
+        blocks.append(self.block1(blocks[-1]))
+        blocks.append(self.block2(blocks[-1]))
+        blocks.append(self.block3(blocks[-1]))
+        blocks.append(self.block4(blocks[-1]))
+
+        out = [blocks[i] for i in returned]
+        return out
+
+class CosineClassifier(nn.Module):
+    def __init__(self, temp=0.05):
+        super(CosineClassifier, self).__init__()
+        self.temp = temp
+
+    def forward(self, img, concept, scale=True):
+        """
+        img: (bs, emb_dim)
+        concept: (n_class, emb_dim)
+        """
+        img_norm = F.normalize(img, dim=-1)
+        concept_norm = F.normalize(concept, dim=-1)
+        pred = torch.matmul(img_norm, concept_norm.transpose(0, 1))
+        if scale:
+            pred = pred / self.temp
+        return pred
+
+class Embedder(nn.Module):
+    """
+    Text and Visual Embedding Model.
+    """
+    def __init__(self,
+                 type_name,
+                 feat_dim = 512,
+                 mid_dim = 1024,
+                 out_dim = 324,
+                 drop_rate = 0.35,
+                 cosine_cls_temp = 0.05,
+                 wordembs = 'glove',
+                 extractor_name = 'resnet18'):
+        super(Embedder, self).__init__()
+
+        mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+        self.type_name = type_name
+        self.feat_dim = feat_dim
+        self.mid_dim = mid_dim
+        self.out_dim = out_dim
+        self.drop_rate = drop_rate
+        self.cosine_cls_temp = cosine_cls_temp
+        self.wordembs = wordembs
+        self.extractor_name = extractor_name
+        self.transform = transforms.Normalize(mean, std)
+        
+        self._setup_word_embedding()
+        self._setup_image_embedding()
+        
+    def _setup_image_embedding(self):
+        # image embedding
+        self.feat_extractor = Backbone(self.extractor_name)
+
+        img_emb_modules = [
+            nn.Conv2d(self.feat_dim, self.mid_dim, kernel_size=1, bias=False),
+            nn.BatchNorm2d(self.mid_dim),
+            nn.ReLU()
+        ]
+        if self.drop_rate > 0:
+            img_emb_modules += [nn.Dropout2d(self.drop_rate)]
+        self.img_embedder = nn.Sequential(*img_emb_modules)
+
+        self.img_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.img_final = nn.Linear(self.mid_dim, self.out_dim)
+
+        self.classifier = CosineClassifier(temp=self.cosine_cls_temp)
+
+    def _setup_word_embedding(self):
+
+        self.type2idx = {self.type_name[i]: i for i in range(len(self.type_name))}
+        self.num_type = len(self.type_name)
+        train_type = [self.type2idx[type_i] for type_i in self.type_name]
+        self.train_type = torch.LongTensor(train_type).to("cuda" if torch.cuda.is_available() else "cpu")
+
+        wordemb, self.word_dim = \
+            initialize_wordembedding_matrix(self.wordembs, self.type_name)
+        
+        self.embedder = nn.Embedding(self.num_type, self.word_dim)
+        self.embedder.weight.data.copy_(wordemb)
+
+        self.mlp = nn.Sequential(
+                nn.Linear(self.word_dim, self.out_dim),
+                nn.ReLU(True)
+            )
+
+    def train_forward(self, batch):
+
+        scene, img = batch[0], self.transform(batch[1])
+        bs = img.shape[0]
+
+        # word embedding
+        scene_emb = self.embedder(self.train_type)
+        scene_weight = self.mlp(scene_emb)
+
+        #image embedding
+        img = self.feat_extractor(img)[0]
+        img = self.img_embedder(img)
+        img = self.img_avg_pool(img).squeeze(3).squeeze(2)
+        img = self.img_final(img)
+
+        pred = self.classifier(img, scene_weight)
+        label_loss = F.cross_entropy(pred, scene)
+        pred = torch.max(pred, dim=1)[1]
+        type_pred = self.train_type[pred]
+        correct_type = (type_pred == scene)
+        out = {
+            'loss_total': label_loss,
+            'acc_type': torch.div(correct_type.sum(),float(bs)), 
+        }
+
+        return out
+    
+    def image_encoder_forward(self, batch):
+        img = self.transform(batch)
+
+        # word embedding
+        scene_emb = self.embedder(self.train_type)
+        scene_weight = self.mlp(scene_emb)
+
+        #image embedding
+        img = self.feat_extractor(img)[0]
+        bs, _, h, w = img.shape
+        img = self.img_embedder(img)
+        img = self.img_avg_pool(img).squeeze(3).squeeze(2)
+        img = self.img_final(img)
+
+        pred = self.classifier(img, scene_weight)
+        pred = torch.max(pred, dim=1)[1]
+
+        out_embedding = torch.zeros((bs,self.out_dim)).to("cuda" if torch.cuda.is_available() else "cpu")
+        for i in range(bs):
+            out_embedding[i,:] = scene_weight[pred[i],:]
+        num_type = self.train_type[pred]
+        text_type = [self.type_name[num_type[i]] for i in range(bs)]
+
+        return out_embedding, num_type, text_type
+    
+    def text_encoder_forward(self, text):
+
+        bs = len(text)
+
+        # word embedding
+        scene_emb = self.embedder(self.train_type)
+        scene_weight = self.mlp(scene_emb)
+
+        num_type = torch.zeros((bs)).to("cuda" if torch.cuda.is_available() else "cpu")
+        for i in range(bs):
+            num_type[i] = self.type2idx[text[i]]
+
+        out_embedding = torch.zeros((bs,self.out_dim)).to("cuda" if torch.cuda.is_available() else "cpu")
+        for i in range(bs):
+            out_embedding[i,:] = scene_weight[int(num_type[i]),:]
+        text_type = text
+
+        return out_embedding, num_type, text_type
+    
+    def text_idx_encoder_forward(self, idx):
+
+        bs = idx.shape[0]
+
+        # word embedding
+        scene_emb = self.embedder(self.train_type)
+        scene_weight = self.mlp(scene_emb)
+
+        num_type = idx
+
+        out_embedding = torch.zeros((bs,self.out_dim)).to("cuda" if torch.cuda.is_available() else "cpu")
+        for i in range(bs):
+            out_embedding[i,:] = scene_weight[int(num_type[i]),:]
+
+        return out_embedding
+
+    def contrast_loss_forward(self, batch):
+
+        img = self.transform(batch)
+
+        #image embedding
+        img = self.feat_extractor(img)[0]
+        img = self.img_embedder(img)
+        img = self.img_avg_pool(img).squeeze(3).squeeze(2)
+        img = self.img_final(img)
+
+        return img
+    
+    def forward(self, x, type = 'image_encoder'):
+
+        if type == 'train':
+            out = self.train_forward(x)
+
+        elif type == 'image_encoder':
+            with torch.no_grad():
+                out = self.image_encoder_forward(x)
+
+        elif type == 'text_encoder':
+            out = self.text_encoder_forward(x)
+        
+        elif type == 'text_idx_encoder':
+            out = self.text_idx_encoder_forward(x)
+
+        elif type == 'visual_embed':
+            x = F.interpolate(x,size=(224,224),mode='bilinear')
+            out = self.contrast_loss_forward(x)
+
+        return out

model/OneRestore.py

@@ -0,0 +1,314 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Jun 20 16:14:37 2021
+
+@author: Administrator
+"""
+
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from torchvision import transforms
+import torch, math
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+import numbers
+
+from thop import profile
+import numpy as np
+import time
+from torchvision import transforms
+
+
+class OneRestore(nn.Module):
+	def __init__(self, channel = 32):
+		super(OneRestore,self).__init__()
+		self.norm = lambda x: (x-0.5)/0.5
+		self.denorm = lambda x: (x+1)/2
+		self.in_conv = nn.Conv2d(3,channel,kernel_size=1,stride=1,padding=0,bias=False)
+		self.encoder = encoder(channel)
+		self.middle = backbone(channel)
+		self.decoder = decoder(channel)
+		self.out_conv = nn.Conv2d(channel,3,kernel_size=1,stride=1,padding=0,bias=False)
+
+	def forward(self,x,embedding):
+		x_in = self.in_conv(self.norm(x))
+		x_l, x_m, x_s, x_ss = self.encoder(x_in, embedding)
+		x_mid = self.middle(x_ss, embedding)
+		x_out = self.decoder(x_mid, x_ss, x_s, x_m, x_l, embedding)
+		out = self.out_conv(x_out) + x
+		return self.denorm(out)
+
+class encoder(nn.Module):
+	def __init__(self,channel):
+		super(encoder,self).__init__()    
+
+		self.el = ResidualBlock(channel)#16
+		self.em = ResidualBlock(channel*2)#32
+		self.es = ResidualBlock(channel*4)#64
+		self.ess = ResidualBlock(channel*8)#128
+
+		self.maxpool = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
+		self.conv_eltem = nn.Conv2d(channel,2*channel,kernel_size=1,stride=1,padding=0,bias=False)#16 32
+		self.conv_emtes = nn.Conv2d(2*channel,4*channel,kernel_size=1,stride=1,padding=0,bias=False)#32 64
+		self.conv_estess = nn.Conv2d(4*channel,8*channel,kernel_size=1,stride=1,padding=0,bias=False)#64 128
+		self.conv_esstesss = nn.Conv2d(8*channel,16*channel,kernel_size=1,stride=1,padding=0,bias=False)#128 256
+
+	def forward(self,x,embedding):
+
+		elout = self.el(x, embedding)#16
+		x_emin = self.conv_eltem(self.maxpool(elout))#32
+		emout = self.em(x_emin, embedding)
+		x_esin = self.conv_emtes(self.maxpool(emout))        
+		esout = self.es(x_esin, embedding)
+		x_esin = self.conv_estess(self.maxpool(esout))        
+		essout = self.ess(x_esin, embedding)#128
+
+		return elout, emout, esout, essout#,esssout
+
+class backbone(nn.Module):
+	def __init__(self,channel):
+		super(backbone,self).__init__()    
+
+		self.s1 = ResidualBlock(channel*8)#128
+		self.s2 = ResidualBlock(channel*8)#128
+
+	def forward(self,x,embedding):
+
+		share1 = self.s1(x, embedding)
+		share2 = self.s2(share1, embedding)
+
+		return share2
+
+class decoder(nn.Module):
+	def __init__(self,channel):
+		super(decoder,self).__init__()    
+
+		self.dss = ResidualBlock(channel*8)#128
+		self.ds = ResidualBlock(channel*4)#64
+		self.dm = ResidualBlock(channel*2)#32
+		self.dl = ResidualBlock(channel)#16
+
+		#self.conv_dssstdss = nn.Conv2d(16*channel,8*channel,kernel_size=1,stride=1,padding=0,bias=False)#256 128
+		self.conv_dsstds = nn.Conv2d(8*channel,4*channel,kernel_size=1,stride=1,padding=0,bias=False)#128 64
+		self.conv_dstdm = nn.Conv2d(4*channel,2*channel,kernel_size=1,stride=1,padding=0,bias=False)#64 32
+		self.conv_dmtdl = nn.Conv2d(2*channel,channel,kernel_size=1,stride=1,padding=0,bias=False)#32 16
+
+	def _upsample(self,x,y):
+		_,_,H0,W0 = y.size()
+		return F.interpolate(x,size=(H0,W0),mode='bilinear')
+
+	def forward(self, x, x_ss, x_s, x_m, x_l, embedding):
+
+		dssout = self.dss(x + x_ss, embedding)
+		x_dsin = self.conv_dsstds(self._upsample(dssout, x_s))        
+		dsout = self.ds(x_dsin + x_s, embedding)
+		x_dmin = self.conv_dstdm(self._upsample(dsout, x_m))
+		dmout = self.dm(x_dmin + x_m, embedding)
+		x_dlin = self.conv_dmtdl(self._upsample(dmout, x_l))
+		dlout = self.dl(x_dlin + x_l, embedding)
+
+		return dlout
+
+
+class ResidualBlock(nn.Module):  # Edge-oriented Residual Convolution Block 面向边缘的残差网络块 解决梯度消失的问题
+	def __init__(self, channel, norm=False):
+		super(ResidualBlock, self).__init__()
+
+		self.el = TransformerBlock(channel, num_heads=8, ffn_expansion_factor=2.66, bias=False, LayerNorm_type='WithBias')
+
+	def forward(self, x,embedding):
+		return self.el(x,embedding)
+
+def to_3d(x):
+	return rearrange(x, 'b c h w -> b (h w) c')
+
+def to_4d(x, h, w):
+	return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+
+
+class BiasFree_LayerNorm(nn.Module):
+	def __init__(self, normalized_shape):
+		super(BiasFree_LayerNorm, self).__init__()
+		if isinstance(normalized_shape, numbers.Integral):
+			normalized_shape = (normalized_shape,)
+		normalized_shape = torch.Size(normalized_shape)
+		assert len(normalized_shape) == 1
+		self.weight = nn.Parameter(torch.ones(normalized_shape))
+		self.normalized_shape = normalized_shape
+
+	def forward(self, x):
+		sigma = x.var(-1, keepdim=True, unbiased=False)
+		return x / torch.sqrt(sigma + 1e-5) * self.weight
+
+class WithBias_LayerNorm(nn.Module):
+	def __init__(self, normalized_shape):
+		super(WithBias_LayerNorm, self).__init__()
+		if isinstance(normalized_shape, numbers.Integral):
+			normalized_shape = (normalized_shape,)
+		normalized_shape = torch.Size(normalized_shape)
+		assert len(normalized_shape) == 1
+
+		self.weight = nn.Parameter(torch.ones(normalized_shape))
+		self.bias = nn.Parameter(torch.zeros(normalized_shape))
+		self.normalized_shape = normalized_shape
+
+	def forward(self, x):
+		mu = x.mean(-1, keepdim=True)
+		sigma = x.var(-1, keepdim=True, unbiased=False)
+		return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
+
+class LayerNorm(nn.Module):
+	def __init__(self, dim, LayerNorm_type):
+		super(LayerNorm, self).__init__()
+		if LayerNorm_type == 'BiasFree':
+			self.body = BiasFree_LayerNorm(dim)
+		else:
+			self.body = WithBias_LayerNorm(dim)
+
+	def forward(self, x):
+		h, w = x.shape[-2:]
+		return to_4d(self.body(to_3d(x)), h, w)
+
+class Cross_Attention(nn.Module):
+    def __init__(self, 
+		 		dim, 
+				num_heads, 
+				bias,
+				q_dim = 324):
+        super(Cross_Attention, self).__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        sqrt_q_dim = int(math.sqrt(q_dim))
+        self.resize = transforms.Resize([sqrt_q_dim, sqrt_q_dim])
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.q = nn.Linear(q_dim, q_dim, bias=bias)
+        
+        self.kv = nn.Conv2d(dim, dim*2, kernel_size=1, bias=bias)
+        self.kv_dwconv = nn.Conv2d(dim*2, dim*2, kernel_size=3, stride=1, padding=1, groups=dim*2, bias=bias)
+        
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+    def forward(self, x, query):
+        b,c,h,w = x.shape
+
+        q = self.q(query)
+        k, v = self.kv_dwconv(self.kv(x)).chunk(2, dim=1) 
+        k = self.resize(k)
+	
+        q = repeat(q, 'b l -> b head c l', head=self.num_heads, c=self.dim//self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+        
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+class Self_Attention(nn.Module):
+    def __init__(self, 
+		 		dim, 
+				num_heads, 
+				bias):
+        super(Self_Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Conv2d(dim, dim*3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(dim*3, dim*3, kernel_size=3, stride=1, padding=1, groups=dim*3, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+    def forward(self, x):
+        b,c,h,w = x.shape
+
+        qkv = self.qkv_dwconv(self.qkv(x))
+        q,k,v = qkv.chunk(3, dim=1)   
+        
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+        
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+class FeedForward(nn.Module):
+	def __init__(self, 
+	      		dim, 
+				ffn_expansion_factor, 
+				bias):
+		super(FeedForward, self).__init__()
+
+		hidden_features = int(dim * ffn_expansion_factor)
+
+		self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
+
+		self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1,
+                                groups=hidden_features * 2, bias=bias)
+
+		self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+	def forward(self, x):
+		x = self.project_in(x)
+		x1, x2 = self.dwconv(x).chunk(2, dim=1)
+		x = F.gelu(x1) * x2
+		x = self.project_out(x)
+		return x
+
+class TransformerBlock(nn.Module):
+	def __init__(self, 
+	      		dim, 
+				num_heads=8, 
+				ffn_expansion_factor=2.66, 
+				bias=False, 
+				LayerNorm_type='WithBias'):
+		super(TransformerBlock, self).__init__()
+		self.norm1 = LayerNorm(dim, LayerNorm_type)
+		self.cross_attn = Cross_Attention(dim, num_heads, bias)
+		self.norm2 = LayerNorm(dim, LayerNorm_type)
+		self.self_attn = Self_Attention(dim, num_heads, bias)
+		self.norm3 = LayerNorm(dim, LayerNorm_type)
+		self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
+
+	def forward(self, x, query):
+		x = x + self.cross_attn(self.norm1(x),query)
+		x = x + self.self_attn(self.norm2(x))
+		x = x + self.ffn(self.norm3(x))
+		return x
+
+if __name__ == '__main__':
+	net = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+	# x = torch.Tensor(np.random.random((2,3,256,256))).to("cuda" if torch.cuda.is_available() else "cpu")
+	# query = torch.Tensor(np.random.random((2, 324))).to("cuda" if torch.cuda.is_available() else "cpu")
+	# out = net(x, query)
+	# print(out.shape)
+	input = torch.randn(1, 3, 512, 512).to("cuda" if torch.cuda.is_available() else "cpu")
+	query = torch.Tensor(np.random.random((1, 324))).to("cuda" if torch.cuda.is_available() else "cpu")
+	macs, _ = profile(net, inputs=(input, query))
+	total = sum([param.nelement() for param in net.parameters()])
+	print('Macs = ' + str(macs/1000**3) + 'G')
+	print('Params = ' + str(total/1e6) + 'M')
+
+	from fvcore.nn import FlopCountAnalysis, parameter_count_table
+	flops = FlopCountAnalysis(net, (input, query))
+	print("FLOPs", flops.total()/1000**3)
+
+

model/loss.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+import cv2 as cv
+import numpy as np
+from matplotlib import pyplot as plt
+from math import exp
+from torchvision import transforms
+from torchvision.models import vgg16
+import torchvision
+'''
+MS-SSIM Loss
+'''
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+
+
+def create_window(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+    return window
+
+
+def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
+    # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
+    if val_range is None:
+        if torch.max(img1) > 128:
+            max_val = 255
+        else:
+            max_val = 1
+
+        if torch.min(img1) < -0.5:
+            min_val = -1
+        else:
+            min_val = 0
+        L = max_val - min_val
+    else:
+        L = val_range
+
+    padd = 0
+    (_, channel, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window(real_size, channel=channel).to(img1.device)
+
+    mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2
+
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+
+    if full:
+        return ret, cs
+    return ret
+
+
+def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
+    weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(img1.device)
+    levels = weights.size()[0]
+    mssim = []
+    mcs = []
+    for _ in range(levels):
+        sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
+        mssim.append(sim)
+        mcs.append(cs)
+
+        img1 = F.avg_pool2d(img1, (2, 2))
+        img2 = F.avg_pool2d(img2, (2, 2))
+
+    mssim = torch.stack(mssim)
+    mcs = torch.stack(mcs)
+
+    # Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
+    if normalize:
+        mssim = (mssim + 1) / 2
+        mcs = (mcs + 1) / 2
+
+    pow1 = mcs ** weights
+    pow2 = mssim ** weights
+    # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
+    output = torch.prod(pow1[:-1] * pow2[-1])
+    return output
+
+
+# Classes to re-use window
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, val_range=None):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.val_range = val_range
+
+        # Assume 1 channel for SSIM
+        self.channel = 1
+        self.window = create_window(window_size)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.dtype == img1.dtype:
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
+            self.window = window
+            self.channel = channel
+
+        return ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
+
+class MSSSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, channel=3):
+        super(MSSSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = channel
+
+    def forward(self, img1, img2):
+        # TODO: store window between calls if possible
+        return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)
+
+class TVLoss(nn.Module):
+    def __init__(self,TVLoss_weight=1):
+        super(TVLoss,self).__init__()
+        self.TVLoss_weight = TVLoss_weight
+ 
+    def forward(self,x):
+        batch_size = x.size()[0]
+        h_x = x.size()[2]
+        w_x = x.size()[3]
+        count_h = self._tensor_size(x[:,:,1:,:])  #算出总共求了多少次差
+        count_w = self._tensor_size(x[:,:,:,1:])
+        h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum()  
+        # x[:,:,1:,:]-x[:,:,:h_x-1,:]就是对原图进行错位，分成两张像素位置差1的图片，第一张图片
+        # 从像素点1开始（原图从0开始），到最后一个像素点，第二张图片从像素点0开始，到倒数第二个            
+        # 像素点，这样就实现了对原图进行错位，分成两张图的操作，做差之后就是原图中每个像素点与相
+        # 邻的下一个像素点的差。
+        w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum()
+        return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size
+ 
+    def _tensor_size(self,t):
+        return t.size()[1]*t.size()[2]*t.size()[3]
+ 
+    def _tensor_size(self,t):
+        return t.size()[1]*t.size()[2]*t.size()[3]
+
+class ContrastLoss(nn.Module):
+    def __init__(self):
+        super(ContrastLoss, self).__init__()
+        self.l1 = nn.L1Loss()
+        self.model = vgg16(weights = torchvision.models.VGG16_Weights.DEFAULT)
+        self.model = self.model.features[:16].to("cuda" if torch.cuda.is_available() else "cpu")
+        for param in self.model.parameters():
+            param.requires_grad = False
+        self.layer_name_mapping = {
+            '3': "relu1_2",
+            '8': "relu2_2",
+            '15': "relu3_3"
+        }
+
+    def gen_features(self, x):
+        output = []
+        for name, module in self.model._modules.items():
+            x = module(x)
+            if name in self.layer_name_mapping:
+                output.append(x)
+        return output
+    def forward(self, inp, pos, neg, out):
+        inp_t = inp
+        inp_x0 = self.gen_features(inp_t)
+        pos_t = pos
+        pos_x0 = self.gen_features(pos_t)
+        out_t = out
+        out_x0 = self.gen_features(out_t)
+        neg_t, neg_x0 = [],[]
+        for i in range(neg.shape[1]):
+            neg_i = neg[:,i,:,:]
+            neg_t.append(neg_i)
+            neg_x0_i = self.gen_features(neg_i)
+            neg_x0.append(neg_x0_i)
+        loss = 0
+        for i in range(len(pos_x0)):
+            pos_term = self.l1(out_x0[i], pos_x0[i].detach())
+            inp_term = self.l1(out_x0[i], inp_x0[i].detach())/(len(neg_x0)+1)
+            neg_term = sum(self.l1(out_x0[i], neg_x0[j][i].detach()) for j in range(len(neg_x0)))/(len(neg_x0)+1)
+            loss = loss + pos_term / (inp_term+neg_term+1e-7)
+        return loss / len(pos_x0)
+
+class Total_loss(nn.Module):
+    def __init__(self, args):
+        super(Total_loss, self).__init__()
+        self.con_loss = ContrastLoss()
+        self.weight_sl1, self.weight_msssim, self.weight_drl = args.loss_weight
+
+    def forward(self, inp, pos, neg, out):
+        smooth_loss_l1 = F.smooth_l1_loss(out, pos)
+        msssim_loss = 1-msssim(out, pos, normalize=True)
+        c_loss = self.con_loss(inp[0], pos, neg, out)
+
+        total_loss = self.weight_sl1 * smooth_loss_l1 + self.weight_msssim * msssim_loss + self.weight_drl * c_loss
+        return total_loss

remove_optim.py

@@ -0,0 +1,32 @@
+import torch, argparse
+from model.OneRestore import OneRestore
+from model.Embedder import Embedder
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--type", type=str, default = 'OneRestore')
+parser.add_argument("--input-file", type=str, default = './ckpts/onerestore_cdd-11.tar')
+parser.add_argument("--output-file", type=str, default = './ckpts/onerestore_cdd-11.tar')
+
+args = parser.parse_args()
+
+if args.type == 'OneRestore':
+    restorer = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+    restorer_info = torch.load(args.input_file, map_location='cuda:0')
+    weights_dict = {}
+    for k, v in restorer_info['state_dict'].items():
+        new_k = k.replace('module.', '') if 'module' in k else k
+        weights_dict[new_k] = v
+    restorer.load_state_dict(weights_dict)
+    torch.save(restorer.state_dict(), args.output_file)
+elif args.type == 'Embedder':
+    combine_type = ['clear', 'low', 'haze', 'rain', 'snow',\
+                    'low_haze', 'low_rain', 'low_snow', 'haze_rain',\
+                        'haze_snow', 'low_haze_rain', 'low_haze_snow']
+    embedder = Embedder(combine_type).to("cuda" if torch.cuda.is_available() else "cpu")
+    embedder_info = torch.load(args.input_file)
+    embedder.load_state_dict(embedder_info['state_dict'])
+    torch.save(embedder.state_dict(), args.output_file)
+else:
+    print('ERROR!')
+

requirements.txt

@@ -0,0 +1,10 @@
+pillow
+numpy
+scikit-image
+pandas
+einops
+thop
+fasttext
+opencv-python
+h5py
+matplotlib

syn_data/syn_data.py

@@ -0,0 +1,86 @@
+import os, argparse, cv2, random
+import numpy as np
+from skimage import exposure
+
+def guideFilter(I, p, winSize, eps):
+    mean_I = cv2.blur(I, winSize)
+    mean_p = cv2.blur(p, winSize)
+    mean_II = cv2.blur(I * I, winSize)
+    mean_Ip = cv2.blur(I * p, winSize)
+    var_I = mean_II - mean_I * mean_I
+    cov_Ip = mean_Ip - mean_I * mean_p
+    a = cov_Ip / (var_I + eps)
+    b = mean_p - a * mean_I
+    mean_a = cv2.blur(a, winSize)
+    mean_b = cv2.blur(b, winSize)
+    q = mean_a * I + mean_b
+    return q
+
+def syn_low(img, light, img_gray, light_max=3,
+            light_min=2, noise_max=0.08, noise_min=0.03):
+    light = guideFilter(light, img_gray,(3,3),0.01)[:, :, np.newaxis]
+    n = np.random.uniform(noise_min, noise_max)
+    R = img / (light + 1e-7)
+    L = (light + 1e-7) ** np.random.uniform(light_min, light_max)
+    return np.clip(R * L + np.random.normal(0, n, img.shape), 0, 1)
+
+def syn_haze(img, depth, beta_max=2.0, beta_min=1.0, A_max=0.9, A_min=0.6,
+                 color_max=0, color_min=0):
+    beta = np.random.rand(1) * (beta_max - beta_min) + beta_min
+    t = np.exp(-np.minimum(1 - cv2.blur(depth,(22,22)),0.7) * beta)
+    A = np.random.rand(1) * (A_max - A_min) + A_min
+    A_random = np.random.rand(3) * (color_max - color_min) + color_min
+    A = A + A_random
+    return np.clip(img * t + A * (1 - t), 0, 1)
+
+def syn_data(hq_file, light_file, depth_file, rain_file, snow_file, out_file, 
+             low, haze, rain, snow):
+    file_list = os.listdir(hq_file)
+    rain_list = os.listdir(rain_file)
+    snow_list = os.listdir(snow_file)
+    num_rain = random.sample(range(0,len(rain_list)),len(rain_list))
+    num_snow = random.sample(range(0,len(snow_list)),len(snow_list))
+    for i in range(1, len(file_list)):
+        img = cv2.imread(hq_file+file_list[i])
+        w, h, _ = img.shape
+        light = cv2.cvtColor(cv2.imread(light_file + file_list[i]), cv2.COLOR_RGB2GRAY) / 255.0
+        depth = cv2.imread(depth_file + file_list[i]) / 255.0
+        rain_mask = cv2.imread(rain_file + rain_list[num_rain[i]]) / 255.0
+        rain_mask = cv2.resize(rain_mask,(h,w))
+        snow_mask = cv2.imread(snow_file + snow_list[num_snow[i]]) / 255.0
+        snow_mask = cv2.resize(snow_mask, (h, w))
+        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)/ 255.0
+        lq = img.copy()/255.0
+        color_dis  = 1
+
+        if low:
+            lq = syn_low(lq, light, img_gray)
+        if rain:
+            lq = lq+rain_mask
+        if snow:
+            lq = lq*(1-snow_mask)+color_dis*snow_mask
+        if haze:
+            lq = syn_haze(lq, depth)
+
+        # out = np.concatenate((lq*255.0,img),1)
+        out = lq*255.0
+        cv2.imwrite(out_file + file_list[i], out)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+     # load model
+    parser.add_argument("--hq-file", type=str, default = './data/clear/')
+    parser.add_argument("--light-file", type=str, default = './data/light_map/')
+    parser.add_argument("--depth-file", type=str, default = './data/depth_map/')
+    parser.add_argument("--rain-file", type=str, default = './data/rain_mask/')
+    parser.add_argument("--snow-file", type=str, default = './data/snow_mask/')
+    parser.add_argument("--out-file", type=str, default = './out/')
+    parser.add_argument("--low", action='store_true')
+    parser.add_argument("--haze", action='store_true')
+    parser.add_argument("--rain", action='store_true')
+    parser.add_argument("--snow", action='store_true')
+    
+    args = parser.parse_args()
+
+    syn_data(args.hq_file, args.light_file, args.depth_file, args.rain_file, 
+             args.snow_file, args.out_file, args.low, args.haze, args.rain, args.snow)

test.py

@@ -0,0 +1,82 @@
+import os, time, argparse
+from PIL import Image
+import numpy as np
+
+
+import torch
+from torchvision import transforms
+
+from torchvision.utils import save_image as imwrite
+from utils.utils import print_args, load_restore_ckpt, load_embedder_ckpt
+
+transform_resize = transforms.Compose([
+        transforms.Resize([224,224]),
+        transforms.ToTensor()
+        ]) 
+
+def main(args):
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    #train
+    print('> Model Initialization...')
+
+    embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=args.embedder_model_path)
+    restorer = load_restore_ckpt(device, freeze_model=True, ckpt_name=args.restore_model_path)
+
+    os.makedirs(args.output,exist_ok=True)
+    
+    files = os.listdir(argspar.input)
+    time_record = []
+    for i in files:
+        lq = Image.open(f'{argspar.input}/{i}')
+
+        with torch.no_grad():
+            lq_re = torch.Tensor((np.array(lq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+            lq_em = transform_resize(lq).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+
+            start_time = time.time()
+            
+            if args.prompt == None:
+                text_embedding, _, [text] = embedder(lq_em,'image_encoder')
+                print(f'This is {text} degradation estimated by visual embedder.')
+            else:
+                text_embedding, _, [text] = embedder([args.prompt],'text_encoder')
+                print(f'This is {text} degradation generated by input text.')
+            
+            out = restorer(lq_re, text_embedding)
+
+            run_time = time.time()-start_time
+            time_record.append(run_time)
+
+            if args.concat:
+                out = torch.cat((lq_re, out), dim=3)
+
+            imwrite(out, f'{args.output}/{i}', range=(0, 1))
+
+            print(f'{i} Running Time: {run_time:.4f}.')
+    print(f'Average time is {np.mean(np.array(run_time))}')
+            
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description = "OneRestore Running")
+
+    # load model
+    parser.add_argument("--embedder-model-path", type=str, default = "./ckpts/embedder_model.tar", help = 'embedder model path')
+    parser.add_argument("--restore-model-path", type=str, default = "./ckpts/onerestore_cdd-11.tar", help = 'restore model path')
+
+    # select model automatic (prompt=False) or manual (prompt=True, text={'clear', 'low', 'haze', 'rain', 'snow',\
+    #                'low_haze', 'low_rain', 'low_snow', 'haze_rain', 'haze_snow', 'low_haze_rain', 'low_haze_snow'})
+    parser.add_argument("--prompt", type=str, default = None, help = 'prompt')
+
+    parser.add_argument("--input", type=str, default = "./image/", help = 'image path')
+    parser.add_argument("--output", type=str, default = "./output/", help = 'output path')
+    parser.add_argument("--concat", action='store_true', help = 'output path')
+
+    argspar = parser.parse_args()
+
+    print_args(argspar)
+
+    main(argspar)

train_Embedder.py

@@ -0,0 +1,104 @@
+import argparse, os, torch, time
+import torch.optim
+
+from utils.utils import load_embedder_ckpt_with_optim, adjust_learning_rate, freeze_text_embedder, AverageMeter
+from utils.utils_data import init_embedding_data
+
+
+
+def train_embedding(cur_epoch, model, optimizer, trainloader, testloader, device, cfg_em):
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.enabled = True
+    
+    acc_train_meter = AverageMeter()
+    acc_test_meter = AverageMeter()
+    loss_train_meter = AverageMeter()
+    loss_test_meter = AverageMeter()
+    time_train_meter = AverageMeter()
+    time_test_meter = AverageMeter()
+
+    freeze_text_embedder(model)
+    for k,v in model.named_parameters():
+        print('{}: {}'.format(k, v.requires_grad))
+    for epoch in range(cur_epoch, cfg_em.epoch+1):
+
+        optimizer = adjust_learning_rate(optimizer, epoch-1, cfg_em.lr_decay)
+        lr = optimizer.param_groups[-1]['lr']
+
+        model.train()
+        for idx, batch in enumerate(trainloader):
+            for i in range(len(batch)):
+                batch[i] = batch[i].to("cuda" if torch.cuda.is_available() else "cpu")
+            time_start = time.time()
+            out = model(batch, 'train')
+            loss = out['loss_total']
+            acc  = out['acc_type']
+            time_train_meter.update(time.time() - time_start)
+
+            acc_train_meter.update(acc)
+            loss_train_meter.update(loss)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            print(f'Epoch:{epoch}|Iter:{idx+1}/{len(trainloader)}|lr:{lr},'
+                f'Loss: {loss_train_meter.avg:.3f},' 
+                f'Acc: {acc_train_meter.avg:.3f},'
+                f'Time: {time_train_meter.avg:.3f},', flush=True)
+        
+        model.eval()
+        for idx, batch in enumerate(testloader):
+            for i in range(len(batch)):
+                batch[i] = batch[i].to("cuda" if torch.cuda.is_available() else "cpu")
+
+            time_start = time.time()
+            out = model(batch, 'train')
+            loss = out['loss_total']
+            acc  = out['acc_type']
+            time_test_meter.update(time.time() - time_start)
+
+            acc_test_meter.update(acc)
+            loss_test_meter.update(loss)
+            print(f'Epoch:{epoch}|Iter:{idx+1}/{len(testloader)}|lr:{lr},'
+                f'Loss: {loss_test_meter.avg:.3f},'
+                f'Acc: {acc_test_meter.avg:.3f},'
+                f'Time: {time_test_meter.avg:.3f},', flush=True)
+        
+        torch.save({'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer' : optimizer.state_dict()},
+                         f'{cfg_em.check_dir}/embedder_model_epoch{epoch}_{acc_train_meter.avg:.3f}_{loss_train_meter.avg:.3f}_{acc_test_meter.avg:.3f}_{loss_test_meter.avg:.3f}.tar')
+        acc_train_meter.reset()
+        acc_test_meter.reset()
+        loss_train_meter.reset()
+        loss_test_meter.reset()
+        time_train_meter.reset()
+        time_test_meter.reset()
+    print('Done!')
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+     # load model
+    parser.add_argument("--seed", type=int, default = 124)
+    parser.add_argument("--pre_weight", type=str, default = '')
+    parser.add_argument("--lr", type=float, default = 0.0001)
+    parser.add_argument("--type_name", type=list, default = ['clear', 'low', 'haze', 'rain',\
+                                        'snow', 'low_haze', 'low_rain', 'low_snow', 'haze_rain',\
+                                        'haze_snow', 'low_haze_rain', 'low_haze_snow'])
+    parser.add_argument("--train-dir", type=str, default = './data/CDD-11_train/')
+    parser.add_argument("--test-dir", type=str, default = './data/CDD-11_test/')
+    parser.add_argument("--batch", type=int, default = 128)
+    parser.add_argument("--num-workers", type=int, default = 0)
+    parser.add_argument("--epoch", type=int, default = 200)
+    parser.add_argument("--lr-decay", type=int, default = 50)
+    parser.add_argument("--check-dir", type=str, default = "./ckpts")
+    
+    args = parser.parse_args()
+
+    os.makedirs(args.check_dir,exist_ok=True)
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    embedder, optimizer, cur_epoch, device = load_embedder_ckpt_with_optim(device, args)
+    trainloader, testloader = init_embedding_data(args, 'train')
+    train_embedding(cur_epoch, embedder, optimizer, trainloader, testloader, device, args)

train_OneRestore_multi-gpu.py

@@ -0,0 +1,153 @@
+import os, time, torch, argparse
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.utils import save_image as imwrite
+import numpy as np
+from torchvision import transforms
+from makedataset import Dataset
+from utils.utils import print_args, load_restore_ckpt_with_optim, load_embedder_ckpt, adjust_learning_rate, data_process, tensor_metric, load_excel, save_checkpoint
+from model.loss import Total_loss
+from model.Embedder import Embedder
+from model.OneRestore import OneRestore
+from torch.utils.data.distributed import DistributedSampler
+from PIL import Image
+
+torch.distributed.init_process_group(backend="nccl")
+local_rank = torch.distributed.get_rank()
+torch.cuda.set_device(local_rank)
+device = torch.device("cuda", local_rank)
+
+
+transform_resize = transforms.Compose([
+        transforms.Resize([224,224]),
+        transforms.ToTensor()
+        ]) 
+
+def main(args):
+    
+
+    print('> Model Initialization...')
+    embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=args.embedder_model_path)
+    restorer, optimizer, cur_epoch = load_restore_ckpt_with_optim(device, local_rank=local_rank, freeze_model=False, ckpt_name=args.restore_model_path, lr=args.lr)
+    loss = Total_loss(args)
+    
+    print('> Loading dataset...')
+    data = Dataset(args.train_input)
+    dataset = DataLoader(dataset=data, batch_size=args.bs,
+                                    shuffle=False, 
+                                    num_workers=args.num_works,
+                                    pin_memory=True,drop_last=False,                                     
+                                    sampler=DistributedSampler(data,shuffle=True))
+    
+    print('> Start training...')
+    start_all = time.time()
+    train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device)
+    end_all = time.time()
+    print('Whloe Training Time:' +str(end_all-start_all)+'s.')
+
+def train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device):
+
+    metric = []
+    for epoch in range(cur_epoch, args.epoch):
+        optimizer = adjust_learning_rate(optimizer, epoch, args.adjust_lr)
+        learnrate = optimizer.param_groups[-1]['lr']
+        restorer.train()
+
+        for i, data in enumerate(dataset,0):
+            pos, inp, neg = data_process(data, args, device)
+
+            text_embedding,_,_ = embedder(inp[1],'text_encoder')
+            out = restorer(inp[0], text_embedding)
+
+            restorer.zero_grad()
+            total_loss = loss(inp, pos, neg, out)
+            total_loss.backward()
+            optimizer.step()
+
+            mse = tensor_metric(pos,out, 'MSE', data_range=1)
+            psnr = tensor_metric(pos,out, 'PSNR', data_range=1)
+            ssim = tensor_metric(pos,out, 'SSIM', data_range=1)
+
+            print("[epoch %d][%d/%d] lr :%f Floss: %.4f MSE: %.4f PSNR: %.4f SSIM: %.4f"%(epoch+1, i+1, \
+                len(dataset), learnrate, total_loss.item(), mse, psnr, ssim))
+            
+
+        psnr_t1, ssim_t1, psnr_t2, ssim_t2 = test(args, restorer, embedder, device, epoch)
+        metric.append([psnr_t1, ssim_t1, psnr_t2, ssim_t2])
+        print("[epoch %d] Test images PSNR1: %.4f SSIM1: %.4f"%(epoch+1, psnr_t1,ssim_t1))
+
+        load_excel(metric)
+        save_checkpoint({'epoch': epoch + 1,'state_dict': restorer.state_dict(),'optimizer' : optimizer.state_dict()},\
+                        args.save_model_path, epoch+1, psnr_t1,ssim_t1,psnr_t2,ssim_t2)
+
+def test(args, restorer, embedder, device, epoch=-1):
+    combine_type = args.degr_type
+    psnr_1, psnr_2, ssim_1, ssim_2 = 0, 0, 0, 0
+    os.makedirs(args.output,exist_ok=True)
+
+    for i in range(len(combine_type)-1):
+        file_list =  os.listdir(f'{args.test_input}/{combine_type[i+1]}/')
+        for j in range(len(file_list)):
+            hq = Image.open(f'{args.test_input}/{combine_type[0]}/{file_list[j]}')
+            lq = Image.open(f'{args.test_input}/{combine_type[i+1]}/{file_list[j]}')
+            restorer.eval()
+            with torch.no_grad():
+                lq_re = torch.Tensor((np.array(lq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+                lq_em = transform_resize(lq).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+                hq = torch.Tensor((np.array(hq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+
+                starttime = time.time()
+
+                text_embedding_1,_,text_1 = embedder([combine_type[i+1]],'text_encoder')
+                text_embedding_2,_, text_2 = embedder(lq_em,'image_encoder')
+                out_1 = restorer(lq_re, text_embedding_1)
+                if text_1 != text_2:
+                    print(text_1, text_2)
+                    out_2 = restorer(lq_re, text_embedding_2)
+                else:
+                    out_2 = out_1
+                
+                endtime1 = time.time()
+
+                imwrite(torch.cat((lq_re, out_1, out_2, hq), dim=3), args.output \
+                    + file_list[j][:-4] + '_' + str(epoch) + '_' + combine_type[i+1] + '.png', range=(0, 1))
+                # due to the vision problem, you can replace above line by
+                # imwrite(torch.cat((lq_re, out_1, out_2, hq), dim=3), args.output \
+                #     + file_list[j][:-4] + '_' + str(epoch) + '_' + combine_type[i+1] + '.png')
+            psnr_1 += tensor_metric(hq, out_1, 'PSNR', data_range=1)
+            ssim_1 += tensor_metric(hq, out_1, 'SSIM', data_range=1)
+            psnr_2 += tensor_metric(hq, out_2, 'PSNR', data_range=1)
+            ssim_2 += tensor_metric(hq, out_2, 'SSIM', data_range=1)
+            print('The ' + file_list[j][:-4] + ' Time:' + str(endtime1 - starttime) + 's.')
+
+    return psnr_1 / (len(file_list)*len(combine_type)), ssim_1 / (len(file_list)*len(combine_type)),\
+        psnr_2 / (len(file_list)*len(combine_type)), ssim_2 / (len(file_list)*len(combine_type))
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description = "OneRestore Training")
+
+    # load model
+    parser.add_argument("--embedder-model-path", type=str, default = "./ckpts/embedder_model.tar", help = 'embedder model path')
+    parser.add_argument("--restore-model-path", type=str, default = None, help = 'restore model path')
+    parser.add_argument("--save-model-path", type=str, default = "./ckpts/", help = 'restore model path')
+
+    parser.add_argument("--epoch", type=int, default = 300, help = 'epoch number')
+    parser.add_argument("--bs", type=int, default = 4, help = 'batchsize')
+    parser.add_argument("--lr", type=float, default = 1e-4, help = 'learning rate')
+    parser.add_argument("--adjust-lr", type=int, default = 30, help = 'adjust learning rate')
+    parser.add_argument("--num-works", type=int, default = 4, help = 'number works')
+    parser.add_argument("--loss-weight", type=tuple, default = (0.6,0.3,0.1), help = 'loss weights')
+    parser.add_argument("--degr-type", type=list, default = ['clear', 'low', 'haze', 'rain', 'snow',\
+        'low_haze', 'low_rain', 'low_snow', 'haze_rain', 'haze_snow', 'low_haze_rain', 'low_haze_snow'], help = 'degradation type')
+    
+    parser.add_argument("--train-input", type=str, default = "./dataset.h5", help = 'train data')
+    parser.add_argument("--test-input", type=str, default = "./data/CDD-11_test", help = 'test path')
+    parser.add_argument("--output", type=str, default = "./result/", help = 'output path')
+
+    argspar = parser.parse_args()
+
+    print_args(argspar)
+    main(argspar)

train_OneRestore_single-gpu.py

@@ -0,0 +1,140 @@
+import os, time, torch, argparse
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.utils import save_image as imwrite
+import numpy as np
+from torchvision import transforms
+from makedataset import Dataset
+from utils.utils import print_args, load_restore_ckpt_with_optim, load_embedder_ckpt, adjust_learning_rate, data_process, tensor_metric, load_excel, save_checkpoint
+from model.loss import Total_loss
+
+from PIL import Image
+
+transform_resize = transforms.Compose([
+        transforms.Resize([224,224]),
+        transforms.ToTensor()
+        ]) 
+
+def main(args):
+    
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    print('> Model Initialization...')
+
+    embedder = load_embedder_ckpt(device, freeze_model=True, ckpt_name=args.embedder_model_path)
+    restorer, optimizer, cur_epoch = load_restore_ckpt_with_optim(device, freeze_model=False, ckpt_name=args.restore_model_path, lr=args.lr)
+    loss = Total_loss(args)
+    
+    print('> Loading dataset...')
+    data = Dataset(args.train_input)
+    dataset = DataLoader(dataset=data, num_workers=args.num_works, batch_size=args.bs, shuffle=True)
+    
+    print('> Start training...')
+    start_all = time.time()
+    train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device)
+    end_all = time.time()
+    print('Whloe Training Time:' +str(end_all-start_all)+'s.')
+
+def train(restorer, embedder, optimizer, loss, cur_epoch, args, dataset, device):
+
+    metric = []
+    for epoch in range(cur_epoch, args.epoch):
+        optimizer = adjust_learning_rate(optimizer, epoch, args.adjust_lr)
+        learnrate = optimizer.param_groups[-1]['lr']
+        restorer.train()
+
+        for i, data in enumerate(dataset,0):
+            pos, inp, neg = data_process(data, args, device)
+
+            text_embedding,_,_ = embedder(inp[1],'text_encoder')
+            out = restorer(inp[0], text_embedding)
+
+            restorer.zero_grad()
+            total_loss = loss(inp, pos, neg, out)
+            total_loss.backward()
+            optimizer.step()
+
+            mse = tensor_metric(pos,out, 'MSE', data_range=1)
+            psnr = tensor_metric(pos,out, 'PSNR', data_range=1)
+            ssim = tensor_metric(pos,out, 'SSIM', data_range=1)
+
+            print("[epoch %d][%d/%d] lr :%f Floss: %.4f MSE: %.4f PSNR: %.4f SSIM: %.4f"%(epoch+1, i+1, \
+                len(dataset), learnrate, total_loss.item(), mse, psnr, ssim))
+            
+
+        psnr_t1, ssim_t1, psnr_t2, ssim_t2 = test(args, restorer, embedder, device, epoch)
+        metric.append([psnr_t1, ssim_t1, psnr_t2, ssim_t2])
+        print("[epoch %d] Test images PSNR1: %.4f SSIM1: %.4f"%(epoch+1, psnr_t1,ssim_t1))
+
+        load_excel(metric)
+        save_checkpoint({'epoch': epoch + 1,'state_dict': restorer.state_dict(),'optimizer' : optimizer.state_dict()},\
+                        args.save_model_path, epoch+1, psnr_t1,ssim_t1,psnr_t2,ssim_t2)
+
+def test(args, restorer, embedder, device, epoch=-1):
+    combine_type = args.degr_type
+    psnr_1, psnr_2, ssim_1, ssim_2 = 0, 0, 0, 0
+    os.makedirs(args.output,exist_ok=True)
+
+    for i in range(len(combine_type)-1):
+        file_list =  os.listdir(f'{args.test_input}/{combine_type[i+1]}/')
+        for j in range(len(file_list)):
+            hq = Image.open(f'{args.test_input}/{combine_type[0]}/{file_list[j]}')
+            lq = Image.open(f'{args.test_input}/{combine_type[i+1]}/{file_list[j]}')
+            restorer.eval()
+            with torch.no_grad():
+                lq_re = torch.Tensor((np.array(lq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+                lq_em = transform_resize(lq).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+                hq = torch.Tensor((np.array(hq)/255).transpose(2, 0, 1)).unsqueeze(0).to("cuda" if torch.cuda.is_available() else "cpu")
+
+                starttime = time.time()
+
+                text_embedding_1,_,text_1 = embedder([combine_type[i+1]],'text_encoder')
+                text_embedding_2,_, text_2 = embedder(lq_em,'image_encoder')
+                out_1 = restorer(lq_re, text_embedding_1)
+                if text_1 != text_2:
+                    print(text_1, text_2)
+                    out_2 = restorer(lq_re, text_embedding_2)
+                else:
+                    out_2 = out_1
+                
+                endtime1 = time.time()
+
+                imwrite(torch.cat((lq_re, out_1, out_2, hq), dim=3), args.output \
+                    + file_list[j][:-4] + '_' + str(epoch) + '_' + combine_type[i+1] + '.png', range=(0, 1))
+            psnr_1 += tensor_metric(hq, out_1, 'PSNR', data_range=1)
+            ssim_1 += tensor_metric(hq, out_1, 'SSIM', data_range=1)
+            psnr_2 += tensor_metric(hq, out_2, 'PSNR', data_range=1)
+            ssim_2 += tensor_metric(hq, out_2, 'SSIM', data_range=1)
+            print('The ' + file_list[j][:-4] + ' Time:' + str(endtime1 - starttime) + 's.')
+
+    return psnr_1 / (len(file_list)*len(combine_type)), ssim_1 / (len(file_list)*len(combine_type)),\
+        psnr_2 / (len(file_list)*len(combine_type)), ssim_2 / (len(file_list)*len(combine_type))
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description = "OneRestore Training")
+
+    # load model
+    parser.add_argument("--embedder-model-path", type=str, default = "./ckpts/embedder_model.tar", help = 'embedder model path')
+    parser.add_argument("--restore-model-path", type=str, default = None, help = 'restore model path')
+    parser.add_argument("--save-model-path", type=str, default = "./ckpts/", help = 'restore model path')
+
+    parser.add_argument("--epoch", type=int, default = 300, help = 'epoch number')
+    parser.add_argument("--bs", type=int, default = 4, help = 'batchsize')
+    parser.add_argument("--lr", type=float, default = 1e-4, help = 'learning rate')
+    parser.add_argument("--adjust-lr", type=int, default = 30, help = 'adjust learning rate')
+    parser.add_argument("--num-works", type=int, default = 4, help = 'number works')
+    parser.add_argument("--loss-weight", type=tuple, default = (0.6,0.3,0.1), help = 'loss weights')
+    parser.add_argument("--degr-type", type=list, default = ['clear', 'low', 'haze', 'rain', 'snow',\
+        'low_haze', 'low_rain', 'low_snow', 'haze_rain', 'haze_snow', 'low_haze_rain', 'low_haze_snow'], help = 'degradation type')
+    
+    parser.add_argument("--train-input", type=str, default = "./dataset.h5", help = 'train data')
+    parser.add_argument("--test-input", type=str, default = "./data/CDD-11_test", help = 'test path')
+    parser.add_argument("--output", type=str, default = "./result/", help = 'output path')
+
+    argspar = parser.parse_args()
+
+    print_args(argspar)
+    main(argspar)

utils/glove.6B.300d.txt

@@ -0,0 +1,5 @@
+clear -0.081023 -0.29179 0.052021 -0.13324 0.028162 -0.0031446 -0.17156 0.063324 0.16568 -2.1722 -0.14127 0.087891 -0.2298 0.069017 0.21673 0.36556 -0.39979 -0.15506 0.099728 0.202 0.16989 0.14807 0.10938 -0.17141 -0.7258 -0.13189 -0.052768 -0.26383 -0.13189 -0.11408 0.081757 0.14773 -0.24342 0.0076364 -1.0992 0.13661 0.19262 -0.30012 0.031524 0.11439 -0.10854 0.21089 -0.037365 0.23449 0.054638 0.21505 0.023071 0.20918 -0.08606 -0.078589 -0.26945 -0.040802 -0.042601 -0.12093 -0.33614 0.25624 -0.35266 -0.17224 0.31018 0.6426 -0.036072 0.1558 0.26609 0.17298 -0.08158 0.0085636 0.13196 -0.11876 -0.19205 -0.32204 -0.092694 -0.19274 0.0056832 0.17194 0.24011 0.014739 0.091188 0.45903 0.0047753 -0.18136 -0.16434 0.012617 0.42791 0.075318 -0.042848 -0.055952 -0.071895 0.086806 0.078092 0.20169 -0.34189 -0.01975 -0.44579 -0.093254 0.23684 0.098079 -0.0018186 -0.13013 0.054252 -0.68408 0.21378 -0.084742 -0.12383 0.36645 -0.46434 0.56799 0.22341 0.31607 -0.23559 0.033889 0.062509 -0.31468 0.27684 -0.13729 -0.027181 0.17143 -0.35535 0.14426 0.14137 -0.27987 0.051007 0.1689 0.48614 0.43247 -0.31014 -0.2273 -0.17253 0.50221 -0.29023 -0.16833 -0.027586 0.25614 0.096051 0.19145 -0.15576 0.50767 0.0064827 -0.047304 0.47358 -0.029665 -0.095882 0.064574 0.1247 -0.3439 -0.59591 -0.17307 0.30627 0.16351 -0.21709 -0.13142 -0.029781 0.079412 0.36018 -0.068721 0.367 0.26454 0.1306 -0.34602 0.22326 0.22999 0.14122 -0.3084 0.22239 -0.13701 0.24538 0.10902 0.33084 0.052159 -0.54817 0.32921 0.33889 -0.060382 -0.16611 -0.26388 0.13997 -0.15486 -0.05013 -0.089628 -0.0080954 0.13155 -0.019735 0.25758 0.37509 -0.012096 -0.49247 0.13436 -0.21072 -0.13763 0.24047 0.13328 -0.043418 0.0070651 0.30496 -0.11184 0.68017 -0.65417 -0.39198 0.075546 -0.2043 0.041099 0.84586 -0.3361 -0.26385 -0.39417 -0.25468 -0.095349 0.19947 -0.30772 -0.53846 0.18258 -0.091379 -0.27183 0.10918 -0.042102 -0.25614 -0.039694 0.34987 -0.24526 -0.011983 -0.024231 0.62785 -0.16641 0.026109 0.029096 -0.16937 0.25329 -0.12066 0.023087 0.16152 -0.14058 0.044846 0.4533 0.34099 -0.028432 -0.39407 -0.068924 -0.29128 -0.012954 0.048176 -0.090455 -0.0098771 -0.022352 0.091535 -0.084673 -0.43955 -0.25237 0.79719 0.21526 0.0019634 -0.10022 -0.075669 -0.25113 -0.12675 0.12179 0.25892 0.026661 -0.38419 -0.18566 -0.15325 0.44484 -0.088815 0.10119 0.0060884 0.293 -0.415 0.26712 0.033683 -0.42317 0.22025 -0.027351 0.40923 -0.013339 -0.29543 0.37699 -0.019656 -0.082896 -1.5198 0.2961 0.81263 -0.18199 0.59082 0.007938 0.2309 0.23573 0.24941 -0.18754 -0.04029 0.17258 0.1948 0.131 -0.21552 0.016352 0.62256 0.41283 0.40387 -0.062911 -0.093159 -0.078137 -0.30083 -0.035913
+low -0.21751 0.43389 0.149 0.14107 0.2574 -0.12448 0.0047523 0.035596 0.10741 -2.1047 0.17181 -0.15079 -0.044546 -0.090869 -0.43288 -0.13611 -0.0058198 -0.064724 0.23531 -0.36224 -0.21305 -0.075476 0.46786 -0.18465 -0.19746 -0.097471 0.39984 -0.084092 -0.53715 0.27303 -0.087786 0.24297 -0.38444 0.28854 -0.7873 0.089192 -0.26376 -0.16287 0.35911 0.30458 0.24502 0.22553 -0.0031653 0.47358 0.31146 -0.13823 0.075685 -0.10776 0.38329 -0.13762 0.51707 -0.16707 -0.037466 -0.7236 -0.4151 -0.42359 0.14354 0.046639 0.17527 0.48721 0.26708 -0.031042 0.86002 -0.3946 -0.50514 -0.51294 0.58527 0.18819 -0.29543 0.68596 -0.1035 0.22565 0.185 0.058375 0.030999 0.11929 0.12353 0.12873 0.42126 0.14188 -0.050079 -0.2683 0.12126 0.32302 0.27623 0.5414 0.074715 -0.1949 -0.47053 0.02313 0.68686 0.60158 -0.16194 -0.3651 0.41796 -0.22905 0.074734 0.17509 -0.44255 0.3518 -0.40079 -0.28305 0.39133 0.32303 -0.63198 -0.1507 -0.16894 0.17169 0.18894 0.027644 -0.36997 -0.26366 0.36344 -0.049584 0.32724 0.049712 0.051381 -0.058867 -0.2621 -0.50359 -0.21435 -0.25527 0.22161 0.66558 0.2224 0.27607 0.58587 -0.3071 0.24905 0.098802 -0.26459 0.77839 0.014585 0.86936 0.2329 -0.0027986 -0.087016 0.10863 0.18987 0.54552 0.24903 0.059293 0.30362 -0.028582 -0.6569 0.1206 -0.055416 -0.093077 -0.0012132 -0.15009 0.11192 -0.62139 -0.035773 0.1165 0.36541 0.55984 -0.19964 -0.065579 0.097118 -0.1672 0.13677 -0.95276 -0.25994 0.064799 -0.042161 0.12046 0.12391 0.0017478 0.29533 0.40176 0.057528 0.57864 -0.9973 0.13805 -0.30689 0.11015 -0.35402 -0.13434 -0.24479 0.50355 -0.18675 -0.22337 0.29573 0.21612 -0.068496 -0.60643 0.79013 -0.26975 -0.15492 0.70849 0.21372 0.62962 -0.0056421 0.53597 -0.54259 -0.34726 -0.29945 -0.51895 0.28471 -0.14973 0.54188 0.53535 -0.11233 0.19291 -0.24707 0.058424 -0.5473 -0.06426 0.47187 0.11149 0.28313 -0.23876 -0.10552 -0.051705 -0.28853 -0.13702 0.040562 -0.032269 0.10368 -0.29381 0.33416 0.038269 0.029697 -0.48604 -0.26334 0.28942 -0.0093944 0.13942 -0.29043 0.27332 0.16614 -0.028973 -0.32829 -0.034614 -0.0012628 0.062871 -0.000894 0.22467 0.16005 0.23141 -0.19918 0.16465 0.15247 0.29742 -1.0225 0.056188 0.91529 -0.47809 -0.24204 -0.3158 0.21033 -0.13616 0.10777 -0.26815 -0.44804 -0.12696 -0.43468 0.17849 -0.48101 0.026114 0.057368 0.26052 -0.030488 0.051275 -0.36344 0.11878 0.2279 -0.086855 -0.01455 0.070256 -0.16753 0.61449 -0.27428 -0.17901 -0.36261 0.093134 -1.5724 0.47192 -0.52493 -0.27512 -0.37945 0.29588 0.020506 0.08707 0.057053 0.37167 -0.056446 -0.38735 -0.31246 0.028304 -0.058202 0.067263 -0.58761 0.074556 0.49917 0.45134 -0.51433 -0.60996 0.076835 -0.078086
+haze -0.0061289 -0.2702 0.16559 -0.29621 -0.66216 -0.1756 0.46686 1.0362 -0.20692 -0.36097 0.98615 0.32297 -0.55094 -0.36163 -0.27046 0.052225 -0.10079 0.22536 -0.095491 0.17188 0.058372 0.083556 -0.28255 0.12623 -0.0094164 -0.028727 -0.20589 -0.3932 -0.2935 -0.36104 1.0595 0.14423 -0.311 -0.20573 0.11827 -0.0048368 -0.8324 -0.10389 0.34491 0.34006 0.10354 0.11593 0.47379 -0.1042 0.38523 -0.57589 0.027253 -0.44913 -0.52822 -0.44094 0.71219 -0.12278 0.034288 -0.6935 -0.57852 0.33917 0.35018 -0.30193 0.55504 0.085603 -0.21189 -0.51958 -0.17589 -0.13369 0.2976 -0.26048 0.068146 0.62144 0.3416 -0.54399 -0.23937 -0.34802 -0.31469 -0.59554 -0.25011 -0.11644 0.19993 -0.1636 0.24289 -0.0022965 0.3064 -0.26188 0.27166 0.1962 0.37527 -0.22408 0.52979 0.59141 0.035196 0.10632 -0.28318 0.18766 -0.12253 0.41932 -0.64713 0.26068 0.67209 -0.23333 0.030945 -0.15135 0.61662 -0.0025061 -0.58374 0.51866 -0.89244 1.0056 0.15919 0.29183 -0.059984 0.10701 -0.32101 -1.0921 -0.050394 -0.074584 0.56258 -0.5915 0.048547 0.085668 -0.39964 -0.40997 0.093632 -0.22538 -0.83102 -0.051418 -0.31192 0.36056 -0.028854 -0.046907 0.09394 0.012504 0.34555 0.56564 0.48111 0.092143 0.82492 -0.20086 -0.27718 0.9004 0.38921 0.028667 0.78904 0.44698 -0.26892 0.073712 -0.73296 -0.46286 0.53386 0.53514 0.04207 -0.11448 0.27771 0.080703 -0.017482 0.43225 0.047742 -0.095399 -0.063173 -0.36341 0.2948 0.15311 -0.55934 -0.88294 0.62005 -0.23936 0.51953 -0.49463 0.41669 0.61169 -0.20471 -0.0056962 -0.29331 0.46269 0.084808 -0.049355 -0.64697 -0.85777 0.34718 -0.16176 0.14756 -0.65658 -0.54259 -0.13124 -0.88851 0.070637 -0.84926 -0.69345 0.4024 -0.5683 -0.68142 -0.1402 -0.36857 0.36013 -0.49769 -0.17478 0.77214 -0.23962 0.32951 1.0984 -0.00011441 0.9649 -0.13312 0.64326 -0.037091 0.35672 0.025156 0.046782 0.19764 -0.22757 -0.39887 -0.3045 -0.45283 -0.0045182 0.032546 -0.076483 0.72189 -0.038917 1.0621 -0.55688 0.56429 0.11264 0.40465 -0.53146 0.16851 0.69236 -0.24456 0.038704 0.69151 0.16591 -0.43451 0.14115 0.84069 0.29081 -0.31053 -0.6849 -0.27188 -0.32813 0.57882 0.13779 0.36621 -0.45935 0.27899 -0.32315 -0.5743 0.19837 0.0046648 0.18459 0.43369 0.22359 0.16652 -0.081114 -0.54539 -1.0103 -0.14539 0.12021 0.078636 -0.26667 -0.65403 0.4096 0.07257 0.036639 0.21757 0.25738 0.51675 -0.031326 -0.3869 0.012763 -0.45692 0.13828 -0.48614 -0.53757 0.50268 0.47865 -0.049528 -0.032281 -0.4486 0.036258 -0.12295 -0.46811 -0.019014 0.035839 -0.55749 0.018281 -0.88963 -0.024676 -0.19482 -0.19364 0.0069875 0.12679 -0.37379 -0.34094 -0.051568 0.55404 -0.29656 0.26045 0.50872 -0.37399 0.20334 0.70298 -0.3271 -0.24116
+rain -0.52618 -0.54041 -0.89537 -0.35598 -0.74356 -0.66838 0.26326 0.89254 0.14362 -0.34904 0.25866 -0.11143 -0.52035 0.1436 -0.075728 -0.84569 -0.28762 0.049872 0.39234 0.52551 -0.39244 -0.2822 -0.097458 -0.12929 -0.38623 0.17261 0.7574 -0.29868 -0.691 -0.36639 0.63951 0.25255 -0.22299 0.16387 -0.83199 -0.30276 -0.32411 -0.36789 -0.073673 0.54726 0.14785 0.26259 0.086208 -0.033827 0.044403 -0.2135 0.3761 0.33816 -0.36696 -0.2096 0.025934 0.47679 0.23046 -0.44333 -0.65379 0.85762 0.62861 -0.70343 1.1284 0.2497 -0.34459 0.17005 0.27826 0.01167 -0.44087 -0.12649 0.31811 0.073688 -0.17127 -0.023486 0.34294 0.18888 -0.15694 -0.37975 -0.58313 -0.45624 -0.5968 0.09743 -0.50593 -0.64092 0.083647 0.38474 -0.15071 0.55042 -0.68742 0.14893 -0.039046 -0.19582 0.61498 -0.066786 0.63395 -0.4659 0.44123 -0.55136 -0.17711 0.97118 0.26321 -0.035901 -0.11096 -0.11161 0.353 1.026 -0.2605 -0.12231 0.31695 0.35807 0.2526 0.21803 -0.47766 -0.13033 -0.36929 -0.88388 -0.1249 0.27972 0.017521 0.19048 0.38647 -0.10236 0.26691 -0.66637 -0.66046 -0.48598 -0.5029 0.59602 -0.23975 -0.054244 0.71177 0.097479 0.18964 0.60496 -0.2421 1.261 0.5195 0.12978 0.28374 0.1499 -0.073072 -0.064345 0.041775 0.20712 -0.13972 0.021692 -0.45101 -0.077633 -0.58888 -0.0062811 0.50587 0.63067 -0.096216 -0.45549 -0.10162 -0.74026 -0.45125 0.16204 0.34589 0.2203 0.73482 -0.72055 0.019937 0.50934 -0.045864 -1.0167 0.4202 0.29336 0.057842 0.19622 0.71137 0.44455 -0.11329 -0.23249 0.3283 0.6458 -0.032498 0.58903 0.067438 -0.21519 0.24967 -0.047893 -0.12095 0.20468 -0.010392 -0.10827 0.5248 -0.013868 -0.40703 -0.2761 0.61498 -0.12118 -0.70097 -0.76415 -0.37243 0.3 -0.32852 -0.13877 0.23339 -0.58504 0.54768 -0.090521 0.30928 -0.19777 0.68883 0.043808 -0.012833 0.25696 0.017598 -0.11323 -0.76201 0.42972 -0.22032 -0.43818 -0.57085 0.23867 -0.098037 -0.4015 0.27659 -0.51578 -0.28637 -0.37785 0.83469 0.10563 1.1508 -0.67165 0.095388 -0.070545 0.039198 0.17726 0.44885 -0.045378 0.22337 -0.24957 0.93144 -0.16601 -0.095582 -0.60227 0.20068 -0.10264 -0.62696 0.048702 0.34737 -0.10634 -0.35068 0.11719 -0.79712 -0.32956 -0.60446 -0.0049038 -0.3351 -0.060065 -0.3063 -0.15462 -0.099521 -0.1788 0.098109 -0.59477 0.53245 -0.15388 0.063044 -0.47686 0.26712 -0.064799 0.2029 -0.093498 -0.44456 0.4692 -0.13718 0.035772 -0.74958 -0.51603 0.47025 -0.65103 0.027106 0.31463 -0.51519 -0.09912 -0.30605 0.2127 -1.6502 -0.34658 -0.19282 0.036578 -0.33871 0.21323 0.54172 -0.17543 -0.60187 -0.14679 0.20983 -0.084584 0.070885 -0.21752 -0.12642 0.030381 0.075461 0.86541 0.30098 0.22916 0.049217 -0.21204 0.32909 -0.021816
+snow -0.6961 -0.3339 -0.66542 -0.16459 -0.70283 0.053264 0.57508 1.1246 -0.41143 -0.93335 -0.397 -0.13949 -0.21725 0.49383 -0.16481 -0.43673 -0.39998 -0.14702 0.5828 0.73123 -0.16808 0.050093 0.20341 0.093283 -0.18944 -0.0092796 0.0064213 -0.5586 0.079708 0.034177 0.503 -0.084123 -0.15241 0.042398 -0.95865 0.13482 0.10695 0.22212 0.16383 0.081416 -0.61437 0.60299 0.53843 0.33915 -0.060046 -0.12329 0.30417 0.067838 -0.058329 -0.24791 -0.28177 0.32273 -0.12639 -0.40664 -0.42578 0.71366 0.18676 -0.49576 0.56635 0.39411 -0.11876 0.62798 0.50193 -0.38534 -0.32333 -0.29613 -0.1984 0.082042 -0.63666 -0.25177 0.070225 0.23886 -0.35341 -0.30615 -0.7898 -0.014515 -0.096662 0.27064 0.37095 -0.3916 0.15589 0.40176 -0.12316 -0.0069311 -0.17538 0.29317 -0.035662 -0.062503 -0.11821 -0.26708 0.33433 -0.41039 -0.44941 -0.058539 -0.5973 -0.060833 0.014623 0.031391 0.041093 0.21223 0.54304 0.51444 -0.2447 -0.034937 -0.61583 0.24116 0.93612 0.29663 -0.01733 0.39864 -0.399 -0.69927 0.010899 0.044804 0.096444 0.20555 0.37109 0.13219 0.29942 -0.28494 -0.071103 -0.45338 -0.22126 -0.31673 -0.10643 0.040453 -0.15324 0.33191 0.27801 -0.25143 -0.41784 1.1352 0.18709 0.57932 0.14912 0.42731 -0.81353 0.35546 0.10287 -0.10858 0.13692 0.11451 -0.68607 -0.17115 -0.52708 0.28953 0.5147 0.25549 -0.23139 -0.44275 0.42679 -0.41475 0.041182 -0.2664 0.60967 0.03783 0.27371 -0.5267 0.12029 0.5208 0.59519 -1.1315 0.19505 -0.2528 0.34636 0.82065 0.63271 0.091682 0.38433 -0.81108 0.18232 0.19068 -0.13031 0.21336 0.074454 -0.094498 0.47594 -0.31026 -0.11718 0.092891 0.22067 -0.16721 0.71703 0.30143 -0.40609 -0.16231 0.31315 -0.59325 -0.53404 -0.1087 -0.23026 0.36507 0.30648 -0.75576 -0.20767 -0.46966 -0.21035 0.0091924 0.5057 0.45564 0.84145 -0.19412 0.23964 0.85852 0.05229 -0.0011899 -0.29387 0.044187 -0.23886 0.19207 -0.0079459 -0.25773 0.31145 -0.47615 -0.00056431 -0.8941 -0.38667 -0.37907 0.52821 -0.45513 0.53567 0.13216 0.39741 -0.4904 0.24118 -0.11714 0.27007 0.15184 0.42316 -0.39708 0.13827 -0.27638 0.29908 -0.76008 0.061752 -0.4452 -0.5132 0.12124 0.15792 -0.57067 -0.68793 -0.33873 -0.43291 -0.46817 -0.84667 -0.65852 -0.59116 -0.043406 -0.013031 0.11246 -0.35374 0.3923 0.1172 -0.56268 0.83477 -0.34675 0.054568 -0.48494 0.12108 -0.15504 -0.047008 -0.2665 0.024593 0.70123 0.21284 -0.077796 0.050835 0.3865 0.37534 -0.48749 -0.013739 0.57852 -0.90425 -0.0062806 -0.28674 -0.017749 -1.0189 -0.71371 -0.36557 -0.73412 -0.027371 -0.071396 0.64792 -0.057281 -0.2512 0.039567 0.076976 0.34572 0.34606 -0.38323 -0.074011 -0.14153 -0.03109 0.53137 -0.35708 -0.28263 0.098663 0.17693 -0.39297 0.27708

utils/utils.py

@@ -0,0 +1,232 @@
+import numpy as np
+import torch
+import os
+from torch.autograd import Variable
+from skimage.metrics import peak_signal_noise_ratio as compare_psnr
+from skimage.metrics import mean_squared_error as compare_mse
+from skimage.metrics import structural_similarity as compare_ssim
+import pandas as pd
+
+from model.OneRestore import OneRestore
+from model.Embedder import Embedder
+
+def load_embedder_ckpt(device, freeze_model=False, ckpt_name=None,
+                                  combine_type = ['clear', 'low', 'haze', 'rain', 'snow',\
+                                            'low_haze', 'low_rain', 'low_snow', 'haze_rain',\
+                                                    'haze_snow', 'low_haze_rain', 'low_haze_snow']):
+    if ckpt_name != None:
+        if torch.cuda.is_available():
+            model_info = torch.load(ckpt_name)
+        else:
+            model_info = torch.load(ckpt_name, map_location=torch.device('cpu'))
+
+        print('==> loading existing Embedder model:', ckpt_name)
+        model = Embedder(combine_type)
+        model.load_state_dict(model_info)
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    else:
+        print('==> Initialize Embedder model.')
+        model = Embedder(combine_type)
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    if freeze_model:
+        freeze(model)
+
+    return model
+
+def load_restore_ckpt(device, freeze_model=False, ckpt_name=None):
+    if ckpt_name != None:
+        if torch.cuda.is_available():
+            model_info = torch.load(ckpt_name)
+        else:
+            model_info = torch.load(ckpt_name, map_location=torch.device('cpu'))
+        print('==> loading existing OneRestore model:', ckpt_name)
+        model = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+        model.load_state_dict(model_info)
+    else:
+        print('==> Initialize OneRestore model.')
+        model = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+        model = torch.nn.DataParallel(model).to("cuda" if torch.cuda.is_available() else "cpu")
+
+    if freeze_model:
+        freeze(model)
+    total = sum([param.nelement() for param in model.parameters()])
+    print("Number of OneRestore parameter: %.2fM" % (total/1e6))
+
+    return model
+
+def load_restore_ckpt_with_optim(device, local_rank=None, freeze_model=False, ckpt_name=None, lr=None):
+    if ckpt_name != None:
+        if torch.cuda.is_available():
+            model_info = torch.load(ckpt_name)
+        else:
+            model_info = torch.load(ckpt_name, map_location=torch.device('cpu'))
+
+        print('==> loading existing OneRestore model:', ckpt_name)
+        model = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+        optimizer = torch.optim.Adam(model.parameters(), lr=lr) if lr != None else None
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True) if local_rank != None else model
+
+        if local_rank != None:
+            model.load_state_dict(model_info['state_dict'])
+        else:
+            weights_dict = {}
+            for k, v in model_info['state_dict'].items():
+                new_k = k.replace('module.', '') if 'module' in k else k
+                weights_dict[new_k] = v
+            model.load_state_dict(weights_dict)
+        optimizer = torch.optim.Adam(model.parameters())
+        optimizer.load_state_dict(model_info['optimizer'])
+        cur_epoch = model_info['epoch']
+    else:
+        print('==> Initialize OneRestore model.')
+        model = OneRestore().to("cuda" if torch.cuda.is_available() else "cpu")
+        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True) if local_rank != None else torch.nn.DataParallel(model)
+        cur_epoch = 0
+
+    if freeze_model:
+        freeze(model)
+    total = sum([param.nelement() for param in model.parameters()])
+    print("Number of OneRestore parameter: %.2fM" % (total/1e6))
+
+    return model, optimizer, cur_epoch
+
+def load_embedder_ckpt_with_optim(device, args, combine_type = ['clear', 'low', 'haze', 'rain', 'snow',\
+    'low_haze', 'low_rain', 'low_snow', 'haze_rain', 'haze_snow', 'low_haze_rain', 'low_haze_snow']):
+    print('Init embedder')
+    # seed
+    if args.seed == -1:
+        args.seed = np.random.randint(1, 10000)
+    seed = args.seed
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    print('Training embedder seed:', seed)
+
+    # embedder model
+    embedder = Embedder(combine_type).to("cuda" if torch.cuda.is_available() else "cpu")
+
+    if args.pre_weight == '':
+        optimizer = torch.optim.Adam(embedder.parameters(), lr=args.lr)
+        cur_epoch = 1
+    else:
+        try:
+            embedder_info = torch.load(f'{args.check_dir}/{args.pre_weight}')
+            if torch.cuda.is_available():
+                embedder_info = torch.load(f'{args.check_dir}/{args.pre_weight}')
+            else:
+                embedder_info = torch.load(f'{args.check_dir}/{args.pre_weight}', map_location=torch.device('cpu'))
+            embedder.load_state_dict(embedder_info['state_dict'])
+            optimizer = torch.optim.Adam(embedder.parameters(), lr=args.lr)
+            optimizer.load_state_dict(embedder_info['optimizer'])
+            cur_epoch = embedder_info['epoch'] + 1
+        except:
+            print('Pre-trained model loading error!')
+    return embedder, optimizer, cur_epoch, device
+
+def freeze_text_embedder(m):
+    """Freezes module m.
+    """
+    m.eval()
+    for name, para in m.named_parameters():
+        if name == 'embedder.weight' or name == 'mlp.0.weight' or name == 'mlp.0.bias':
+            print(name)
+            para.requires_grad = False
+            para.grad = None
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+def data_process(data, args, device):
+    combine_type = args.degr_type
+    b,n,c,w,h = data.size()
+
+    pos_data = data[:,0,:,:,:]
+
+    inp_data = torch.zeros((b,c,w,h))
+    inp_class = []
+
+    neg_data = torch.zeros((b,n-2,c,w,h))
+
+    index = np.random.randint(1, n, (b))
+    for i in range(b):
+        k = 0
+        for j in range(n):
+            if j == 0:
+                continue
+            elif index[i] == j:
+                inp_class.append(combine_type[index[i]])
+                inp_data[i, :, :, :] = data[i, index[i], :, :,:]
+            else:
+                neg_data[i,k,:,:,:] = data[i, j, :, :,:]
+                k=k+1
+    return pos_data.to("cuda" if torch.cuda.is_available() else "cpu"), [inp_data.to("cuda" if torch.cuda.is_available() else "cpu"), inp_class], neg_data.to("cuda" if torch.cuda.is_available() else "cpu")
+
+def print_args(argspar):
+    print("\nParameter Print")
+    for p, v in zip(argspar.__dict__.keys(), argspar.__dict__.values()):
+        print('\t{}: {}'.format(p, v))
+    print('\n')
+
+def adjust_learning_rate(optimizer, epoch, lr_update_freq):
+    if not epoch % lr_update_freq and epoch:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = param_group['lr'] /2
+    return optimizer
+
+
+def tensor_metric(img, imclean, model, data_range=1):
+
+    img_cpu = img.data.cpu().numpy().astype(np.float32).transpose(0,2,3,1)
+    imgclean = imclean.data.cpu().numpy().astype(np.float32).transpose(0,2,3,1)
+    
+    SUM = 0
+    for i in range(img_cpu.shape[0]):
+        
+        if model == 'PSNR':
+            SUM += compare_psnr(imgclean[i, :, :, :], img_cpu[i, :, :, :],data_range=data_range)
+        elif model == 'MSE':
+            SUM += compare_mse(imgclean[i, :, :, :], img_cpu[i, :, :, :])
+        elif model == 'SSIM':
+            SUM += compare_ssim(imgclean[i, :, :, :], img_cpu[i, :, :, :], data_range=data_range, multichannel = True)
+            # due to the skimage vision problem, you can replace above line by
+            # SUM += compare_ssim(imgclean[i, :, :, :], img_cpu[i, :, :, :], data_range=data_range, channel_axis=-1)
+        else:
+            print('Model False!')
+        
+    return SUM/img_cpu.shape[0]
+
+def save_checkpoint(stateF, checkpoint, epoch, psnr_t1,ssim_t1,psnr_t2,ssim_t2, filename='model.tar'):
+    torch.save(stateF, checkpoint + 'OneRestore_model_%d_%.4f_%.4f_%.4f_%.4f.tar'%(epoch,psnr_t1,ssim_t1,psnr_t2,ssim_t2))
+
+def load_excel(x):
+    data1 = pd.DataFrame(x)
+
+    writer = pd.ExcelWriter('./mertic_result.xlsx')	
+    data1.to_excel(writer, 'PSNR-SSIM', float_format='%.5f')
+    # writer.save()
+    writer.close()
+
+def freeze(m):
+    """Freezes module m.
+    """
+    m.eval()
+    for p in m.parameters():
+        p.requires_grad = False
+        p.grad = None

utils/utils_data.py

@@ -0,0 +1,92 @@
+import torch, os
+from PIL import Image
+import numpy as np
+import torchvision.transforms as transforms
+import torch.utils.data as data
+from einops import rearrange
+
+class ImageLoader:
+    def __init__(self, root):
+        self.img_dir = root
+
+    def __call__(self, img):
+        file = f'{self.img_dir}/{img}'
+        img = Image.open(file).convert('RGB')
+        return img
+
+def imagenet_transform(phase):
+
+    if phase == 'train':
+        transform = transforms.Compose([
+            transforms.RandomResizedCrop(224),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+       
+    elif phase == 'test':
+        transform = transforms.Compose([
+            transforms.Resize([224,224]),
+            transforms.ToTensor()
+        ])
+
+    return transform
+
+class Dataset_embedding(data.Dataset):
+    def __init__(self, cfg_data, phase='train'):
+
+        self.transform = imagenet_transform(phase)
+        self.type_name = cfg_data.type_name
+        self.type2idx = {self.type_name[i]: i for i in range(len(self.type_name))}
+
+        if phase == 'train':
+            self.loader = ImageLoader(cfg_data.train_dir)
+            name = os.listdir(f'{cfg_data.train_dir}/{self.type_name[0]}')
+            self.data = []
+            for i in range(len(self.type_name)):
+                for j in range(len(name)):
+                    self.data.append([self.type_name[i], name[j]])
+        elif phase == 'test':
+            self.loader = ImageLoader(cfg_data.test_dir)
+            name = os.listdir(f'{cfg_data.test_dir}/{self.type_name[0]}')
+            self.data = []
+            for i in range(1, len(self.type_name)):
+                for j in range(len(name)):
+                    self.data.append([self.type_name[i], name[j]])
+        print(f'The amount of {phase} data is {len(self.data)}')
+
+    def __getitem__(self, index):
+
+        type_name, image_name = self.data[index]
+        scene = self.type2idx[type_name]
+        image = self.transform(self.loader(f'{type_name}/{image_name}'))
+
+        return (scene, image)
+
+    def __len__(self):
+        return len(self.data)
+
+def init_embedding_data(cfg_em,  phase):
+    if phase == 'train':
+        train_dataset = Dataset_embedding(cfg_em, 'train')
+        test_dataset = Dataset_embedding(cfg_em, 'test')
+        train_loader = data.DataLoader(train_dataset,
+                                    batch_size=cfg_em.batch,
+                                    shuffle=True, 
+                                    num_workers=cfg_em.num_workers,
+                                    pin_memory=True)
+        test_loader = data.DataLoader(test_dataset,
+                                    batch_size=cfg_em.batch,
+                                    shuffle=False, 
+                                    num_workers=cfg_em.num_workers,
+                                    pin_memory=True)
+        print(len(train_dataset),len(test_dataset))
+            
+    elif phase == 'inference':
+        test_dataset = Dataset_embedding(cfg_em, 'test')
+        test_loader = data.DataLoader(test_dataset, 
+                                    batch_size=1,
+                                    shuffle=False, 
+                                    num_workers=cfg_em.num_workers,
+                                    pin_memory=True)
+    
+    return train_loader, test_loader

utils/utils_word_embedding.py

@@ -0,0 +1,157 @@
+import torch
+import numpy as np
+import fasttext.util
+from gensim import models
+
+def load_word_embeddings(emb_file, vocab):
+    embeds = {}
+    for line in open(emb_file, 'rb'):
+        line = line.decode().strip().split(' ')
+        wvec = torch.FloatTensor(list(map(float, line[1:])))
+        embeds[line[0]] = wvec
+
+    # for zappos (should account for everything)
+    custom_map = {
+        'Faux.Fur':'fake_fur', 'Faux.Leather':'fake_leather', 'Full.grain.leather':'thick_leather', 
+        'Hair.Calf':'hair_leather', 'Patent.Leather':'shiny_leather', 'Nubuck':'grainy_leather', 
+        'Boots.Ankle':'ankle_boots', 'Boots.Knee.High':'knee_high_boots', 'Boots.Mid-Calf':'midcalf_boots', 
+        'Shoes.Boat.Shoes':'boat_shoes', 'Shoes.Clogs.and.Mules':'clogs_shoes', 'Shoes.Flats':'flats_shoes',
+        'Shoes.Heels':'heels', 'Shoes.Loafers':'loafers', 'Shoes.Oxfords':'oxford_shoes',
+        'Shoes.Sneakers.and.Athletic.Shoes':'sneakers'}
+    custom_map_vaw = {
+        'selfie': 'photo'
+    }
+
+    E = []
+    for k in vocab:
+        if k in custom_map:
+            print(f'Change {k} to {custom_map[k]}')
+            k = custom_map[k]
+        k = k.lower()
+        if '_' in k:
+            toks = k.split('_')
+            emb_tmp = torch.zeros(300).float()
+            for tok in toks:
+                if tok in custom_map_vaw:
+                    tok = custom_map_vaw[tok]
+                emb_tmp += embeds[tok]
+            emb_tmp /= len(toks)
+            E.append(emb_tmp)
+        else:
+            E.append(embeds[k])
+
+    embeds = torch.stack(E)
+    print ('Loaded embeddings from file %s' % emb_file, embeds.size())
+
+    return embeds
+
+def load_fasttext_embeddings(emb_file,vocab):
+    custom_map = {
+        'Faux.Fur': 'fake fur',
+        'Faux.Leather': 'fake leather',
+        'Full.grain.leather': 'thick leather',
+        'Hair.Calf': 'hairy leather',
+        'Patent.Leather': 'shiny leather',
+        'Boots.Ankle': 'ankle boots',
+        'Boots.Knee.High': 'kneehigh boots',
+        'Boots.Mid-Calf': 'midcalf boots',
+        'Shoes.Boat.Shoes': 'boatshoes',
+        'Shoes.Clogs.and.Mules': 'clogs shoes',
+        'Shoes.Flats': 'flats shoes',
+        'Shoes.Heels': 'heels',
+        'Shoes.Loafers': 'loafers',
+        'Shoes.Oxfords': 'oxford shoes',
+        'Shoes.Sneakers.and.Athletic.Shoes': 'sneakers',
+        'traffic_light': 'traficlight',
+        'trash_can': 'trashcan',
+        'dry-erase_board' : 'dry_erase_board',
+        'black_and_white' : 'black_white',
+        'eiffel_tower' : 'tower'
+    }
+    vocab_lower = [v.lower() for v in vocab]
+    vocab = []
+    for current in vocab_lower:
+        if current in custom_map:
+            vocab.append(custom_map[current])
+        else:
+            vocab.append(current)
+
+    
+    ft = fasttext.load_model(emb_file) #DATA_FOLDER+'/fast/cc.en.300.bin')
+    embeds = []
+    for k in vocab:
+        if '_' in k:
+            ks = k.split('_')
+            emb = np.stack([ft.get_word_vector(it) for it in ks]).mean(axis=0)
+        else:
+            emb = ft.get_word_vector(k)
+        embeds.append(emb)
+
+    embeds = torch.Tensor(np.stack(embeds))
+    print('Fasttext Embeddings loaded, total embeddings: {}'.format(embeds.size()))
+    return embeds
+
+def load_word2vec_embeddings(emb_file,vocab):
+    # vocab = [v.lower() for v in vocab]
+
+    
+    model = models.KeyedVectors.load_word2vec_format(emb_file,binary=True)
+        #DATA_FOLDER+'/w2v/GoogleNews-vectors-negative300.bin', binary=True)
+
+    custom_map = {
+        'Faux.Fur': 'fake_fur',
+        'Faux.Leather': 'fake_leather',
+        'Full.grain.leather': 'thick_leather',
+        'Hair.Calf': 'hair_leather',
+        'Patent.Leather': 'shiny_leather',
+        'Boots.Ankle': 'ankle_boots',
+        'Boots.Knee.High': 'knee_high_boots',
+        'Boots.Mid-Calf': 'midcalf_boots',
+        'Shoes.Boat.Shoes': 'boat_shoes',
+        'Shoes.Clogs.and.Mules': 'clogs_shoes',
+        'Shoes.Flats': 'flats_shoes',
+        'Shoes.Heels': 'heels',
+        'Shoes.Loafers': 'loafers',
+        'Shoes.Oxfords': 'oxford_shoes',
+        'Shoes.Sneakers.and.Athletic.Shoes': 'sneakers',
+        'traffic_light': 'traffic_light',
+        'trash_can': 'trashcan',
+        'dry-erase_board' : 'dry_erase_board',
+        'black_and_white' : 'black_white',
+        'eiffel_tower' : 'tower'
+    }
+
+    embeds = []
+    for k in vocab:
+        if k in custom_map:
+            k = custom_map[k]
+        if '_' in k and k not in model:
+            ks = k.split('_')
+            emb = np.stack([model[it] for it in ks]).mean(axis=0)
+        else:
+            emb = model[k]
+        embeds.append(emb)
+    embeds = torch.Tensor(np.stack(embeds))
+    print('Word2Vec Embeddings loaded, total embeddings: {}'.format(embeds.size()))
+    return embeds
+
+
+
+def initialize_wordembedding_matrix(name, vocab):
+    """
+    Args:
+    - name: hyphen separated word embedding names: 'glove-word2vec-conceptnet'.
+    - vocab: list of attributes/objects.
+    """
+    wordembs = name.split('+')
+    result = None
+
+    for wordemb in wordembs:
+        if wordemb == 'glove':
+            wordemb_ = load_word_embeddings(f'./utils/glove.6B.300d.txt', vocab)
+        if result is None:
+            result = wordemb_
+        else:
+            result = torch.cat((result, wordemb_), dim=1)
+    dim = 300 * len(wordembs)
+    return result, dim