Improve model card for MaskDCPT with abstract, results, usage, and BasicSR library_name

README.md

@@ -1,13 +1,122 @@
 ---
 license: cc-by-nc-4.0
+pipeline_tag: image-to-image
 tags:
 - image-restoration
 - all-in-one image restoration
 - unified image restoration
 - universal image restoration
-pipeline_tag: image-to-image
+library_name: basicsr
 ---
 
-Paper: https://arxiv.org/abs/2510.13282
+# Universal Image Restoration Pre-training via Masked Degradation Classification (MaskDCPT)
+
+**Paper**: [Universal Image Restoration Pre-training via Masked Degradation Classification](https://arxiv.org/abs/2510.13282)
+**Code / Project Page**: [https://github.com/MILab-PKU/MaskDCPT](https://github.com/MILab-PKU/MaskDCPT)
+
+## Abstract
+This study introduces a Masked Degradation Classification Pre-Training method (MaskDCPT), designed to facilitate the classification of degradation types in input images, leading to comprehensive image restoration pre-training. Unlike conventional pre-training methods, MaskDCPT uses the degradation type of the image as an extremely weak supervision, while simultaneously leveraging the image reconstruction to enhance performance and robustness. MaskDCPT includes an encoder and two decoders: the encoder extracts features from the masked low-quality input image. The classification decoder uses these features to identify the degradation type, whereas the reconstruction decoder aims to reconstruct a corresponding high-quality image. This design allows the pre-training to benefit from both masked image modeling and contrastive learning, resulting in a generalized representation suited for restoration tasks. Benefit from the straightforward yet potent MaskDCPT, the pre-trained encoder can be used to address universal image restoration and achieve outstanding performance. Implementing MaskDCPT significantly improves performance for both convolution neural networks (CNNs) and Transformers, with a minimum increase in PSNR of 3.77 dB in the 5D all-in-one restoration task and a 34.8% reduction in PIQE compared to baseline in real-world degradation scenarios. It also emergences strong generalization to previously unseen degradation types and levels. In addition, we curate and release the UIR-2.5M dataset, which includes 2.5 million paired restoration samples across 19 degradation types and over 200 degradation levels, incorporating both synthetic and real-world data. The dataset, source code, and models are available at this https URL .
+
+## Model Overview
+MaskDCPT is a novel approach for universal image restoration. It employs a Masked Degradation Classification Pre-Training method (MaskDCPT) to classify degradation types in input images, coupled with image reconstruction to boost performance and robustness. The architecture features an encoder for feature extraction from masked low-quality inputs, a classification decoder to identify degradation types, and a reconstruction decoder to generate high-quality images. This design effectively combines masked image modeling and contrastive learning, resulting in a robust and generalized representation suitable for a wide array of restoration tasks.
+
+![The oversall pipeline of MaskDCPT.](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_pipeline.png)
+
+## Results
+MaskDCPT demonstrates outstanding performance across various image restoration tasks, achieving significant improvements in PSNR and PIQE. Below are some visual results from the paper:
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_figures.png)
+
+<details>
+<summary><strong>5D All-in-one Image Restoration</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_5d.png)
+
+</details>
+
+<details>
+<summary><strong>12D All-in-one Image Restoration</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_12d.png)
+
+</details>
+
+<details>
+<summary><strong>Out-of-distribution Image Restoration (Gaussian Denoising)</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_ooddn.png)
+
+</details>
+
+<details>
+<summary><strong>Real-world All-in-one Image Restoration</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_realworld.png)
+
+</details>
+
+<details>
+<summary><strong>Mixed Degradation Image Restoration</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_mixed.png)
+
+</details>
+
+## UIR-2.5M Dataset
+We introduce and release the UIR-2.5M dataset, the largest universal image restoration dataset to date. This dataset features:
+-   **2.5 million** image pairs.
+-   **19** degradation types.
+-   **more than 200** degradation levels, incorporating both synthetic and real-world data.
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_dataset.png)
+
+<details>
+<summary><strong>Details about the UIR-2.5M Dataset.</strong> (click to expand) </summary>
+
+![](https://github.com/MILab-PKU/MaskDCPT/raw/main/assets/maskdcpt_dataset_table.png)
+
+</details>
+
+## Quick Start / Sample Usage
+
+### Setup
+Clone the repository and set up the environment:
+
+```shell
+git clone https://github.com/MILab-PKU/MaskDCPT.git
+cd MaskDCPT
+```
+
+You can also create a new environment to avoid conflicts:
+
+```
+conda env create -f environment.yml
+conda activate maskdcpt
+```
+
+> !!! Remember to remove `basicsr` if it's already in your python environment to avoid conflicts !!!
+
+### Test / Inference
+Download [**pretrained models**](https://huggingface.co/Jiakui/MaskDCPT) and place them in the `./pretrained_models` directory.
+
+To reproduce the results or run inference, please modify the dataset path and model path in the configuration file, then execute:
+
+```shell
+python basicsr/test.py -opt options/all_in_one/test/test_NAFNet_5d.yml
+# You can modify the path of the config file after `-opt` for different tests.
+```
+
+## Citation
+If you find this repository useful, please consider giving a star ⭐ and citing our work:
 
-Inference code: https://github.com/MILab-PKU/MaskDCPT
+```bibtex
+@misc{hu2025universal,
+      title={Universal Image Restoration Pre-training via Masked Degradation Classification},
+      author={JiaKui Hu and Zhengjian Yao and Lujia Jin and Yinghao Chen and Yanye Lu},
+      year={2025},
+      eprint={2510.13282},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2510.13282},
+}
+```