Source-linked AI summary
Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning
Ke Yu, Chao Dong, Liang Lin, Chen Change Loy
TL;DR
Image restoration must handle mixed and unknown distortions without always applying a large, fixed network. RL-Restore learns an agent to select a sequence of lightweight specialized CNN tools, jointly training the agent and tools. It achieves competitive restoration with less computational complexity, including 82.2% fewer computations than a single large CNN for the same performance.
Problem
Existing restoration methods commonly use specialized or potentially large single networks, while mixed distortions require sequential restoration and can create difficult intermediate states.
Method
RL-Restore formulates restoration as sequential decision-making, using an agent to select task-specific lightweight CNN tools and jointly training the agent and tools.
Results
82.2% fewer computations achieve the same performance as a single large CNN, while RL-Restore performs competitively against strong human-designed networks with less computational complexity.
Takeaways & Limitations
The dynamically formed toolchain can restore complex and mixed distortions while adapting tool use to the image and revealing restoration steps.
Takeaways & Limitations
Restoration is constrained by a maximum toolchain length, with T ≤ Tmax, to prevent excessive restoration.
Abstract
from arXiv · showhide
We investigate a novel approach for image restoration by reinforcement learning. Unlike existing studies that mostly train a single large network for a specialized task, we prepare a toolbox consisting of small-scale convolutional networks of different complexities and specialized in different tasks. Our method, RL-Restore, then learns a policy to select appropriate tools from the toolbox to progressively restore the quality of a corrupted image. We formulate a step-wise reward function proportional to how well the image is restored at each step to learn the action policy. We also devise a joint learning scheme to train the agent and tools for better performance in handling uncertainty. In comparison to conventional human-designed networks, RL-Restore is capable of restoring images corrupted with complex and unknown distortions in a more parameter-efficient manner using the dynamically formed toolchain.
1. Introduction
Existing restoration systems commonly rely on specialized or increasingly deep CNNs, whereas RL-Restore uses lightweight task-specific tools selected sequentially by an agent. The framework targets mixed distortions and achieves competitive restoration with substantially lower computation.
- Motivation: CNN restoration models are typically specialized for tasks or degradation settings, while handling multiple distortions often requires deeper networks that process every image identically.Some images are less difficult and could be restored more cheaply.
- RL-Restore: RL-Restore treats restoration as sequential decision-making, with an agent adaptively selecting task-specific lightweight CNN tools and optionally stopping when quality is satisfactory.The tools address assignments including deblurring, denoising, and JPEG artifact reduction.
- Contributions: RL-Restore jointly trains the agent and tools to improve handling of new and unknown artifacts that emerge during processing.
- Results: 82.2% fewer computations achieve the same performance as a single large CNN, while RL-Restore is superior to CNN approaches at similar complexity.The comparison uses a CNN baseline with similar parameter count.
2. Related Work
Prior image-restoration research largely uses single CNNs specialized for particular tasks, with some extensions addressing multiple degradations. RL-Restore instead learns a policy to select appropriate CNNs, aiming to improve parameter and computational efficiency while adapting to distortion type.
- CNN for Image Restoration: CNN methods have achieved strong performance across deblurring, denoising, JPEG artifact reduction, and super-resolution, but most train one network for the task at hand.
- Multiple Degradations: Several prior models handle multiple degradations using deeper or broader single networks, including multi-scale super-resolution and multiple restoration tasks.Examples include 20-layer CNNs and one-to-many architectures.
- Efficient Restoration: Network compression and recursive architectures reduce parameters or target computational efficiency, but recursive restoration can retain high computational cost because of many recursions.
- Deep Reinforcement Learning: RL-Restore is presented as an orthogonal alternative that saves parameters and computation by learning a policy to select appropriate CNNs rather than compressing an existing network.
3. Learning a Restoration Toolchain
RL-Restore formulates mixed-distortion image restoration as sequential tool selection, using specialized CNN tools and an agent that adaptively builds a toolchain. The framework addresses order sensitivity and intermediate-state artifacts through reinforcement learning and joint training.
- Problem Definition: Mixed distortions are modeled as sequential compositions, and restoration aims to recover a clear image by applying tools step by step.The distortion process may combine blur, noise, and JPEG compression, requiring sequential restoration rather than a single-task restorer.
- Challenges: Toolchain order and restoration level strongly affect final performance, making the agent’s tool choice at each step crucial.Rearranging tools or changing their levels can produce substantially different restoration results.
- Framework: RL-Restore treats tool selection as a Markov Decision Process and uses an agent to choose tools iteratively until selecting a stopping action.The state includes the current image and the previous action-value vector; the highest-valued action is selected at each step.
- Toolbox: The toolbox contains 12 lightweight, task-specific CNN tools covering Gaussian blur, Gaussian noise, and JPEG compression at different levels.Tools use either three-layer or eight-layer CNNs according to the targeted distortion.
- Reward: A stepwise reward equals the PSNR improvement between consecutive restoration states, so cumulative reward measures overall PSNR gain.The framework can also use other image-quality metrics as rewards, although those alternatives are outside this paper’s focus.
- Training: Joint end-to-end training refines the agent and tools so specialized restorers better handle artifacts introduced in intermediate states.The tools are jointly fine-tuned along selected toolchains using final restoration loss, while training data also receives slight noise and JPEG perturbations.
4. Experiments
Experiments evaluate RL-Restore on synthetic mixed distortions, real-world images, complexity, tool selection, and ablations. The results show competitive or superior restoration with lower complexity, adaptive toolchains, and benefits from stopping and training design choices.
- Datasets and Evaluation: The DIV2K experiments use mixed Gaussian blur, Gaussian noise, and JPEG compression at random levels, grouped into five distortion-severity categories.Two extreme cases that were too easy or too hard for restoration are discarded.
- Complexity: 0.19 million parameters make RL-Restore about one-third as complex as DnCNN and VDSR, including the agent and selected tools.The comparison uses 63 × 63 input images and also evaluates multiplication counts; VDSR-s is introduced as a similarly complex baseline.
- Synthetic Evaluation: On severe test images, RL-Restore surpasses DnCNN by 0.2 dB and VDSR by 0.3 dB, despite these distortions being absent from training.On mild and moderate sets, it is superior to VDSR-s and comparable to DnCNN and VDSR.
- Tool Selection: 47% of final-step actions are stopping actions; on unseen mild and severe sets, stopping occurs at 60% and 38%, respectively.Deblurring and denoising tools are preferred first, while denoising and de-JPEG tools are frequently selected second; the reported ratios indicate that severe distortions require longer toolchains.
- Real-World Evaluation: RL-Restore restores real-world smartphone images with mixed blur, noise, exposure, and compression artifacts without test-time fine-tuning.The stopping action is selected on several examples when the agent is confident in restored quality.
- Ablation Studies: Removing the stopping action drops PSNR by around 0.15 dB, with a larger gap on mild distortions because images are more easily over-restored.The stopping mechanism also prevents over-restoration and saves computation.
5. Conclusion
RL-Restore frames image restoration as reinforcement learning, using an agent to dynamically select a toolchain for progressively restoring images corrupted by complex and mixed distortions. Results on synthetic and real-world images validate its effectiveness and indicate extensibility to more challenging restoration tasks.
- RL-Restore dynamically selects a toolchain to progressively restore images corrupted by complex and mixed distortions.
- Extensive results on synthetic and real-world images validate the effectiveness of the proposed approach.
- The framework can be applied to more challenging restoration tasks or other low-level vision problems by developing powerful tools and an appropriate reward.