:fire: RAM-Jittor

This is the official Jittor version of codes for the paper. For the PyTorch version, please refer to RAM.

Restore Anything with Masks：Leveraging Mask Image Modeling for Blind All-in-One Image Restoration
Chujie Qin, Ruiqi Wu, Zikun Liu, Xin Lin, Chunle Guo, Hyun Hee Park, Chongyi Li^* （ * indicates corresponding author)
European Conference on Computer Vision (ECCV), 2024

:rocket: Highlights:

• RAM is a Blind All-In-One Image Restoration framework that can simultaneously handle 7 Restoration Tasks and achieve SOTA performance !
• RAM focus on tackling how to extract Image Prior instead of degradation prior from diverse corrupted images by Leveraging Mask Image Modeling.

:wrench: Dependencies and Installation

1. Clone and enter our repository:

   git clone https://github.com/Dragonisss/RAM-Jittor.git RAM
   cd RAM-Jittor
   

2. Simply run the install.sh for installation!

   source install.sh
   

3. Activate the environment whenever you test!

   conda activate RAM_jit
   

:sparkles: Datasets and Pretrained Models

Given the number of datasets involved, we plan to offer a unified download link in the future to make it easier to access all datasets.

We combine datasets from various restoration tasks to form the training set. Here are the relevant links for all the datasets used:

Dataset	Phase	Source	Task for
OTS_ALPHA	Train	[Baidu Cloud(f1zz)]	Dehaze
Rain-13k	Train & Test	[Google Drive]	Derain
LOL-v2	Train & Test	[Real Subset Baidu Cloud(65ay)] / [Synthetic Subset Baidu Cloud(b14u)]	Low Light Enhancement
GoPro	Train & Test	[Download]	Motion Deblur
LSDIR	Train & Test	[HomePage]	Denoise DeJPEG DeBlur
SOTS	Test	[Download]	Denoise DeJPEG DeBlur
CBSD68	Test	[Download]	Denoise

You need to collect required datasets above and place them under the `./datasets` Directory.

Symbolic links is a recommended approach, allowing you to place the datasets anywhere you prefer!

The final directory structure will be arranged as:

datasets
    |- CBSD68
        |- CBSD68
          |- noisy5
          |- noisy10
          |- ...
    |- gopro
        |- test
        |- train
    |- LOL-v2
        |- Real_captured
        |- Synthetic
    |- LSDIR
        |- 0001000
        |- 0002000
        |- ...
    |- OTS_ALPHA
        |- clear
        |- depth
        |- haze
    |- LSDIR-val
        |- 0000001.png
        |- 0000002.png
        |- ...
    |- rain13k
        |- test
        |- train
    |- SOTS
        |- outdoor

Our pipeline can be applied to any image restoration network. We provide the pre-trained and fine-tuned model files for SwinIR and PromptIR mentioned in the paper:

Method	Phase	Framework	Download Links	Config File
RAM	Pretrain	SwinIR	[GoogleDrive]	[options/RAM_SwinIR/ram_swinir_pretrain.yaml]
RAM	Finetune	SwinIR	[GoogleDrive]	[options/RAM_SwinIR/ram_swinir_finetune.yaml]
RAM	Pretrain	PromptIR	[GoogleDrive]	[options/RAM_PromptIR/ram_promptir_pretrain.yaml]
RAM	Finetune	PromptIR	[GoogleDrive]	[options/RAM_PromptIR/ram_promptir_finetune.yaml]

:camera: Quick Demo

We provide scripts for inference your own images in inference.py.
You could run python inference.py --help to get detailed information of this scripts.

:robot: Training RAM With Jittor From Scratch!

Before proceeding, please ensure that the relevant datasets have been prepared as required. Besides, Multi-GPU training and inference in Jittor depends on OpenMPI. The user can install OpenMPI according to the Jittor Documentation.

1.Pretraining with MIM We use the collected datasets for model training. First, we execute the following command:

mpirun -np [num of gpus] python -m train \
-opt [OPT]

# e.g.
mpirun -np 8 python -m train \
-opt options/RAM_SwinIR/ram_swinir_pretrain.yaml

2.Mask Attribute Conductance Analysis

We use proposed Mask Attribute Conductance Analysis to analyze the importance of different layers for finetuning. You can run the following command to conduct MAC analysis:

python scripts/mac_analysis.py -opt [OPT]

# e.g.
python scripts/mac_analysis.py \
-opt options/RAM_SwinIR/ram_swinir_mac.yml

For convenience, we have provided the analysis results of the two models, RAM-SwinIR and RAM-PromptIR, mentioned in the paper. You can find them in ./mac_analysis_result/

3.Finetuning

mpirun -np [num of gpus] python -m train \
-opt [OPT]

# e.g.
mpirun -np 8 python -m train \
-opt options/RAM_SwinIR/ram_swinir_finetune.yaml

You can also add CUDA_DEVICE_VISIBLE= to choose gpu you want to use.

:chart_with_upwards_trend: Evaluation

We have provided a script for fast evaluation:

python ram/test.py \
-opt [OPT]

To benchmark the performance of RAM on the test dataset, you can run the following command:

# RAM-SwinIR
python test.py \
-opt options/test/ram_swinir_benchmark.yml

# RAM-PromptIR
python test.py \
-opt options/test/ram_promptir_benchmark.yml

:book: Citation

If you find our repo useful for your research, please consider citing our paper:

@inproceedings{qin2024restore,
  title={Restore Anything with Masks: Leveraging Mask Image Modeling for Blind All-in-One Image Restoration},
  author={Qin, Chu-Jie and Wu, Rui-Qi and Liu, Zikun and Lin, Xin and Guo, Chun-Le and Park, Hyun Hee and Li, Chongyi},
  booktitle={European Conference on Computer Vision},
  pages={364--380},
  year={2024},
  organization={Springer}
}

:postbox: Contact

For technical questions, please contact chujie.qin[AT]mail.nankai.edu.cn