Path-Restore: Learning Network Path Selection for Image Restoration

TPAMI 2021

Paper

Abstract

Very deep Convolutional Neural Networks (CNNs) have greatly improved the performance on various image restoration tasks. However, this comes at a price of increasing computational burden, hence limiting their practical usages. We observe that some corrupted image regions are inherently easier to restore than others since the distortion and content vary within an image. To leverage this, we propose Path-Restore, a multi-path CNN with a pathfinder that can dynamically select an appropriate route for each image region. We train the pathfinder using reinforcement learning with a difficulty-regulated reward. This reward is related to the performance, complexity and "the difficulty of restoring a region". A policy mask is further investigated to jointly process all the image regions. We conduct experiments on denoising and mixed restoration tasks. The results show that our method achieves comparable or superior performance to existing approaches with less computational cost. In particular, Path-Restore is effective for real-world denoising, where the noise distribution varies across different regions on a single image. Compared to the state-of-the-art RIDNet, our method achieves comparable performance and runs 2.7x faster on the realistic Darmstadt Noise Dataset.

Path-Restore handles different image regions (gradual Gaussian noise are added) using different network paths. [Top] Before and after restoration. [Bottom] Path-Restore automatically selects shorter network paths for 'easier' regions and longer paths for 'harder' regions.

The Framework of

Path-Restore

Path-Restore is composed of a multi-path CNN and a pathfinder. The multi-path CNN contains N dynamic blocks, each of which has M optional paths. The pathfinder is able to dynamically select paths for different image regions according to their contents and distortions.

The Architecture of

Path-Restore-Mask

We propose Path-Restore-Mask to process a high-resolution input at image level with higher efficiency. The pathfinder predicts a policy mask with multiple channels, which are then up-sampled by nearest-neighbor interpolation to select eligible regions for the corresponding paths.

Mixed

Distortions

Qualitative results of real-world denoising on the Darmstadt Noise Dataset. Path-Restore recovers clean results with sharp edges. With a deeper architecture and more training data, Path-Restore-Ext further mitigates the artifacts along sharp edges.

Gaussian

Denoising

Qualitative results of spatially variant Gaussian denoising. Path-Restore recovers textured regions better than the unified DnCNN since these regions are processed with long paths.

Mixed

Distortions

Qualitative results of addressing mixed distortions. Path-Restore is able to handle diverse distortions and produce more spatially consistent results compared to RL-Restore.

Network Path

Selections

Path-Restore handles different image regions using different network paths. [Left] Before and after restoration. [Right] Path-Restore automatically selects shorter network paths for 'easier' regions and longer paths for 'harder' regions.

Paper

Citation

@article{yu2021path,
 title={Path-Restore: Learning Network Path Selection for Image Restoration},
 author={Yu, Ke and Wang, Xintao and Dong, Chao and Tang, Xiaoou and Loy, Chen Change},
 journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
 year={2021},
 publisher={IEEE}
}

Related

Projects

  • Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning
    K. Yu, C. Dong, L. Lin, C. C. Loy
    in Proceedings of IEEE Conference on Computer Vision and Pattern Recognition, 2018 (CVPR)
    [arXiv] [Project Page]
  • Recovering Realistic Texture in Image Super-resolution by Deep Spatial Feature Transform
    X. Wang, K. Yu, C. Dong, C. C. Loy
    in Proceedings of IEEE Conference on Computer Vision and Pattern Recognition, 2018 (CVPR)
    [arXiv] [Project Page]

Contact


Ke Yu
Email: ericyu16 at hotmail.com