A streamlined implementation of inference for the megvii-research/NAFNet model, focusing solely on straightforward prediction tasks without additional features like training or fine-tuning.
NAFNet Image Restoration
This repository contains implementations of the NAFNet (Nonlinear Activation Free Network) for image restoration tasks such as deblurring and denoising. The core of the project includes a modular framework for easy processing of images using pre-trained models.
Features
Image Restoration: The core functionality includes deblurring and denoising using state-of-the-art models.
Modular Design: Separate processors for different tasks, with a base class for common functionality.
Model Management: Automatically downloads and manages model files for easy usage.
Installation
Via PyPI
You can easily install the nafnetlib package from PyPI:
pip install nafnetlib
Requirements
Python 3.x
PyTorch
PIL (Pillow)
Other dependencies listed in requirements.txt
You can install the dependencies using pip:
pip install -r requirements.txt
These dependencies will be automatically installed when you install nafnetlib via pip.
Usage
Initialization and Image Processing
To use the deblurring or denoising functionality, you can initialize the corresponding processor class (DeblurProcessor or DenoiseProcessor), and then use the process() method to process images.
Here’s an example of how to use the DeblurProcessor and DenoiseProcessor:
python
1from PIL import Image
2from nafnetlib import DeblurProcessor, DenoiseProcessor
34# Initialize the deblurring processor5db_processor = DeblurProcessor(6 model_id="reds_width64",7 model_dir="/absolute/path/to/model/directory",8 device="cuda"9)1011# Print available models12print(db_processor.available_models())1314# Initialize the denoising processor15db_processor = DenoiseProcessor(16 model_id="sidd_width64",17 model_dir="/absolute/path/to/model/directory",18 device="cuda"19)2021# Load image for processing22img_path ="path/to/image/file"23image = Image.open(img_path)2425# Process the image (e.g., deblur or denoise)26db_processor.process(image).show()
Supported Models
The following models are available for different image restoration tasks:
Image Deblurring (GoPro):
Models: gopro_width64, gopro_width32
Image Denoising (SIDD):
Models: sidd_width64, sidd_width32
Image Deblurring with JPEG Artifacts (REDS):
Models: reds_width64
These models are derived from the megvii-research/nafnet GitHub repository.
Not-supported Models
Baseline-GoPro-*, Baseline-SIDD-* and NAFSSR-L_*
Model Download
When initializing the processors, models are automatically downloaded if they are not found in the specified directory.
Supported models come from megvii-research/nafnet.
Contributing
Feel free to fork the repository and submit issues or pull requests. Contributions are welcome!
License
This project is licensed under the MIT License - see the LICENSE file for details.
Citations
If NAFNet helps your research or work, please consider citing NAFNet.
@article{chen2022simple,
title={Simple Baselines for Image Restoration},
author={Chen, Liangyu and Chu, Xiaojie and Zhang, Xiangyu and Sun, Jian},
journal={arXiv preprint arXiv:2204.04676},
year={2022}
}
If NAFSSR helps your research or work, please consider citing NAFSSR.
@InProceedings{chu2022nafssr,
author = {Chu, Xiaojie and Chen, Liangyu and Yu, Wenqing},
title = {NAFSSR: Stereo Image Super-Resolution Using NAFNet},
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
month = {June},
year = {2022},
pages = {1239-1248}
}
