├── .gitignore ├── .gitmodules ├── .travis.yml ├── CHANGELOG.md ├── Data ├── README.md ├── credentials.json ├── crf.npz └── datasets.yaml ├── Docs ├── Benchmark (paper).md ├── Benchmark (reproduce).md ├── Collections CVPR 2019.md ├── Collections CVPR 2020.md ├── Collections ECCV 2020.md ├── Collections ICCV 2019.md ├── Collections NeurIPS 2020.md ├── CookBook │ ├── 1. Inference using SRCNN.ipynb │ ├── 2. Building data pipeline with DataLoader.ipynb │ └── 3. Calling model with executor.ipynb └── HowTo │ ├── Adding new model.md │ └── Change backend.md ├── LICENSE ├── README.md ├── Tests ├── correlation_test.py ├── data │ ├── crf_s.npz │ ├── fake_datasets.yml │ ├── flying_chair │ │ ├── flow │ │ │ └── 0000.flo │ │ └── pair │ │ │ └── 0000 │ │ │ ├── img0.png │ │ │ └── img1.png │ ├── kitti_car │ │ ├── c_10.png │ │ ├── c_11.png │ │ └── f_01.png │ ├── set5_x2 │ │ ├── img_001_SRF_2_LR.png │ │ ├── img_002_SRF_2_LR.png │ │ ├── img_003_SRF_2_LR.png │ │ ├── img_004_SRF_2_LR.png │ │ └── img_005_SRF_2_LR.png │ ├── vespcn.yaml │ └── video │ │ ├── custom_pair │ │ ├── hr │ │ │ └── xiuxian │ │ │ │ ├── xiuxian001.png │ │ │ │ ├── xiuxian002.png │ │ │ │ └── xiuxian003.png │ │ └── lr │ │ │ └── xiuxian │ │ │ ├── xiuxian001.png │ │ │ ├── xiuxian002.png │ │ │ └── xiuxian003.png │ │ └── raw_32x32.yv12 ├── dataset_test.py ├── googledrive_test.py ├── image_test.py ├── imfilter_test.py ├── initializer_test.py ├── loader_test.py ├── model_test.py ├── motion_test.py ├── space_to_depth_test.py ├── training_test.py ├── utility_test.py ├── vgg_test.py └── virtualfile_test.py ├── Tools ├── CelebA.py ├── DND.py ├── FFmpegHelper.py ├── FastMetrics.py ├── Image2Raw.py ├── MakeHDF.py ├── Misc.py ├── NtireHelper.py ├── Raw2Image.py ├── SeqVisual.py ├── Vimeo.py └── YoukuPackage.py ├── Train ├── README.md ├── check_dataset.py ├── eval.py ├── par │ ├── keras │ │ └── srcnn.yaml │ ├── pytorch │ │ ├── carn.yml │ │ ├── cubic.yml │ │ ├── dbpn.yml │ │ ├── dncnn.yml │ │ ├── drcn.yml │ │ ├── drn.yml │ │ ├── drrn.yml │ │ ├── edrn.yml │ │ ├── edsr.yml │ │ ├── espcn.yml │ │ ├── esrgan.yml │ │ ├── ffdnet.yml │ │ ├── frvsr.yml │ │ ├── mldn.yml │ │ ├── msrn.yml │ │ ├── ntire19.yml │ │ ├── qprn.yml │ │ ├── rbpn.yml │ │ ├── rcan.yml │ │ ├── realsr.yml │ │ ├── rsr.yml │ │ ├── sofvsr.yml │ │ ├── spmc.yml │ │ ├── srcnn.yml │ │ ├── srfeat.yml │ │ ├── srmd.yml │ │ ├── tecogan.yml │ │ ├── vdsr.yml │ │ └── vespcn.yml │ └── tensorflow │ │ ├── carn.yaml │ │ ├── dbpn.yaml │ │ ├── dcscn.yaml │ │ ├── dncnn.yaml │ │ ├── drcn.yaml │ │ ├── drrn.yaml │ │ ├── drsr.yaml │ │ ├── drsr_v2.yaml │ │ ├── duf.yaml │ │ ├── edsr.yaml │ │ ├── espcn.yaml │ │ ├── ffdnet.yaml │ │ ├── frvsr.yaml │ │ ├── gangp.yaml │ │ ├── idn.yaml │ │ ├── lapsrn.yaml │ │ ├── lsgan.yaml │ │ ├── memnet.yaml │ │ ├── msrn.yaml │ │ ├── nlrn.yaml │ │ ├── ragan.yaml │ │ ├── ragangp.yaml │ │ ├── ralsgan.yaml │ │ ├── rcan.yaml │ │ ├── rdn.yaml │ │ ├── rgan.yaml │ │ ├── rgangp.yaml │ │ ├── rlsgan.yaml │ │ ├── sgan.yaml │ │ ├── srcnn.yaml │ │ ├── srdensenet.yaml │ │ ├── srfeat.yaml │ │ ├── srgan.yaml │ │ ├── vdsr.yaml │ │ ├── vespcn.yaml │ │ ├── wgan.yaml │ │ └── wgangp.yaml └── train.py ├── VSR ├── Backend │ ├── Keras │ │ ├── Framework │ │ │ ├── Environment.py │ │ │ ├── Trainer.py │ │ │ └── __init__.py │ │ ├── Models │ │ │ ├── Model.py │ │ │ ├── Srcnn.py │ │ │ └── __init__.py │ │ └── __init__.py │ ├── TF │ │ ├── Arch │ │ │ ├── Dense.py │ │ │ ├── Discriminator.py │ │ │ ├── Residual.py │ │ │ └── __init__.py │ │ ├── Framework │ │ │ ├── GAN.py │ │ │ ├── LayersHelper.py │ │ │ ├── Motion.py │ │ │ ├── Noise.py │ │ │ ├── SuperResolution.py │ │ │ ├── Trainer.py │ │ │ └── __init__.py │ │ ├── Models │ │ │ ├── Carn.py │ │ │ ├── Crdn.py │ │ │ ├── Dbpn.py │ │ │ ├── Dcscn.py │ │ │ ├── DnCnn.py │ │ │ ├── Drcn.py │ │ │ ├── Drrn.py │ │ │ ├── Drsr.py │ │ │ ├── Drsr_v2.py │ │ │ ├── Duf.py │ │ │ ├── Edsr.py │ │ │ ├── Espcn.py │ │ │ ├── FFDNet.py │ │ │ ├── Gan.py │ │ │ ├── Idn.py │ │ │ ├── LapSrn.py │ │ │ ├── MemNet.py │ │ │ ├── Msrn.py │ │ │ ├── Nlrn.py │ │ │ ├── Rcan.py │ │ │ ├── Rdn.py │ │ │ ├── SRDenseNet.py │ │ │ ├── SRFeat.py │ │ │ ├── SrGan.py │ │ │ ├── Srcnn.py │ │ │ ├── Vdsr.py │ │ │ ├── Vespcn.py │ │ │ └── __init__.py │ │ ├── Util.py │ │ └── __init__.py │ ├── Torch │ │ ├── Framework │ │ │ ├── Environment.py │ │ │ ├── Summary.py │ │ │ ├── Trainer.py │ │ │ └── __init__.py │ │ ├── Models │ │ │ ├── Bicubic.py │ │ │ ├── Carn.py │ │ │ ├── Classic.py │ │ │ ├── Contrib │ │ │ │ ├── __init__.py │ │ │ │ ├── ntire19 │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── denoise.py │ │ │ │ │ ├── edrn.py │ │ │ │ │ ├── frn.py │ │ │ │ │ └── ran2.py │ │ │ │ └── ntire20 │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── xiaozhong │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── ops │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── discriminator.py │ │ │ │ │ ├── loss.py │ │ │ │ │ └── network.py │ │ │ ├── Crdn.py │ │ │ ├── Dbpn.py │ │ │ ├── Drn.py │ │ │ ├── Edsr.py │ │ │ ├── Esrgan.py │ │ │ ├── Ffdnet.py │ │ │ ├── Frvsr.py │ │ │ ├── Mldn.py │ │ │ ├── Model.py │ │ │ ├── Msrn.py │ │ │ ├── NTIRE19.py │ │ │ ├── NTIRE20.py │ │ │ ├── Ops │ │ │ │ ├── Blocks.py │ │ │ │ ├── Discriminator.py │ │ │ │ ├── Distortion.py │ │ │ │ ├── Initializer.py │ │ │ │ ├── Loss.py │ │ │ │ ├── Motion.py │ │ │ │ ├── Scale.py │ │ │ │ └── __init__.py │ │ │ ├── Optim │ │ │ │ └── SISR.py │ │ │ ├── Qprn.py │ │ │ ├── Rbpn.py │ │ │ ├── Rcan.py │ │ │ ├── SRFeat.py │ │ │ ├── Sofvsr.py │ │ │ ├── Spmc.py │ │ │ ├── Srmd.py │ │ │ ├── TecoGAN.py │ │ │ ├── Vespcn.py │ │ │ └── __init__.py │ │ ├── README.md │ │ ├── Util │ │ │ ├── Metrics.py │ │ │ ├── Utility.py │ │ │ └── __init__.py │ │ └── __init__.py │ └── __init__.py ├── DataLoader │ ├── Crop.py │ ├── Dataset.py │ ├── FloDecoder.py │ ├── Loader.py │ ├── NVDecoder.py │ ├── Transform.py │ ├── VirtualFile.py │ ├── YVDecoder.py │ └── __init__.py ├── Model │ └── __init__.py ├── Util │ ├── Config.py │ ├── Ensemble.py │ ├── GoogleDriveDownloader.py │ ├── Hook.py │ ├── ImageProcess.py │ ├── LearningRateScheduler.py │ ├── Math.py │ ├── PcaPrecompute.py │ ├── Utility.py │ ├── VisualizeOpticalFlow.py │ └── __init__.py └── __init__.py ├── azure-pipelines.yml ├── prepare_data.py └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Temp work dir 2 | seafile-ignore.txt 3 | .idea/ 4 | .vscode/ 5 | .pytest_cache/ 6 | .vsr/ 7 | __pycache__/ 8 | 9 | # Install info. 10 | *.egg-info/ 11 | 12 | # Intermediate output files 13 | Results/ 14 | Train/*.sh 15 | build/ 16 | dist/ 17 | 18 | # NDA 19 | Exp/ 20 | Exp 21 | 22 | # Images 23 | .png 24 | .jpg 25 | .jpeg 26 | .bmp 27 | .yuv 28 | .rgb 29 | .nv12 30 | .pyc 31 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/.gitmodules -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - "3.7" 4 | cache: pip 5 | 6 | install: 7 | - pip install -q tensorflow==1.15.0 google-api-python-client oauth2client 8 | - pip install -q torch==1.4.0+cpu torchvision==0.5.0+cpu -f https://download.pytorch.org/whl/torch_stable.html 9 | - pip install -U -q -e . 10 | script: 11 | - pytest -v Tests/dataset_test.py Tests/virtualfile_test.py --disable-warnings 12 | - pytest -v Tests/utility_test.py --disable-warnings 13 | - pytest -v Tests/image_test.py --disable-warnings 14 | - pytest -v Tests/initializer_test.py --disable-warnings 15 | - pytest -v Tests/loader_test.py --disable-warnings 16 | - pytest -v Tests/model_test.py --disable-warnings 17 | - pytest -v Tests/motion_test.py --disable-warnings 18 | - pytest -v Tests/vgg_test.py --disable-warnings 19 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | 1.0.6.2 2 | 3 | ## 1.0.6 4 | ## 2020-07 5 | - Update TF backend 6 | - Add support to tensorflow 2.0 (both legacy and eager mode) 7 | - Refactor torch backend models 8 | - Add `--caching_dataset` to cache transformed data into memory (ignored when `memory_limit` set). 9 | - Fix FastMetrics multi-threads issue 10 | - Fix loading issue when inferring with VSR models 11 | 12 | ## 1.0.5 13 | ## 2020-05 14 | - Fix bugs of DataLoader #108, #109 15 | 16 | ## 1.0.4 17 | ## 2020-04 18 | - Fix an error that dataloader may mess up the file order 19 | - Add a dataset checker to help verify the DDF 20 | - Fix a bug of memory_limit [issue](https://github.com/LoSealL/VideoSuperResolution/issues/102) 21 | 22 | ## 1.0.3 23 | ## 2020-02 24 | - Add --export to Train/eval.py 25 | 26 | ## 1.0.2 27 | ## 2020-02-16 28 | - Fix wrong data format in ImageProcess 29 | - Google Drive download api expired, fallback to a link message 30 | 31 | ## 1.0.1 32 | ## 2020-02-16 33 | - Add SRMD (CVPR 2018) by Kai Zhang et. al. 34 | - Upload to [PyPI](https://pypi.org/project/VSR/). 35 | 36 | ## 1.0.0 37 | ### 2020-01-31 38 | - Trace changelog since 1.0.0 39 | - Move VSRTorch package into VSR 40 | - Use `~/.vsr/config.yml` (`C:\Users\\.vsr\config.yml` for **windows**) to config the package 41 | - Provide common API for users to add into their own project easily 42 | - Refresh detailed [documents](./Docs) 43 | -------------------------------------------------------------------------------- /Data/README.md: -------------------------------------------------------------------------------- 1 | ## DDF 2 | Dataset Description Format 3 | 4 | ===changelog=== 5 | - 2018.10.17 Mitigate DDF from json to yaml 6 | - 2018.08.21 Add "Path_Tracked" section to alias test set 7 | 8 | #### Path 9 | In `Path` section, a file glob pattern is bound to a name, the pattern is from `pathlib.Path.glob`, you can find 10 | more information [here](https://docs.python.org/3/library/pathlib.html#pathlib.Path.glob). 11 | 12 | #### Path_Tracked 13 | `Path_Tracked` section is almost the same as `Path`. It is recently added to automatically add these paths to 14 | individual test dataset. Paths under this section is usually test-only data. 15 | 16 | #### Dataset 17 | In `Dataset` section, we bind a number of **Path** to a dataset name. 18 | In each dataset, you can specify a `train` set, a `val` set a `test` set and an `infer` set. Every set can point 19 | to a single path of a path array, or even a pure file glob pattern. 20 | For RAW data (such as RGB, YUV et.al.) you should explicitly specify `mode`, `width` and `height` (see `MCL-V` for example). 21 | 22 | ## Note 23 | [dataset.json](datasets.json) has been deprecated now, but is kept for compatibility. 24 | Please use [dataset.yaml](datasets.yaml) instead. 25 | 26 | ## Google API 27 | I use google drive api for downloading drive shared files. This is what `credentials.json` works for. 28 | -------------------------------------------------------------------------------- /Data/credentials.json: -------------------------------------------------------------------------------- 1 | { 2 | "installed": { 3 | "client_id": "478543842142-v29iqrpjg8oc54vabqbbtng19gto09b6.apps.googleusercontent.com", 4 | "project_id": "videosuperresolution", 5 | "auth_uri": "https://accounts.google.com/o/oauth2/auth", 6 | "token_uri": "https://www.googleapis.com/oauth2/v3/token", 7 | "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", 8 | "client_secret": "_PpBH6mb3lH1AjK5JGq4ds0K", 9 | "redirect_uris": [ 10 | "urn:ietf:wg:oauth:2.0:oob", 11 | "http://localhost" 12 | ] 13 | } 14 | } -------------------------------------------------------------------------------- /Data/crf.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Data/crf.npz -------------------------------------------------------------------------------- /Data/datasets.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018-2020 Wenyi Tang 2 | # VSR Dataset Description File 3 | # Date: Oct 17th 2018 4 | # Ver: v1.3 5 | 6 | --- 7 | # Add root dir to dataset. Take effect on all patterns below. 8 | Root: /mnt/data/datasets 9 | # Collect your dataset directory and name them! 10 | Path: 11 | 91-IMAGE: 91-image/ 12 | BSD100: BSD100_SR/image_SRF_4/*HR.* 13 | BSD500-Train: BSR_bsds500/BSR/BSDS500/data/images/train/*.jpg 14 | BSD500-Val: BSR_bsds500/BSR/BSDS500/data/images/val/*.jpg 15 | BSD500-Test: BSR_bsds500/BSR/BSDS500/data/images/test/*.jpg 16 | GOPRO-Train[video]: GOPRO_Large_all/train 17 | GOPRO-Val[video]: GOPRO_Large_all/test 18 | WATERLOO: exploration_database_and_code/pristine_images/ 19 | DIV2K-Train: DIV2K/DIV2K_train_HR/ 20 | DIV2K-Raw: DIV2K/DIV2K_train_LR/ 21 | DIV2K-Val: DIV2K/DIV2K_valid_HR/ 22 | SET5: Set5_SR/Set5/image_SRF_4/*HR.* 23 | SET14: Set14_SR/Set14/image_SRF_4/*HR.* 24 | URBAN100: Urban100_SR/image_SRF_4/*HR.* 25 | SUNHAY80: SunHays80_SR/image_SRF_8/*HR.* 26 | VID4[video]: vid4/original/ 27 | YOTRAIN-HR[video]: youku/train/hr/png 28 | YOTRAIN-LR[video]: youku/train/lr/png 29 | YOVAL-HR[video]: youku/val/hr/png 30 | YOVAL-LR[video]: youku/val/lr/png 31 | 32 | # bind datasets to a name, called in scripts 33 | Dataset: 34 | NONE: # empty set, do nothing 35 | train: [] 36 | val: [] 37 | test: [] 38 | 39 | # The training data is collected from list of `train`. 40 | # They are treated as the ground-truth HR images, and LR 41 | # counterparts are automatically generated using bicubic interpolation. 42 | BSD: # Combined BSD100 and BSD500 data 43 | train: [BSD100, BSD500-Train] # collected in array 44 | val: BSD500-Val # point as a single set 45 | test: [BSD500-Test] 46 | 47 | 91-IMAGE: # Yang's 91 images 48 | train: 91-IMAGE 49 | val: [SET5] 50 | test: [SET5, SET14] 51 | 52 | WATERLOO: # https://ece.uwaterloo.ca/~k29ma/exploration/ 53 | train: WATERLOO 54 | val: [SET5, SET14] 55 | test: [URBAN100, SUNHAY80] 56 | 57 | DIV2K: # NTIRE-2017 Challenge 58 | train: 59 | hr: DIV2K-Train 60 | lr: DIV2K-Raw 61 | val: [DIV2K-Val] 62 | 63 | DW2K: # Combined DIV2K & Waterloo 64 | train: [DIV2K-Train, WATERLOO, BSD500-Train] 65 | val: [DIV2K-Val] 66 | 67 | GOPRO[video]: # https://github.com/SeungjunNah/DeepDeblur_release 68 | train: [GOPRO-Train] 69 | val: [GOPRO-Val] 70 | test: [VID4] 71 | 72 | # If LR is pre-generated from HR or somewhere else, one can specify 73 | # customized LR data like this. 74 | YOUKU[video]: 75 | train: 76 | hr: YOTRAIN-HR 77 | lr: YOTRAIN-LR 78 | val: 79 | hr: YOVAL-HR 80 | lr: YOVAL-LR 81 | -------------------------------------------------------------------------------- /Docs/Benchmark (reproduce).md: -------------------------------------------------------------------------------- 1 | # Benchmark 2 | _NOTE: all results are evaluated from retrained model here._ 3 | 4 | ## SISR (Single-Image Super-Resolution)* 5 | 6 | |Model |Scale Factor | Set5 | Set14 | BSD100 | Urban100 | Framework | 7 | |:-------|:------------|:-----------|:-----------|:-----------|:-----------|:----------| 8 | |Bicubic | x2 |33.66/0.9299|30.24/0.8688|29.56/0.8431|26.88/0.8403| - | 9 | |--------|-------------|------------|------------|------------|------------|-----------| 10 | |Bicubic | x4 |28.42/0.8104|26.00/0.7027|25.96/0.6675|23.14/0.6577| - | 11 | |SRCNN | x4 |30.03/0.8514|26.94/0.7453|26.70/0.7054|24.08/0.7048|PyTorch | 12 | |ESPCN | x4 |29.01/0.8324|26.38/0.7348|26.34/0.7002|23.62/0.6872|PyTorch | 13 | |VDSR* | x4 |28.39/0.8108|25.76/0.7056|25.95/0.6698|23.12/0.6591|PyTorch | 14 | |DRCN | x4 |30.90/0.8748|27.59/0.7650|27.11/0.7204| -/- |PyTorch | 15 | |DRRN | x4 |31.07/0.8771|27.62/0.7663|27.15/0.7221|24.94/0.7438|PyTorch | 16 | |SRCNN | x4 |29.78/0.8441|26.83/0.7409|26.61/0.7014|23.94/0.6984|Tensorflow | 17 | |VDSR* | x4 |28.38/0.8104|25.84/0.7054|25.95/0.6694|23.12/0.6590|Tensorflow | 18 | 19 | * **Training hyper-params:** 20 | - Steps: 40k 21 | - Dataset: DIV2K 22 | - Batch: [16,96,96,3] (HR) 23 | * VDSR seems fail to reach its global minima. 24 | 25 | ## VSR (Video Super-Resolution) 26 | |Model |Scale| VID4 | TBD.. | TBD.. | TBD.. | Framework | 27 | |:-------|:----|:-----------|:-----------|:-----------|:-----------|:----------| 28 | |Bicubic |x4 |23.58/0.6363| | | | - | 29 | |Bayesian|x4 |26.04/0.8151| | | |Matlab | 30 | |VESPCN-1|x4 |25.46/0.7432| | | |Tensorflow | 31 | |VESPCN-3|x4 |25.04/0.7143| | | |PyTorch | 32 | |SPMC |x4 |25.65/0.7513| | | |PyTorch | 33 | |SOFVSR |x4 |26.04/0.7753| | | |PyTorch | 34 | |FRVSR |x4 |26.10/0.7755| | | |PyTorch | 35 | |RBPN |x4 |26.23/0.7843| | | |PyTorch | 36 | 37 | * Compare Y-channel only, no border pixel shaved off. 38 | 39 | ## Image Denoise 40 | |Model | AWGN Level | BSD68 | Set12 | Urban100 | 41 | |:-------|:-----------|:-----------|:-----------|:-----------| 42 | |BM3D | 15 |31.08/0.8722|32.37/0.8952|32.35/0.9220| 43 | |--------|------------|------------|------------|------------| 44 | |BM3D | 25 |28.57/0.8013|29.97/0.8504|29.70/0.8777| 45 | |--------|------------|------------|------------|------------| 46 | |BM3D | 50 |25.62/0.6864|26.72/0.7676|25.95/0.7791| 47 | 48 | ## Degradation Image Super-Resolution 49 | |Model | Scale | AWGN | Set5 | Set14 | Urban100 | BSD100 | 50 | |:-------|:------|:-----|:-----------|:-----------|:-----------|:-----------| 51 | |VDSR | x4 | 15 |26.11/0.7400|24.01/0.6293|22.11/0.6365|24.23/0.6037| 52 | |VDSR-B | x4 | 15 |25.73/0.7232|23.63/0.6090|21.82/0.6196|24.03/0.5915| 53 | |--------|-------|------|------------|------------|------------|------------| 54 | |VDSR | x4 | 25 |24.66/0.6912|22.93/0.5818|21.27/0.5915|23.35/0.5611| 55 | |VDSR-B | x4 | 25 |24.31/0.6729|22.50/0.5602|21.03/0.5742|23.13/0.5485| 56 | |--------|-------|------|------------|------------|------------|------------| 57 | |VDSR | x4 | 50 |22.27/0.5977|20.92/0.5005|19.74/0.5049|21.73/0.4907| 58 | |VDSR-B | x4 | 50 |22.02/0.5811|20.52/0.4792|19.53/0.4862|21.53/0.4783| 59 | 60 | 61 | ## TF v.s. Torch (x4) 62 | |TF/Torch | Set5 | Set14 | BSD100 | Urban100 | 63 | |:--------|:-----------|:-----------|:-----------|:-----------| 64 | |Baseline |28.42 |26.00 |25.96 |23.14 | 65 | 66 | **Winner:** TF(), Torch() 67 | -------------------------------------------------------------------------------- /Docs/Collections CVPR 2020.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Docs/Collections CVPR 2020.md -------------------------------------------------------------------------------- /Docs/Collections ECCV 2020.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Docs/Collections ECCV 2020.md -------------------------------------------------------------------------------- /Docs/Collections ICCV 2019.md: -------------------------------------------------------------------------------- 1 | # Collections Specially for ICCV 2019 2 | *No workshop publication included* 3 | 4 | Terms: 5 | - SR: super-resolution 6 | - SOTA: state-of-the-art 7 | - DSLR: digital single-lens reflex camera 8 | 9 | ## SISR 10 | - [Meta-SR](http://openaccess.thecvf.com/content_CVPR_2019/papers/Hu_Meta-SR_A_Magnification-Arbitrary_Network_for_Super-Resolution_CVPR_2019_paper.pdf) 11 | - Early published [arXiv](https://arxiv.org/abs/1903.00875) 12 | - Codes: [Pytorch](https://github.com/XuecaiHu/Meta-SR-Pytorch) 13 | - Remark: Face++, arbitrary scale factor 14 | 15 | ## VSR 16 | - [RBPN](http://openaccess.thecvf.com/content_CVPR_2019/papers/Haris_Recurrent_Back-Projection_Network_for_Video_Super-Resolution_CVPR_2019_paper.pdf) 17 | - Early published [arXiv](https://arxiv.org/abs/1903.10128) 18 | - Codes: [Pytorch](https://github.com/alterzero/RBPN-PyTorch) 19 | - Remark: **SOTA**, 7 frames, no explicit warping 20 | 21 | ## Denoising 22 | - [Noise2Void](http://openaccess.thecvf.com/content_CVPR_2019/papers/Krull_Noise2Void_-_Learning_Denoising_From_Single_Noisy_Images_CVPR_2019_paper.pdf) 23 | - Early published [arXiv](https://arxiv.org/abs/1811.10980) 24 | - Remark: A step further than Noise2Noise 25 | 26 | ## Deblurring 27 | - [Blind Image Deblurring with Local Maximum Gradient Prior](http://openaccess.thecvf.com/content_CVPR_2019/papers/Chen_Blind_Image_Deblurring_With_Local_Maximum_Gradient_Prior_CVPR_2019_paper.pdf) 28 | - [Dynamic Scene Deblurring with Parameter Selective Sharing and Nested Skip Connections](http://openaccess.thecvf.com/content_CVPR_2019/papers/Gao_Dynamic_Scene_Deblurring_With_Parameter_Selective_Sharing_and_Nested_Skip_CVPR_2019_paper.pdf) 29 | - [Deep Stacked Hierarchical Multi-patch Network for Image Deblurring](http://openaccess.thecvf.com/content_CVPR_2019/papers/Zhang_Deep_Stacked_Hierarchical_Multi-Patch_Network_for_Image_Deblurring_CVPR_2019_paper.pdf) 30 | - [Phase-only Image Based Kernel Estimation for Single Image Blind Deblurring](http://openaccess.thecvf.com/content_CVPR_2019/papers/Pan_Phase-Only_Image_Based_Kernel_Estimation_for_Single_Image_Blind_Deblurring_CVPR_2019_paper.pdf) 31 | - [Recurrent Neural Networks with Intra-Frame Iterations for Video Deblurring](http://openaccess.thecvf.com/content_CVPR_2019/papers/Nah_Recurrent_Neural_Networks_With_Intra-Frame_Iterations_for_Video_Deblurring_CVPR_2019_paper.pdf) 32 | - [A variational EM framework with adaptive edge selection for blind motion deblurring](http://openaccess.thecvf.com/content_CVPR_2019/papers/Yang_A_Variational_EM_Framework_With_Adaptive_Edge_Selection_for_Blind_CVPR_2019_paper.pdf) 33 | - [Unsupervised Domain-Specific Deblurring via Disentangled Representations](http://openaccess.thecvf.com/content_CVPR_2019/papers/Lu_Unsupervised_Domain-Specific_Deblurring_via_Disentangled_Representations_CVPR_2019_paper.pdf) 34 | - [DAVANet](http://openaccess.thecvf.com/content_CVPR_2019/papers/Zhou_DAVANet_Stereo_Deblurring_With_View_Aggregation_CVPR_2019_paper.pdf) 35 | 36 | ## Restoration 37 | - [DuRN](http://openaccess.thecvf.com/content_CVPR_2019/papers/Liu_Dual_Residual_Networks_Leveraging_the_Potential_of_Paired_Operations_for_CVPR_2019_paper.pdf) 38 | - Early published [arXiv](https://arxiv.org/abs/1903.08817) 39 | - Remark: Dual residual block, dehaze/denoise/deblur 40 | 41 | ## Todo 42 | ### Implement in VSR: 43 | - [ ] IKC 44 | -------------------------------------------------------------------------------- /Docs/Collections NeurIPS 2020.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Docs/Collections NeurIPS 2020.md -------------------------------------------------------------------------------- /Docs/HowTo/Adding new model.md: -------------------------------------------------------------------------------- 1 | # Add Models to VSR 2 | version 0.0.1-alpha 3 | 4 | ## Write new model with VSR 5 | 6 | - Create a new python file in VSR.Models. 7 | - Copy codes from `Edsr.py` and rename the class `EDSR`. 8 | - Write model graph in function `build_graph` 9 | - Write loss in function `build_loss` (You can also write in `build_graph` and ignore `build_loss`, that's your choice.) 10 | - (Optional) Write summaries and savers. (There're default summary and saver) 11 | 12 | ## Register model into VSR 13 | 14 | - Open `VSR.Models.__init__.py`. 15 | - Add an entry to `models` dict. The entry is `alias: (file-name, class-name)`. 16 | 17 | ---- 18 | 19 | ## Write new model with VSRTorch 20 | 21 | - Create a new python file in VSRTorch.Models. 22 | - Copy codes from `Espcn.py` and rename the class `ESPCN`. 23 | - Write modules. 24 | - Write forward and backward data-path in function `train`. 25 | - Write forward data-path in function `eval`. 26 | 27 | ## Register model into VSRTorch 28 | 29 | - Open `VSRTorch.Models.__init__.py`. 30 | - Add an entry to `models` dict. The entry is `alias: (file-name, class-name)`. -------------------------------------------------------------------------------- /Docs/HowTo/Change backend.md: -------------------------------------------------------------------------------- 1 | # How to change VSR backend 2 | The VSR uses PyTorch as the default backend. And VSR also supports tensorflow 3 | for some of models. 4 | 5 | Edit config file `~/.vsr/config.yml`, If you'd like to change to tensorflow: 6 | (create one if not exist) 7 | ```yaml 8 | # the backend could be 'tensorflow', 'keras', 'pytorch' 9 | # the `keras` represents tensorflow v2.0 10 | backend: tensorflow 11 | # the verbose could be 'error', 'warning', 'info', 'debug' 12 | verbose: info 13 | ``` 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Tang, Wenyi 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Tests/correlation_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import unittest 3 | 4 | if not os.getcwd().endswith('Tests'): 5 | os.chdir('Tests') 6 | 7 | import tensorflow as tf 8 | from VSR.Backend.TF.Util import _make_vector, _make_displacement, correlation 9 | 10 | 11 | class CorrelationTest(unittest.TestCase): 12 | @staticmethod 13 | def constant(): 14 | return tf.constant([[ 15 | [[1, 1.1, 1.2], [2, 2.1, 2.2], [3, 3.1, 3.2]], 16 | [[4, 4.1, 4.2], [5, 5.1, 5.2], [6, 6.1, 6.2]], 17 | [[7, 7.1, 7.2], [8, 8.1, 8.2], [9, 9.1, 9.2]] 18 | ]], 'float32') 19 | 20 | def test_correlation(self): 21 | with tf.Session() as sess: 22 | vec = _make_vector(self.constant()).eval() 23 | disp = _make_displacement(self.constant()).eval() 24 | x = self.constant() 25 | corr = correlation(x, x, 3, 1).eval() 26 | x = tf.ones([1, 5, 5, 1], 'float32') 27 | corr_stride = correlation(x, x, 3, 2, 2, 2).eval() 28 | self.assertEqual(vec.shape, [1, 3, 3, 27]) 29 | self.assertEqual(disp.shape, [1, 3, 3, 1]) 30 | self.assertEqual(corr.shape, [1, 3, 3, 4]) 31 | self.assertEqual(corr_stride.shape, [1, 3, 3, 4]) 32 | 33 | 34 | if __name__ == '__main__': 35 | unittest.main() 36 | -------------------------------------------------------------------------------- /Tests/data/crf_s.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/crf_s.npz -------------------------------------------------------------------------------- /Tests/data/fake_datasets.yml: -------------------------------------------------------------------------------- 1 | # Fake VSR Dataset Description File 2 | # Unit Test Only!! 3 | # Ver: v1.3.0 4 | --- 5 | Root: ./ 6 | 7 | Path: 8 | FOO: flying_chair/**/*.png 9 | BAR: set5_x2/*.png 10 | XIUXIAN[video]: video/custom_pair/lr 11 | 12 | Dataset: 13 | NONE: 14 | train: [] 15 | val: [] 16 | test: [] 17 | 18 | NORMAL: 19 | train: [FOO, BAR] 20 | val: [BAR] 21 | test: ["set5_x2/img_001_SRF_2_LR.png"] 22 | 23 | RAW[video]: 24 | train: [video/raw_32x32.yv12] 25 | val: [video/raw_32x32.yv12] 26 | test: [video/raw_32x32.yv12] 27 | 28 | PAIR: 29 | train: 30 | hr: [FOO] 31 | lr: [FOO] 32 | 33 | VIDEOPAIR[video]: 34 | train: 35 | hr: "video/custom_pair/hr" 36 | lr: XIUXIAN 37 | val: 38 | hr: [XIUXIAN] 39 | lr: "video/custom_pair/hr" 40 | test: 41 | hr: "video/custom_pair/hr" 42 | lr: "video/custom_pair/lr" 43 | 44 | GHOST: 45 | train: ["/what-the-fuck"] 46 | -------------------------------------------------------------------------------- /Tests/data/flying_chair/flow/0000.flo: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/flying_chair/flow/0000.flo -------------------------------------------------------------------------------- /Tests/data/flying_chair/pair/0000/img0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/flying_chair/pair/0000/img0.png -------------------------------------------------------------------------------- /Tests/data/flying_chair/pair/0000/img1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/flying_chair/pair/0000/img1.png -------------------------------------------------------------------------------- /Tests/data/kitti_car/c_10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/kitti_car/c_10.png -------------------------------------------------------------------------------- /Tests/data/kitti_car/c_11.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/kitti_car/c_11.png -------------------------------------------------------------------------------- /Tests/data/kitti_car/f_01.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/kitti_car/f_01.png -------------------------------------------------------------------------------- /Tests/data/set5_x2/img_001_SRF_2_LR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/set5_x2/img_001_SRF_2_LR.png -------------------------------------------------------------------------------- /Tests/data/set5_x2/img_002_SRF_2_LR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/set5_x2/img_002_SRF_2_LR.png -------------------------------------------------------------------------------- /Tests/data/set5_x2/img_003_SRF_2_LR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/set5_x2/img_003_SRF_2_LR.png -------------------------------------------------------------------------------- /Tests/data/set5_x2/img_004_SRF_2_LR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/set5_x2/img_004_SRF_2_LR.png -------------------------------------------------------------------------------- /Tests/data/set5_x2/img_005_SRF_2_LR.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/set5_x2/img_005_SRF_2_LR.png -------------------------------------------------------------------------------- /Tests/data/vespcn.yaml: -------------------------------------------------------------------------------- 1 | # VESPCN need data depth to be 3, 5, or 7. (2n + 1) 2 | --- 3 | vespcn: 4 | depth: 3 # this `depth` is to inform graph builder 5 | beta: 1 6 | gamma: 0.01 7 | scale: 4 8 | channel: 3 9 | weight_decay: 0 10 | batch: 4 11 | patch_size: 32 12 | depth: 3 # must be same as the model depth 13 | lr: 1.0e-4 14 | lr_decay: 15 | method: multistep 16 | decay_step: [] 17 | decay_rate: 1 18 | -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/hr/xiuxian/xiuxian001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/hr/xiuxian/xiuxian001.png -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/hr/xiuxian/xiuxian002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/hr/xiuxian/xiuxian002.png -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/hr/xiuxian/xiuxian003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/hr/xiuxian/xiuxian003.png -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/lr/xiuxian/xiuxian001.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/lr/xiuxian/xiuxian001.png -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/lr/xiuxian/xiuxian002.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/lr/xiuxian/xiuxian002.png -------------------------------------------------------------------------------- /Tests/data/video/custom_pair/lr/xiuxian/xiuxian003.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/custom_pair/lr/xiuxian/xiuxian003.png -------------------------------------------------------------------------------- /Tests/data/video/raw_32x32.yv12: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/Tests/data/video/raw_32x32.yv12 -------------------------------------------------------------------------------- /Tests/dataset_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 16 5 | 6 | import os 7 | import unittest 8 | 9 | if not os.getcwd().endswith('Tests'): 10 | os.chdir('Tests') 11 | from VSR.DataLoader.Dataset import Dataset, load_datasets 12 | 13 | 14 | class DatasetTest(unittest.TestCase): 15 | def test_image_data(self): 16 | d = Dataset('data/set5_x2') 17 | data = d.compile() 18 | self.assertEqual(len(data), 5) 19 | self.assertEqual(data.capacity, 983040) 20 | 21 | def test_video_data(self): 22 | d = Dataset('data/video/custom_pair').use_like_video() 23 | data = d.compile() 24 | self.assertEqual(len(data), 2) 25 | 26 | def test_multi_url(self): 27 | d = Dataset('data/set5_x2', 'data/kitti_car') 28 | data = d.compile() 29 | self.assertEqual(len(data), 8) 30 | 31 | def test_include_exclude(self): 32 | d = Dataset('data') 33 | d.include_('xiuxian*') 34 | data1 = d.compile() 35 | d = Dataset('data') 36 | d.include_reg_('set5') 37 | data2 = d.compile() 38 | d = Dataset('data').include_reg('set5').exclude('png') 39 | data3 = d.compile() 40 | 41 | self.assertEqual(len(data1), 6) 42 | self.assertEqual(len(data2), 5) 43 | self.assertEqual(len(data3), 0) 44 | 45 | def test_dataset_desc_file(self): 46 | ddf = 'data/fake_datasets.yml' 47 | datasets = load_datasets(ddf) 48 | self.assertEqual(len(datasets), 9) 49 | self.assertEqual(len(datasets.NONE.train.hr.compile()), 0) 50 | self.assertEqual(len(datasets.NORMAL.train.hr.compile()), 7) 51 | self.assertEqual(len(datasets.NORMAL.val.hr.compile()), 5) 52 | self.assertEqual(len(datasets.NORMAL.test.hr.compile()), 1) 53 | self.assertEqual(len(datasets.PAIR.train.hr.compile()), 2) 54 | self.assertEqual(len(datasets.PAIR.train.lr.compile()), 2) 55 | self.assertEqual(len(datasets.VIDEOPAIR.train.hr.compile()), 1) 56 | self.assertEqual(len(datasets.VIDEOPAIR.train.lr.compile()), 1) 57 | self.assertEqual(len(datasets.VIDEOPAIR.val.hr.compile()), 1) 58 | self.assertEqual(len(datasets.VIDEOPAIR.val.lr.compile()), 1) 59 | self.assertEqual(len(datasets.VIDEOPAIR.test.hr.compile()), 1) 60 | self.assertEqual(len(datasets.VIDEOPAIR.test.lr.compile()), 1) 61 | self.assertEqual(len(datasets.FOO.test.hr.compile()), 2) 62 | self.assertEqual(len(datasets.BAR.test.hr.compile()), 5) 63 | self.assertTrue(datasets.VIDEOPAIR.train.hr.as_video) 64 | self.assertTrue(datasets.XIUXIAN.test.hr.as_video) 65 | 66 | raw = load_datasets(ddf, 'RAW') 67 | self.assertEqual(len(raw.train.hr.compile()), 1) 68 | self.assertEqual(len(raw.val.hr.compile()), 1) 69 | self.assertEqual(len(raw.test.hr.compile()), 1) 70 | self.assertTrue(raw.train.hr.as_video) 71 | 72 | 73 | if __name__ == '__main__': 74 | unittest.main() 75 | -------------------------------------------------------------------------------- /Tests/googledrive_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 16 5 | 6 | import os 7 | import unittest 8 | 9 | if not os.getcwd().endswith('Tests'): 10 | os.chdir('Tests') 11 | 12 | try: 13 | from googleapiclient.discovery import build 14 | from httplib2 import Http 15 | from oauth2client import file, client, tools 16 | except ImportError as ex: 17 | print(f"[!] {ex}") 18 | exit(-1) 19 | 20 | # If modifying these scopes, delete the file token.json. 21 | SCOPES = 'https://www.googleapis.com/auth/drive.readonly' 22 | 23 | 24 | class FetchGoogleDriveTest(unittest.TestCase): 25 | def test_downloads(self): 26 | """Shows basic usage of the Drive v3 API. 27 | Prints the names and ids of the first 10 files the user has access to. 28 | """ 29 | # The file token.json stores the user's access and refresh tokens, and is 30 | # created automatically when the authorization flow completes for the first 31 | # time. 32 | store = file.Storage('/tmp/token.json') 33 | creds = store.get() 34 | if not creds or creds.invalid: 35 | flow = client.flow_from_clientsecrets('../Data/credentials.json', SCOPES) 36 | creds = tools.run_flow(flow, store) 37 | service = build('drive', 'v3', http=creds.authorize(Http())) 38 | file_id = '1H0PIXvJH4c40pk7ou6nAwoxuR4Qh_Sa2' 39 | request = service.files().get_media(fileId=file_id) 40 | request.execute() 41 | 42 | 43 | if __name__ == '__main__': 44 | unittest.main() 45 | -------------------------------------------------------------------------------- /Tests/imfilter_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 13 5 | 6 | import unittest 7 | 8 | import numpy as np 9 | 10 | from VSR.Util.Math import (anisotropic_gaussian_kernel, gaussian_kernel) 11 | from VSR.Backend import BACKEND 12 | 13 | _K1 = gaussian_kernel(15, 2) 14 | _K2 = anisotropic_gaussian_kernel(15, 1, 5, 3) 15 | 16 | 17 | class ImFilter(unittest.TestCase): 18 | # For ones([4, 4]) 19 | y_gold = np.array([ 20 | [0.3151, 0.3776, 0.3776, 0.3151], 21 | [0.3776, 0.4524, 0.4524, 0.3776], 22 | [0.3776, 0.4524, 0.4524, 0.3776], 23 | [0.3151, 0.3776, 0.3776, 0.3151] 24 | ]) 25 | z_gold = np.array([ 26 | [0.3391, 0.3950, 0.3774, 0.2950], 27 | [0.3850, 0.4627, 0.4557, 0.3677], 28 | [0.3677, 0.4557, 0.4627, 0.3850], 29 | [0.2950, 0.3774, 0.3950, 0.3391] 30 | ]) 31 | 32 | def test_torch(self): 33 | if BACKEND != 'pytorch': 34 | return 35 | import torch 36 | from VSR.Backend.Torch.Util.Utility import imfilter 37 | 38 | tk1 = torch.tensor(_K1, dtype=torch.float32) 39 | tk2 = torch.tensor(_K2, dtype=torch.float32) 40 | x = torch.ones(2, 3, 4, 4, dtype=torch.float32) 41 | y = imfilter(x, tk1) 42 | z = imfilter(x, torch.stack([tk1, tk2])) 43 | y_ = y.detach().numpy() 44 | z_ = z.detach().numpy() 45 | self.assertTrue(np.all(y_[0] == z_[0])) 46 | self.assertTrue(np.all(np.abs(y_[0, 0] - self.y_gold) <= 1e-4)) 47 | self.assertTrue(np.all(np.abs(z_[1, 0] - self.z_gold) <= 1e-4)) 48 | 49 | def test_tf(self): 50 | if BACKEND != 'tensorflow': 51 | return 52 | import tensorflow as tf 53 | from VSR.Backend.TF.Util import imfilter 54 | 55 | tk1 = tf.constant(_K1, dtype=tf.float32) 56 | tk2 = tf.constant(_K2, dtype=tf.float32) 57 | x = tf.ones([2, 4, 4, 3], dtype=tf.float32) 58 | y = imfilter(x, tk1) 59 | z = imfilter(x, tk2) 60 | with tf.Session() as sess: 61 | y_, z_ = sess.run([y, z]) 62 | self.assertTrue(np.all(np.abs(y_[0, ..., 0] - self.y_gold) <= 1e-4)) 63 | self.assertTrue(np.all(np.abs(z_[0, ..., 0] - self.z_gold) <= 1e-4)) 64 | 65 | 66 | if __name__ == '__main__': 67 | unittest.main() 68 | -------------------------------------------------------------------------------- /Tests/initializer_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019 - 3 - 5 5 | 6 | import os 7 | import unittest 8 | 9 | if not os.getcwd().endswith('Tests'): 10 | os.chdir('Tests') 11 | import numpy as np 12 | 13 | try: 14 | from torch import nn 15 | import tensorflow as tf 16 | 17 | tf.enable_eager_execution() 18 | except (ImportError, AttributeError) as ex: 19 | print(f"[!] Cross validation is needed") 20 | print(ex) 21 | exit(0) 22 | 23 | from VSR.Backend.TF.Util import TorchInitializer 24 | 25 | 26 | class TestInitializer(unittest.TestCase): 27 | def test_torch_initializer(self): 28 | x = np.ones([4, 16, 16, 16], np.float32) 29 | c2dtf = tf.layers.Conv2D(16, 3, padding='same', 30 | kernel_initializer=TorchInitializer(), 31 | bias_initializer=TorchInitializer(9 * 16)) 32 | c2dtf.build(x.shape) 33 | w1 = c2dtf.kernel 34 | y1 = c2dtf.apply(x) 35 | 36 | c2dnn = nn.Conv2d(16, 16, 3, padding=1) 37 | w2 = c2dnn.weight 38 | # TODO: how to test distribution? 39 | assert True 40 | -------------------------------------------------------------------------------- /Tests/model_test.py: -------------------------------------------------------------------------------- 1 | # ############################################################################## 2 | # Copyright (c) 2020. LoSealL All Rights Reserved. 3 | # Author: Wenyi Tang 4 | # Email: wenyitang@outlook.com 5 | # Date: 2020 - 2 - 5 6 | # ############################################################################## 7 | import os 8 | import unittest 9 | 10 | if not os.getcwd().endswith('Tests'): 11 | os.chdir('Tests') 12 | from VSR.DataLoader import Dataset, Loader 13 | from VSR.Model import get_model 14 | from VSR.Backend import DATA_FORMAT 15 | 16 | 17 | class ModelTest(unittest.TestCase): 18 | def test_train_srcnn(self): 19 | data = Dataset('data').include_reg('set5') 20 | ld = Loader(data, scale=2) 21 | ld.set_color_space('lr', 'L') 22 | ld.set_color_space('hr', 'L') 23 | m = get_model('srcnn')(scale=2, channel=1) 24 | with m.executor as t: 25 | config = t.query_config({}) 26 | config.epochs = 5 27 | config.steps = 10 28 | if DATA_FORMAT == 'channels_first': 29 | config.batch_shape = [16, 1, 16, 16] 30 | else: 31 | config.batch_shape = [16, 16, 16, 1] 32 | t.fit([ld, None], config) 33 | 34 | def test_infer_srcnn(self): 35 | m = get_model('srcnn')(scale=2, channel=3) 36 | data = Dataset('data').include_reg('set5') 37 | ld = Loader(data, scale=2) 38 | with m.executor as t: 39 | config = t.query_config({}) 40 | t.infer(ld, config) 41 | 42 | def test_train_vespcn(self): 43 | data = Dataset('data/video').include_reg("xiuxian").use_like_video() 44 | ld = Loader(data, scale=2) 45 | m = get_model('vespcn')(scale=2, channel=3) 46 | with m.executor as t: 47 | config = t.query_config({}) 48 | config.epochs = 1 49 | config.steps = 10 50 | if DATA_FORMAT == 'channels_first': 51 | config.batch_shape = [16, 3, 3, 16, 16] 52 | else: 53 | config.batch_shape = [16, 3, 16, 16, 3] 54 | t.fit([ld, None], config) 55 | 56 | 57 | if __name__ == '__main__': 58 | unittest.main() 59 | -------------------------------------------------------------------------------- /Tests/space_to_depth_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 16 5 | 6 | import os 7 | import unittest 8 | 9 | if not os.getcwd().endswith('Tests'): 10 | os.chdir('Tests') 11 | 12 | try: 13 | import torch 14 | import torchvision 15 | from torch.nn import PixelShuffle 16 | from VSR.Backend.Torch.Models.Ops.Scale import SpaceToDim 17 | except ImportError: 18 | exit(0) 19 | 20 | import numpy as np 21 | from PIL import Image 22 | 23 | 24 | class SpaceToDimTest(unittest.TestCase): 25 | def test_space_to_depth(self): 26 | f1 = SpaceToDim(2, dim=1) 27 | ff = PixelShuffle(2) 28 | x = Image.open('data/set5_x2/img_001_SRF_2_LR.png') 29 | g = torchvision.transforms.ToTensor() 30 | h = torchvision.transforms.ToPILImage() 31 | z = f1(g(x).unsqueeze(0)) 32 | y = h(ff(z)[0]) 33 | self.assertTrue(np.all(np.array(x) == np.array(y))) 34 | 35 | def dummy_test_space_to_x(self): 36 | f1 = SpaceToDim(2, (1, 2), dim=3) 37 | x = torch.ones(1, 4, 4, 3) 38 | y = f1(x) 39 | self.assertEqual(y.shape, torch.Size([1, 2, 2, 12])) 40 | f2 = SpaceToDim(2, (1, 2), dim=0) 41 | y = f2(x) 42 | self.assertEqual(y.shape, torch.Size([4, 2, 2, 3])) 43 | 44 | 45 | if __name__ == '__main__': 46 | unittest.main() 47 | -------------------------------------------------------------------------------- /Tests/training_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/17 下午5:53 5 | 6 | import os 7 | import subprocess 8 | from VSR.Model import list_supported_models 9 | 10 | if not os.getcwd().endswith('Tests'): 11 | os.chdir('Tests') 12 | 13 | _WORKDIR = r"/tmp/vsr/utest/" 14 | _TCMD = ("python train.py {} --data_config=../Tests/data/fake_datasets.yml" 15 | "--dataset=normal --epochs=1 --steps=1 --save_dir={} --val_steps=1") 16 | _ECMD = r"python eval.py {} --save_dir={} --ensemble -t=../Tests/data/set5_x2" 17 | 18 | 19 | def train(model_name: str): 20 | cmd = _TCMD.format(str(model_name), _WORKDIR) 21 | cwd = '../Train' 22 | subprocess.call(cmd, stderr=subprocess.DEVNULL, cwd=cwd, shell=True) 23 | 24 | 25 | def eval(model_name: str): 26 | cmd = _ECMD.format(str(model_name), _WORKDIR) 27 | cwd = '../Train' 28 | subprocess.call(cmd, stderr=subprocess.DEVNULL, cwd=cwd, shell=True) 29 | 30 | 31 | def test_train_srcnn(): 32 | train('srcnn') 33 | eval('srcnn') 34 | 35 | 36 | def test_train_espcn(): 37 | train('espcn') 38 | eval('espcn') 39 | 40 | 41 | def test_other_models(): 42 | for k in list_supported_models(): 43 | if k in ( 44 | 'sofvsr', 'vespcn', 'frvsr', 'qprn', 'ufvsr', 'yovsr', 'tecogan', 45 | 'spmc', 'rbpn' 46 | ): 47 | # skip video model 48 | continue 49 | train(k) 50 | eval(k) 51 | 52 | 53 | if __name__ == '__main__': 54 | test_other_models() 55 | -------------------------------------------------------------------------------- /Tests/utility_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import unittest 3 | 4 | if not os.getcwd().endswith('Tests'): 5 | os.chdir('Tests') 6 | 7 | from VSR.Util import str_to_bytes, Config 8 | 9 | TEST_STR = ( 10 | '1.3', '2kb', '3 mb', '4GB', '9Zb', '2.3pB' 11 | ) 12 | ANS = ( 13 | 1.3, 2048.0, 3145728.0, 4294967296.0, 10625324586456701730816.0, 14 | 2589569785738035.2 15 | ) 16 | 17 | 18 | class UtilityTest(unittest.TestCase): 19 | def test_str_to_bytes(self): 20 | for t, a in zip(TEST_STR, ANS): 21 | ans = str_to_bytes(t) 22 | self.assertEqual(ans, a, f"{t} != {a}") 23 | 24 | def test_config(self): 25 | d = Config(a=1, b=2) 26 | self.assertTrue(hasattr(d, 'a')) 27 | self.assertTrue(hasattr(d, 'b')) 28 | self.assertTrue(hasattr(d, 'non-exist')) 29 | self.assertIs(d.a, 1) 30 | self.assertIs(d.b, 2) 31 | d.update(a=2, b=3) 32 | self.assertIs(d.a, 2) 33 | self.assertIs(d.b, 3) 34 | d.a = 9 35 | self.assertIs(d.a, 9) 36 | d.update(Config(b=6, f=5)) 37 | self.assertIs(d.b, 6) 38 | self.assertIs(d.f, 5) 39 | d.pop('b') 40 | self.assertIsNone(d.b) 41 | 42 | 43 | if __name__ == '__main__': 44 | unittest.main() 45 | -------------------------------------------------------------------------------- /Tests/vgg_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 16 5 | 6 | import os 7 | import unittest 8 | 9 | if not os.getcwd().endswith('Tests'): 10 | os.chdir('Tests') 11 | 12 | import numpy as np 13 | from PIL import Image 14 | 15 | from VSR.Backend.TF.Util import Vgg 16 | 17 | URL = 'data/set5_x2/img_001_SRF_2_LR.png' 18 | image_boy = np.asarray(Image.open(URL)) 19 | 20 | 21 | class VggTest(unittest.TestCase): 22 | def import_tf(self): 23 | import tensorflow as tf 24 | return tf 25 | 26 | def test_vgg16(self): 27 | vgg = Vgg(False, vgg=Vgg.VGG16) 28 | x = np.random.normal(size=[16, 128, 128, 3]) 29 | y = vgg(x) 30 | self.assertEqual(y.shape, (16,)) 31 | 32 | def test_vgg19(self): 33 | vgg = Vgg(False, vgg=Vgg.VGG19) 34 | x = np.random.normal(size=[16, 128, 128, 3]) 35 | y = vgg(x, 'block2_conv2') 36 | self.assertEqual(y.shape, (16, 64, 64, 128)) 37 | 38 | def test_vgg_classify(self): 39 | vgg16 = Vgg(True, vgg=Vgg.VGG16) 40 | vgg19 = Vgg(True, vgg=Vgg.VGG19) 41 | x = np.expand_dims(image_boy, 0) 42 | y1 = vgg16(x) 43 | y2 = vgg19(x) 44 | tf = self.import_tf() 45 | with tf.Session() as sess: 46 | y1, y2 = sess.run([y1, y2]) 47 | self.assertEqual(y2[0].tolist().index(y2.max()), 48 | y1[0].tolist().index(y1.max())) 49 | 50 | def test_multiple_call(self): 51 | vgg1 = Vgg(False, vgg=Vgg.VGG16) 52 | vgg2 = Vgg(False, vgg=Vgg.VGG16) 53 | x = np.expand_dims(image_boy, 0) 54 | y1 = vgg1(x) 55 | y2 = vgg2(x) 56 | y3 = vgg2(x.copy()) 57 | tf = self.import_tf() 58 | with tf.Session() as sess: 59 | y1, y2, y3 = sess.run([y1, y2, y3]) 60 | self.assertEqual(y1, y2) 61 | self.assertEqual(y2, y3) 62 | 63 | 64 | if __name__ == '__main__': 65 | unittest.main() 66 | -------------------------------------------------------------------------------- /Tests/virtualfile_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Unit test for DataLoader.VirtualFile 3 | """ 4 | import os 5 | import unittest 6 | 7 | if not os.getcwd().endswith('Tests'): 8 | os.chdir('Tests') 9 | from VSR.DataLoader.VirtualFile import * 10 | from VSR.Util.ImageProcess import img_to_array 11 | 12 | RAW = 'data/video/raw_32x32.yv12' 13 | IMG = 'data/set5_x2/img_001_SRF_2_LR.png' 14 | VID = 'data/video/custom_pair/lr/xiuxian' 15 | 16 | 17 | class VirtualFileTest(unittest.TestCase): 18 | def test_image_read(self): 19 | vf = ImageFile(IMG) 20 | self.assertEqual(vf.name, 'img_001_SRF_2_LR') 21 | img = vf.read_frame() 22 | self.assertIsInstance(img, list) 23 | self.assertEqual(len(img), 1) 24 | self.assertIsInstance(img[0], Image.Image) 25 | self.assertEqual(img[0].width, 256) 26 | self.assertEqual(img[0].height, 256) 27 | 28 | def test_video_read(self): 29 | vf = ImageFile(VID) 30 | self.assertEqual(vf.name, 'xiuxian') 31 | vid = vf.read_frame(3) 32 | self.assertIsInstance(vid, list) 33 | self.assertEqual(len(vid), 3) 34 | self.assertEqual(vid[0].width, 240) 35 | self.assertEqual(vid[0].height, 135) 36 | 37 | def test_raw_read(self): 38 | vf = RawFile(RAW, 'YV12', [32, 32]) 39 | self.assertEqual(vf.name, 'raw_32x32') 40 | raw = vf.read_frame(vf.frames) 41 | self.assertEqual(len(raw), vf.frames) 42 | self.assertEqual(raw[0].width, 32) 43 | self.assertEqual(raw[0].height, 32) 44 | 45 | def test_image_seek(self): 46 | vf = ImageFile(IMG, False) 47 | f1 = vf.read_frame(1)[0] 48 | vf.seek(0, SEEK_SET) 49 | f2 = vf.read_frame(1)[0] 50 | vf.seek(-1, SEEK_CUR) 51 | f3 = vf.read_frame(1)[0] 52 | vf.seek(-1, SEEK_END) 53 | f4 = vf.read_frame(1)[0] 54 | vf.seek(-2, SEEK_END) 55 | f5 = vf.read_frame(1)[0] 56 | 57 | F = [f1, f2, f3, f4, f5] 58 | F = [img_to_array(f) for f in F] 59 | self.assertTrue(np.all(F[0] == F[1])) 60 | self.assertTrue(np.all(F[1] == F[2])) 61 | self.assertTrue(np.all(F[3] == F[4])) 62 | 63 | def test_vid_seek(self): 64 | vf = ImageFile(VID, False) 65 | f1 = vf.read_frame(1)[0] 66 | vf.seek(0, SEEK_SET) 67 | f2 = vf.read_frame(1)[0] 68 | vf.seek(-1, SEEK_CUR) 69 | f3 = vf.read_frame(1)[0] 70 | vf.seek(-1, SEEK_END) 71 | f4 = vf.read_frame(1)[0] 72 | vf.seek(2, SEEK_SET) 73 | f5 = vf.read_frame(1)[0] 74 | 75 | F = [f1, f2, f3, f4, f5] 76 | F = [img_to_array(f) for f in F] 77 | self.assertTrue(np.all(F[0] == F[1])) 78 | self.assertTrue(np.all(F[1] == F[2])) 79 | self.assertTrue(np.all(F[3] == F[4])) 80 | 81 | def test_raw_seek(self): 82 | vf = RawFile(RAW, 'YV12', [32, 32]) 83 | f1 = vf.read_frame(1)[0] 84 | vf.seek(0, SEEK_SET) 85 | f2 = vf.read_frame(1)[0] 86 | vf.seek(-1, SEEK_CUR) 87 | f3 = vf.read_frame(1)[0] 88 | vf.seek(-1, SEEK_END) 89 | f4 = vf.read_frame(1)[0] 90 | vf.seek(-2, SEEK_END) 91 | vf.seek(1, SEEK_CUR) 92 | f5 = vf.read_frame(1)[0] 93 | 94 | F = [f1, f2, f3, f4, f5] 95 | F = [img_to_array(f) for f in F] 96 | self.assertTrue(np.all(F[0] == F[1])) 97 | self.assertTrue(np.all(F[1] == F[2])) 98 | self.assertTrue(np.all(F[3] == F[4])) 99 | 100 | def test_vf_copy(self): 101 | import copy 102 | vf0 = ImageFile(IMG, False) 103 | vf1 = copy.deepcopy(vf0) 104 | vf0.read_frame(1) 105 | try: 106 | vf0.read_frame(1) 107 | self.assertFalse(True, "Unreachable code") 108 | except EOFError: 109 | pass 110 | vf1.read_frame(1) 111 | 112 | 113 | if __name__ == '__main__': 114 | unittest.main() 115 | -------------------------------------------------------------------------------- /Tools/CelebA.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2019 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Jan 3rd, 2019 6 | 7 | Pre-processing CelebA dataset: 8 | - Crop and resize to WxH 9 | - Randomly split into (train, test) 10 | """ 11 | 12 | # Copyright (c): Wenyi Tang 2017-2019. 13 | # Author: Wenyi Tang 14 | # Email: wenyi.tang@intel.com 15 | # Update Date: 2019/4/3 下午5:03 16 | 17 | import argparse 18 | from pathlib import Path 19 | 20 | import numpy as np 21 | from PIL import Image 22 | 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument("celeba", help="CelebA root folder.") 25 | parser.add_argument("-W", type=int, default=64, help="width") 26 | parser.add_argument("-H", type=int, default=64, help="height") 27 | parser.add_argument("--n_test", type=int, default=10000, help="test samples") 28 | args = parser.parse_args() 29 | 30 | 31 | def main(): 32 | root = Path(args.celeba) 33 | if not root.exists(): 34 | raise FileNotFoundError("Root of CelebA does not exist!") 35 | 36 | images = list(root.rglob('*.jpg')) 37 | resize_dir = root / 'resize{}'.format(args.W) 38 | test_dir = root / 'test{}'.format(args.W) 39 | resize_dir.mkdir(parents=True, exist_ok=False) 40 | test_dir.mkdir(parents=True, exist_ok=False) 41 | 42 | np.random.shuffle(images) 43 | for img in images[:args.n_test]: 44 | x = Image.open(img) 45 | dw = (x.width - args.W) // 2 46 | dh = (x.height - args.H) // 2 47 | box = [dw, dh, x.width - dw, x.height - dh] 48 | x.crop(box).save(str(test_dir) + '/{}.png'.format(img.stem)) 49 | 50 | for img in images[args.n_test:]: 51 | x = Image.open(img) 52 | dw = (x.width - args.W) // 2 53 | dh = (x.height - args.H) // 2 54 | box = [dw, dh, x.width - dw, x.height - dh] 55 | x.crop(box).save(str(resize_dir) + '/{}.png'.format(img.stem)) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /Tools/FFmpegHelper.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/3 下午5:03 5 | 6 | import argparse 7 | import subprocess 8 | import uuid 9 | from pathlib import Path 10 | 11 | import tqdm 12 | 13 | parser = argparse.ArgumentParser( 14 | description="Helper tool to use ffmpeg for transcoding.") 15 | parser.add_argument("input_dir", help="root folder of the raw videos.") 16 | parser.add_argument("output_dir", help="root folder of the targets.") 17 | parser.add_argument("--gop", type=int, default=1, 18 | help="[Enc] group of pictures (1)") 19 | parser.add_argument("--bf", type=int, default=0, help="[Enc] # B frames (0).") 20 | parser.add_argument("--codec", default="libx264", 21 | help="[Enc] encoder codec (libx264).") 22 | parser.add_argument("--qp", type=int, default=0, 23 | help="[Enc] quality index. [0, 51]") 24 | FLAGS = parser.parse_args() 25 | 26 | 27 | def _check_ffmpeg(): 28 | import shutil 29 | if shutil.which('ffmpeg') is None: 30 | raise FileNotFoundError("Couldn't find ffmpeg!") 31 | 32 | 33 | def parse_inputfile(path): 34 | filename = Path(path).stem 35 | suffix = Path(path).suffix 36 | size = filename.split('_')[-1] 37 | _size = [int(i) for i in size.split('x')] 38 | assert len(_size) == 2 39 | return size, suffix[1:].lower() 40 | 41 | 42 | def encode(file, work_dir): 43 | cmd = 'ffmpeg -f rawvideo' 44 | size, fmt = parse_inputfile(file) 45 | tmp_name = work_dir / f'{str(uuid.uuid4())}.264' 46 | cmd += f' -pix_fmt {fmt}' 47 | cmd += f' -s:v {size}' 48 | cmd += f' -i {str(file)}' 49 | cmd += f' -vcodec {FLAGS.codec}' 50 | cmd += f' -g {FLAGS.gop}' 51 | cmd += f' -bf {FLAGS.bf}' 52 | cmd += f' -qp {FLAGS.qp}' 53 | cmd += f' -f rawvideo {str(tmp_name)}' 54 | # print(cmd) 55 | subprocess.call(cmd.split(' '), stderr=subprocess.DEVNULL) 56 | return tmp_name, fmt 57 | 58 | 59 | def decode(file, output_dir, name, fmt): 60 | cmd = f'ffmpeg -i {str(file)}' 61 | output_name = output_dir / f'{name}_{FLAGS.qp}.{fmt}' 62 | cmd += f' -f rawvideo -pix_fmt {fmt}' 63 | cmd += f' {str(output_name)} -y' 64 | # print(cmd) 65 | subprocess.call(cmd.split(' '), stderr=subprocess.DEVNULL) 66 | 67 | 68 | def main(): 69 | _check_ffmpeg() 70 | raw_videos = filter(lambda f: f.is_file(), Path(FLAGS.input_dir).rglob('*')) 71 | raw_videos = list(raw_videos) 72 | if not raw_videos: 73 | raw_videos = filter(lambda f: f.is_file(), [Path(FLAGS.input_dir)]) 74 | tmp_dir = Path('/tmp/vsr/tools/_ffmpeg') 75 | tmp_dir.mkdir(exist_ok=True, parents=True) 76 | save_dir = Path(FLAGS.output_dir) 77 | save_dir.mkdir(exist_ok=True, parents=True) 78 | with tqdm.tqdm(sorted(raw_videos), ascii=True, unit=' video') as r: 79 | for fp in r: 80 | r.set_postfix({'name': fp.name}) 81 | stream, fmt = encode(fp, tmp_dir) 82 | decode(stream, save_dir, fp.stem, fmt) 83 | 84 | 85 | if __name__ == '__main__': 86 | main() 87 | -------------------------------------------------------------------------------- /Tools/Image2Raw.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/4 下午2:42 5 | 6 | import argparse 7 | from pathlib import Path 8 | 9 | import numpy as np 10 | import tqdm 11 | from PIL import Image 12 | 13 | from Tools.Misc import YUVConverter 14 | 15 | _YUV_COLOR = ("NV12", "NV21", "YV12", "YV21") 16 | _RGB_COLOR = ("RGBA",) 17 | 18 | parser = argparse.ArgumentParser( 19 | usage=r'''python Image2Raw.py input_dir output_dir [--options]''', 20 | description=r'''Convert a folder of images to raw video format (FOURCC).''') 21 | parser.add_argument("input_dir", help="root of the input folder, " 22 | "the leaf will be the last child-folder " 23 | "containing individual image files.") 24 | parser.add_argument("output_dir", help="root of the output path.") 25 | parser.add_argument("--color_fmt", choices=_YUV_COLOR + _RGB_COLOR, 26 | default='NV12', 27 | help="output color format") 28 | FLAGS = parser.parse_args() 29 | 30 | 31 | def parse_video_clips(path): 32 | _path = Path(path) 33 | if not _path.exists(): 34 | raise FileNotFoundError(f"{path} doesn't exist!!") 35 | 36 | files = _path.rglob('*') 37 | parents = set(f.parent for f in filter(lambda f: f.is_file(), files)) 38 | return parents 39 | 40 | 41 | def read_video_frames(path): 42 | _path = Path(path) 43 | files = sorted(_path.glob('*')) 44 | images = [Image.open(f) for f in files] 45 | if FLAGS.color_fmt in _YUV_COLOR: 46 | images = [m.convert('YCbCr') for m in images] 47 | mat = np.stack(images).transpose([0, 3, 1, 2]) # [NCHW] 48 | return { 49 | 'data': mat, 50 | 'length': mat.shape[0], 51 | 'name': path.stem, 52 | 'width': mat.shape[3], 53 | 'height': mat.shape[2], 54 | } 55 | 56 | 57 | def main(): 58 | videos = parse_video_clips(FLAGS.input_dir) 59 | root = Path(FLAGS.output_dir) 60 | root.mkdir(exist_ok=True, parents=True) 61 | print(f" [*] Total videos found: {len(videos)}.") 62 | with tqdm.tqdm(videos, ascii=True, unit=' video') as r: 63 | for fp in r: 64 | data = read_video_frames(fp) 65 | r.set_postfix({"name": data['name']}) 66 | if FLAGS.color_fmt in _YUV_COLOR: 67 | cvt = YUVConverter(data['data']) 68 | bytes = cvt.to(FLAGS.color_fmt).getbuffer().tobytes() 69 | else: 70 | raise NotImplementedError 71 | nm = f"{data['name']}_{cvt.width}x{cvt.height}.{FLAGS.color_fmt}" 72 | save_path = root / nm 73 | with save_path.open('wb') as fd: 74 | fd.write(bytes) 75 | 76 | 77 | if __name__ == '__main__': 78 | main() 79 | exit(0) 80 | -------------------------------------------------------------------------------- /Tools/Raw2Image.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/3 下午5:03 5 | 6 | import argparse 7 | import multiprocessing as mp 8 | import re 9 | from pathlib import Path 10 | 11 | import tqdm 12 | 13 | from VSR.DataLoader.VirtualFile import RawFile 14 | 15 | parser = argparse.ArgumentParser( 16 | description="Convert a raw video to a folder of images.") 17 | parser.add_argument("input_dir", help="root folder of raw videos.") 18 | parser.add_argument("output_dir", help="root folder of images") 19 | parser.add_argument("--width", type=int, default=0, 20 | help="default 0. Auto detect from file name.") 21 | parser.add_argument("--height", type=int, default=0, 22 | help="default 0. Auto detect from file name.") 23 | parser.add_argument("--overwrite", action='store_true', 24 | help="overwrite existing files with same name.") 25 | FLAGS = parser.parse_args() 26 | 27 | 28 | def guess_file_size(file): 29 | name = file.name 30 | rept = re.compile("\d+[xX]\d+") 31 | for i in name.split('_'): 32 | ans = re.findall(rept, i) 33 | if ans: 34 | size = ans[0].lower().split('x') 35 | return int(size[0]), int(size[1]) 36 | return -1, -1 37 | 38 | 39 | def parse_format(fmt): 40 | if fmt.upper() in ('YUV', 'YUV420P'): 41 | return 'YV12' 42 | return fmt.upper() 43 | 44 | 45 | def encode(file, save_dir): 46 | w = FLAGS.width 47 | h = FLAGS.height 48 | if w == 0 or h == 0: 49 | w, h = guess_file_size(file) 50 | if w <= 0 or h <= 0: 51 | raise ValueError("No width/height can be retrieved!") 52 | fmt = file.suffix[1:] 53 | fmt = parse_format(fmt) 54 | save_dir /= file.stem 55 | save_dir.mkdir(exist_ok=FLAGS.overwrite, parents=True) 56 | fd = RawFile(file, fmt, [w, h]) 57 | frames = fd.read_frame(fd.frames) 58 | for i, f in enumerate(frames): 59 | f.convert('RGB').save(f'{str(save_dir)}/{i:05d}.png') 60 | return file.stem 61 | 62 | 63 | def main(): 64 | input_dir = Path(FLAGS.input_dir) 65 | if input_dir.is_dir(): 66 | raw_videos = filter(lambda f: f.is_file(), input_dir.rglob('*')) 67 | else: 68 | assert input_dir.is_file() 69 | raw_videos = [input_dir] 70 | raw_videos = sorted(raw_videos) 71 | save_dir = Path(FLAGS.output_dir) 72 | save_dir.mkdir(exist_ok=True, parents=True) 73 | pool = mp.pool.ThreadPool() 74 | results = [] 75 | for fp in raw_videos: 76 | results.append(pool.apply_async(encode, (fp, save_dir))) 77 | with tqdm.tqdm(results, ascii=True, unit='image') as r: 78 | for i in r: 79 | name = i.get() 80 | r.set_postfix({'name': name}) 81 | pool.close() 82 | pool.join() 83 | 84 | 85 | if __name__ == '__main__': 86 | main() 87 | -------------------------------------------------------------------------------- /Tools/SeqVisual.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/22 下午3:54 5 | 6 | import argparse 7 | from pathlib import Path 8 | 9 | import numpy as np 10 | from PIL import Image 11 | from VSR.DataLoader.VirtualFile import ImageFile 12 | 13 | parser = argparse.ArgumentParser(description="Sequential Visualizer") 14 | parser.add_argument("input_dir", help="input folder") 15 | parser.add_argument("output", help="output file path") 16 | parser.add_argument("--row", '-r', type=int, default=-1, help="row number") 17 | parser.add_argument("--col", '-c', type=int, default=-1, help="column number") 18 | parser.add_argument("--zoom", nargs=4, type=int, help='zoom coordinate') 19 | parser.add_argument("--compose_id", type=int, default=None) 20 | 21 | 22 | def main(): 23 | flags = parser.parse_args() 24 | fp = ImageFile(flags.input_dir) 25 | frames = np.stack(fp.read_frame(fp.frames)) 26 | if flags.zoom: 27 | frames = frames[:, flags.zoom[1]: flags.zoom[3], 28 | flags.zoom[0]: flags.zoom[2]] 29 | savedir = Path(flags.output) 30 | if flags.compose_id is not None: 31 | compose = frames[flags.compose_id] 32 | else: 33 | compose = None 34 | if 0 <= flags.row < frames.shape[1]: 35 | sliced = frames[:, flags.row] 36 | if compose is not None: 37 | sliced = np.concatenate([compose, sliced], axis=0) 38 | if savedir.is_dir(): 39 | savedir.mkdir(exist_ok=True, parents=True) 40 | savedir /= f'{fp.name}_slice_row{flags.row}.png' 41 | elif 0 <= flags.col < frames.shape[2]: 42 | sliced = frames[:, :, flags.col] 43 | sliced = np.transpose(sliced, [1, 0, 2]) 44 | if compose: 45 | sliced = np.concatenate([compose, sliced], axis=1) 46 | if savedir.is_dir(): 47 | savedir.mkdir(exist_ok=True, parents=True) 48 | savedir /= f'{fp.name}_slice_col{flags.col}.png' 49 | Image.fromarray(sliced, 'RGB').save(str(savedir)) 50 | 51 | 52 | if __name__ == '__main__': 53 | main() 54 | -------------------------------------------------------------------------------- /Tools/YoukuPackage.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/5/22 下午1:25 5 | 6 | import argparse 7 | import os 8 | import subprocess 9 | from pathlib import Path 10 | 11 | import tqdm 12 | 13 | parser = argparse.ArgumentParser(description="Youku VSR Packager") 14 | parser.add_argument("-i", "--input_dir") 15 | parser.add_argument("-o", "--output_dir") 16 | parser.add_argument("--full_percentage", default=0.1) 17 | parser.add_argument("--extract_interval", default=25) 18 | parser.add_argument("-v", action='store_true', help="Show debug information") 19 | FLAGS = parser.parse_args() 20 | 21 | 22 | def _check_ffmpeg(): 23 | import shutil 24 | if shutil.which('ffmpeg') is None: 25 | raise FileNotFoundError("Couldn't find ffmpeg!") 26 | 27 | 28 | def gen_y4m_full(url, num): 29 | output = f'{FLAGS.output_dir}/Youku_{num}_h_Res.y4m' 30 | cmd = f'ffmpeg -i {str(url)} -pix_fmt yuv420p -vsync 0 {str(output)} -y' 31 | if FLAGS.v: 32 | print(cmd) 33 | subprocess.call(cmd, stderr=subprocess.DEVNULL, shell=True) 34 | 35 | 36 | def gen_y4m_part(url, num): 37 | output = f'{FLAGS.output_dir}/Youku_{num}_h_Sub25_Res.y4m' 38 | cmd = f'ffmpeg -i {str(url)} -pix_fmt yuv420p ' 39 | cmd += f"-vf select='not(mod(n\\,{FLAGS.extract_interval}))' " 40 | cmd += f"-vsync 0 {str(output)} -y" 41 | if FLAGS.v: 42 | print(cmd) 43 | subprocess.call(cmd, stderr=subprocess.DEVNULL, shell=True) 44 | 45 | 46 | def parse_video_clips(path): 47 | _path = Path(path) 48 | if not _path.exists(): 49 | raise FileNotFoundError(f"{path} doesn't exist!!") 50 | 51 | files = _path.rglob('*') 52 | parents = set(f.parent for f in filter(lambda f: f.is_file(), files)) 53 | return sorted(parents) 54 | 55 | 56 | def parse_url(path): 57 | _path = Path(path) 58 | if not _path.exists(): 59 | raise FileNotFoundError(f"{path} doesn't exist!!") 60 | files = filter(lambda f: f.is_file(), _path.glob('*')) 61 | files = sorted(files) 62 | for i, fp in enumerate(files): 63 | target = _path / f'frames_{i:04d}{fp.suffix}' 64 | fp.rename(target) 65 | assert target.exists() 66 | return _path / f'frames_%04d{fp.suffix}' 67 | 68 | 69 | def zip(url): 70 | url = Path(url) 71 | os.chdir(url) 72 | cmd = 'zip youku_results.zip *.y4m' 73 | subprocess.call(cmd, shell=True) 74 | subprocess.call("rm *.y4m", shell=True) 75 | 76 | 77 | def main(): 78 | _check_ffmpeg() 79 | videos = parse_video_clips(FLAGS.input_dir) 80 | root = Path(FLAGS.output_dir) 81 | root.mkdir(exist_ok=True, parents=True) 82 | print(f" [*] Total videos found: {len(videos)}.") 83 | with tqdm.tqdm(videos, ascii=True, unit=' video') as r: 84 | for i, fp in enumerate(r): 85 | _, num, _ = fp.name.split('_') 86 | url = parse_url(fp) 87 | if i < FLAGS.full_percentage * len(videos): 88 | gen_y4m_full(url, num) 89 | else: 90 | gen_y4m_part(url, num) 91 | r.set_postfix({"name": fp.stem}) 92 | zip(root) 93 | 94 | 95 | if __name__ == '__main__': 96 | main() 97 | -------------------------------------------------------------------------------- /Train/README.md: -------------------------------------------------------------------------------- 1 | # Model Tools 2 | VSR provides **2** tools along with code repo, that help users train and test models. 3 | 4 | To select a proper backend as you wish, please read [Change-backend](../Docs/HowTo/Change-backend.md). 5 | 6 | ## 1. Train 7 | ### Usage: `py train.py -h` 8 | ```bash 9 | usage: train.py [-h] [-p PARAMETER] [--save_dir SAVE_DIR] 10 | [--data_config DATA_CONFIG] [--dataset DATASET] 11 | [--epochs EPOCHS] [--steps STEPS] [--val_steps VAL_STEPS] 12 | [--cuda] [--threads THREADS] [--memory_limit MEMORY_LIMIT] 13 | [--traced_val] [--pretrain PRETRAIN] [--export EXPORT] 14 | [-c COMMENT] 15 | {espcn,srcnn,vdsr,dncnn,drcn,drrn,ffdnet,edsr,carn,dbpn,rcan,srfeat,esrgan,msrn,crdn,mldn,drn,sofvsr,vespcn,frvsr,qprn,ufvsr,yovsr,tecogan,spmc,rbpn,didn,dhdn,grdn,resunet,edrn,frn,ran} 16 | ``` 17 | The bracket {} lists all supported models in the current backend. 18 | 19 | ### Dataset 20 | The dataset used in this tool is described in [datasets.yaml](../Data/datasets.yaml). 21 | You can read [README.md](../Data/README.md) in that folder for details. 22 | 23 | ### Samples 24 | ```bash 25 | python train.py srcnn --dataset 91image --epochs 100 --cuda 26 | ``` 27 | 28 | If your system memory is not enough to hold all training data, and you want to running other process, add `--memory_limit=???` to constrain the memory usage (roughly). 29 | ```bash 30 | python train.py vdsr --dataset div2k --epochs 100 --cuda --memory_limit=4GB 31 | ``` 32 | 33 | If you want to continue to train from an external checkpoint, you can explicitly specify the checkpoint by adding `--pretrain=`. 34 | ```bash 35 | python train.py carn --dataset div2k --epochs 100 --cuda --pretrain=/model/carn/carn.pth 36 | ``` 37 | 38 | 39 | ## 2. Evaluate 40 | Evaluation is almost the same as training. 41 | ### Usage: `py eval.py -h` 42 | ```bash 43 | usage: eval.py [-h] [-p PARAMETER] [-t [TEST [TEST ...]]] 44 | [--save_dir SAVE_DIR] [--data_config DATA_CONFIG] 45 | [--pretrain PRETRAIN] [--ensemble] [--video] [--cuda] 46 | [--threads THREADS] [--output_index OUTPUT_INDEX] 47 | [--auto_rename] [-c COMMENT] 48 | {espcn,srcnn,vdsr,dncnn,drcn,drrn,ffdnet,edsr,carn,dbpn,rcan,srfeat,esrgan,msrn,crdn,mldn,drn,sofvsr,vespcn,frvsr,qprn,ufvsr,yovsr,tecogan,spmc,rbpn,didn,dhdn,grdn,resunet,edrn,frn,ran} 49 | ``` 50 | The bracket {} lists all supported models in the current backend. 51 | 52 | ### Flag `-t` 53 | You can specify a named dataset through `-t`, or an existing path to `-t`. 54 | `-t` also supports a list of arguments. 55 | ```bash 56 | python eval.py vdsr --cuda -t set5 set14 57 | ``` 58 | ```bash 59 | python eval.py vdsr --cuda -t /data/datasets/test/my-photos/*01.png 60 | python eval.py vdsr --cuda -t /data/datasets/test/bsd100 --pretrain=/model/vdsr/vdsr.pth 61 | ``` 62 | 63 | ### Flag `--video` 64 | If you are evaluating a video SR model, to read external data as video stream, parse `--video`: 65 | ```bash 66 | python eval.py vespcn --cuda -t /data/video/my-videos --video 67 | ``` 68 | However, if evaluate with a named dataset, `--video` is not needed: 69 | ```bash 70 | # vid4 in datasets.yaml is tagged with [video] 71 | python eval.py vespcn --cuda -t vid4 72 | ``` 73 | 74 | ## 3. Changing model parameters 75 | All model parameters are in [par](./par/). 76 | You can overwrite model parameters by adding arguments accordingly. 77 | 78 | For example, `srcnn` has parameters under [srcnn.yml](./par/pytorch/srcnn.yml). 79 | ```yaml 80 | srcnn: 81 | scale: 4 82 | channel: 3 83 | filters: [9, 5, 5] 84 | upsample: true 85 | image_weight: 1 86 | feature_weight: 0 # perceptual loss 87 | ``` 88 | Then you can overwrite `channel` by: 89 | ```bash 90 | python train.py srcnn --cuda --dataset 91image --epochs 200 --channel 1 91 | ``` 92 | -------------------------------------------------------------------------------- /Train/check_dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 4 - 17 5 | 6 | import argparse 7 | from pathlib import Path 8 | 9 | from VSR.DataLoader import load_datasets 10 | 11 | 12 | def main(entry: str, ddf_file): 13 | entry = entry.upper() 14 | all_data = load_datasets(ddf_file) 15 | if entry not in all_data: 16 | raise KeyError(f"The dataset `{entry}` not found in the DDF") 17 | 18 | data = all_data.get(entry) 19 | print(f"Dataset: {data.name}") 20 | 21 | def _check(name: str): 22 | print(f"\n========= CHECKING {name} =========\n") 23 | if name in data and data[name] is not None: 24 | print(f"Found `{name}` set in \"{data.name}\":") 25 | _hr = data[name].hr 26 | _lr = data[name].lr 27 | video_type = _hr.as_video 28 | if video_type: 29 | print(f"\"{data.name}\" is video data") 30 | if _hr is not None: 31 | _hr = _hr.compile() 32 | print(f"Found {len(_hr)} ground-truth {name} data") 33 | if _lr is not None: 34 | _lr = _lr.compile() 35 | print(f"Found {len(_lr)} custom degraded {name} data") 36 | if len(_hr) != len(_lr): 37 | print( 38 | f" [E] Ground-truth data and degraded data quantity not matched!!") 39 | elif video_type: 40 | for x, y in zip(_hr, _lr): 41 | if x.frames != y.frames: 42 | print(f" [E] Video clip {x.name}|{y.name} quantity not matched!!") 43 | else: 44 | print(f"{data.name} doesn't contain any {name} data.") 45 | 46 | _check('train') 47 | _check('val') 48 | _check('test') 49 | 50 | 51 | if __name__ == '__main__': 52 | CWD = Path(__file__).resolve().parent.parent 53 | parser = argparse.ArgumentParser( 54 | description="Check the dataset and print out its content") 55 | parser.add_argument("dataset", type=str, 56 | help="The name of the dataset, case insensitive.") 57 | parser.add_argument("--description-file", default=f"{CWD}/Data/datasets.yaml", 58 | help="DDF file") 59 | flags = parser.parse_args() 60 | main(flags.dataset, flags.description_file) 61 | -------------------------------------------------------------------------------- /Train/par/keras/srcnn.yaml: -------------------------------------------------------------------------------- 1 | # srcnn 9-5-5 2 | --- 3 | srcnn: 4 | layers: 3 5 | filters: 6 | - 9 7 | - 1 8 | - 5 9 | scale: 4 10 | channel: 1 11 | 12 | batch: 4 13 | patch_size: 16 14 | lr: 1.0e-4 15 | lr_decay: 16 | method: multistep 17 | decay_step: [10000, 15000] 18 | decay_rate: 0.1 19 | -------------------------------------------------------------------------------- /Train/par/pytorch/carn.yml: -------------------------------------------------------------------------------- 1 | carn: 2 | scale: 4 3 | channel: 3 4 | multi_scale: 1 # change to 1 if use official pth file 5 | group: 1 6 | clip: 10 7 | 8 | batch_shape: [16, 3, 64, 64] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [500] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/cubic.yml: -------------------------------------------------------------------------------- 1 | cubic: 2 | scale: 4 3 | channel: 3 4 | depth: 1 5 | 6 | batch_shape: [16, 3, 32, 32] 7 | -------------------------------------------------------------------------------- /Train/par/pytorch/dbpn.yml: -------------------------------------------------------------------------------- 1 | dbpn: 2 | scale: 4 3 | channel: 3 4 | base_filter: 64 5 | feat: 256 6 | num_stages: 7 7 | 8 | batch_shape: [16, 3, 32, 32] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [500] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/dncnn.yml: -------------------------------------------------------------------------------- 1 | dncnn: 2 | scale: 1 3 | channel: 3 4 | noise: 25 # noise level added during training 5 | image_weight: 1 6 | feature_weight: 0 # perceptual loss 7 | 8 | batch_shape: [64, 3, 96, 96] 9 | lr: 1.0e-3 10 | lr_decay: 11 | method: multistep 12 | decay_step: [50] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/drcn.yml: -------------------------------------------------------------------------------- 1 | drcn: 2 | scale: 4 3 | channel: 3 4 | n_recur: 16 5 | 6 | batch_shape: [16, 3, 32, 32] 7 | lr: 1.0e-4 8 | lr_decay: 9 | method: multistep 10 | decay_step: [500] 11 | decay_rate: 0.1 12 | -------------------------------------------------------------------------------- /Train/par/pytorch/drn.yml: -------------------------------------------------------------------------------- 1 | drn: 2 | scale: 4 3 | channel: 3 4 | n_cb: 4 5 | noise: 0 6 | offset: 0 7 | 8 | batch_shape: [16, 3, 32, 32] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [500] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/drrn.yml: -------------------------------------------------------------------------------- 1 | drrn: 2 | scale: 4 3 | channel: 3 4 | n_ru: 3 5 | n_rb: 3 6 | clip: 0.01 7 | 8 | batch_shape: [16, 3, 32, 32] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [500] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/edrn.yml: -------------------------------------------------------------------------------- 1 | --- 2 | edrn: 3 | scale: 1 4 | channel: 3 5 | G0: 64 6 | EDRNkSize: 3 7 | EDRNconfig: B 8 | rgb_range: 255 9 | 10 | batch_shape: [16, 3, 32, 32] 11 | lr: 1.0e-4 12 | lr_decay: 13 | method: multistep 14 | decay_step: [500] 15 | decay_rate: 0.1 16 | -------------------------------------------------------------------------------- /Train/par/pytorch/edsr.yml: -------------------------------------------------------------------------------- 1 | edsr: 2 | scale: 4 3 | channel: 3 4 | n_resblocks: 16 5 | n_feats: 64 6 | 7 | batch_shape: [8, 3, 48, 48] 8 | lr: 1.0e-4 9 | lr_decay: 10 | method: multistep 11 | decay_step: [500] 12 | decay_rate: 0.1 13 | -------------------------------------------------------------------------------- /Train/par/pytorch/espcn.yml: -------------------------------------------------------------------------------- 1 | espcn: 2 | scale: 4 3 | channel: 3 4 | image_weight: 1 5 | feature_weight: 0 # perceptual loss 6 | 7 | batch_shape: [16, 3, 16, 16] 8 | lr: 1.0e-2 9 | lr_decay: 10 | method: multistep 11 | decay_step: [50, 75] 12 | decay_rate: 0.1 13 | -------------------------------------------------------------------------------- /Train/par/pytorch/esrgan.yml: -------------------------------------------------------------------------------- 1 | esrgan: 2 | scale: 4 3 | channel: 3 4 | nf: 64 5 | nb: 23 6 | gc: 32 7 | weights: [0.01, 1, 5.0e-3] 8 | patch_size: 128 9 | 10 | batch_shape: [16, 3, 32, 32] 11 | lr: 1.0e-4 12 | lr_decay: 13 | method: multistep 14 | decay_step: [250, 500, 1000, 1500] 15 | decay_rate: 0.5 16 | -------------------------------------------------------------------------------- /Train/par/pytorch/ffdnet.yml: -------------------------------------------------------------------------------- 1 | ffdnet: 2 | scale: 1 3 | channel: 3 4 | n_layers: 10 5 | level: 0 6 | training: true 7 | 8 | batch_shape: [16, 3, 32, 32] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [500] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/frvsr.yml: -------------------------------------------------------------------------------- 1 | --- 2 | frvsr: 3 | scale: 4 4 | channel: 3 5 | depth: 7 6 | weights: [1, 1, 1.0e-3] 7 | 8 | batch_shape: [16, 7, 3, 32, 32] 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [200, 300, 400] 13 | decay_rate: 0.5 14 | -------------------------------------------------------------------------------- /Train/par/pytorch/mldn.yml: -------------------------------------------------------------------------------- 1 | mldn: 2 | finetune: true 3 | scale: 1 4 | 5 | batch_shape: [16, 3, 32, 32] 6 | lr: 1.0e-4 7 | lr_decay: 8 | method: multistep 9 | decay_step: [500] 10 | decay_rate: 0.1 11 | -------------------------------------------------------------------------------- /Train/par/pytorch/msrn.yml: -------------------------------------------------------------------------------- 1 | msrn: 2 | scale: 4 3 | channel: 3 4 | rgb_range: 255 5 | 6 | batch_shape: [16, 3, 32, 32] 7 | lr: 1.0e-4 8 | lr_decay: 9 | method: multistep 10 | decay_step: [500] 11 | decay_rate: 0.1 12 | -------------------------------------------------------------------------------- /Train/par/pytorch/ntire19.yml: -------------------------------------------------------------------------------- 1 | --- 2 | frn: 3 | scale: 1 4 | channel: 3 5 | rgb_range: 255 6 | n_resgroups: 10 7 | n_resblocks: 10 8 | n_feats: 64 9 | reduction: 0.1 10 | res_scale: 1 11 | px: 1 12 | 13 | ran2: 14 | scale: 1 15 | channel: 3 16 | rgb_range: 255 17 | n_feats: 64 18 | n_resgroups: 10 19 | n_resblocks: 10 20 | nvis: false 21 | 22 | edrn: 23 | scale: 1 24 | channel: 3 25 | G0: 64 26 | EDRNkSize: 3 27 | EDRNconfig: B 28 | rgb_range: 255 29 | 30 | didn: 31 | scale: 1 32 | channel: 3 33 | filters: 64 34 | umodule: 8 35 | 36 | dhdn: 37 | scale: 1 38 | channel: 3 39 | filters: 128 40 | 41 | grdn: 42 | scale: 1 43 | channel: 3 44 | filters: 64 45 | grdb: 10 46 | rdb: 4 47 | 48 | resunet: 49 | scale: 1 50 | channel: 3 51 | filters: 128 52 | rb: 10 53 | 54 | batch_shape: [16, 3, 32, 32] 55 | lr: 1.0e-4 56 | lr_decay: 57 | method: multistep 58 | decay_step: [200, 300] 59 | decay_rate: 0.1 60 | -------------------------------------------------------------------------------- /Train/par/pytorch/qprn.yml: -------------------------------------------------------------------------------- 1 | --- 2 | qprn: 3 | gain: 64 4 | depth: 10 5 | channel: 3 6 | scale: 1 7 | weights: [1, 10, 1.0e-4, 0.1, 0] 8 | debug: 9 | see_warp: false 10 | see_flow: false 11 | reverse: false 12 | ablation: N 13 | gan: false 14 | 15 | batch_shape: [16, 3, 32, 32] 16 | lr: 1.0e-4 17 | lr_decay: 18 | method: multistep 19 | decay_step: [100, 250, 500] 20 | decay_rate: 0.5 21 | -------------------------------------------------------------------------------- /Train/par/pytorch/rbpn.yml: -------------------------------------------------------------------------------- 1 | --- 2 | rbpn: 3 | scale: 4 4 | channel: 3 5 | depth: 7 6 | residual: false 7 | base_filter: 256 8 | feat: 64 9 | num_stages: 3 10 | n_resblock: 5 11 | 12 | batch_shape: [4, 3, 16, 16] 13 | lr: 1.0e-4 14 | lr_decay: 15 | method: multistep 16 | decay_step: [500] 17 | decay_rate: 0.1 18 | -------------------------------------------------------------------------------- /Train/par/pytorch/rcan.yml: -------------------------------------------------------------------------------- 1 | rcan: 2 | scale: 4 3 | channel: 3 4 | n_resgroups: 10 5 | n_resblocks: 20 6 | n_feats: 64 7 | reduction: 16 8 | rgb_range: 255 9 | 10 | batch_shape: [16, 3, 32, 32] 11 | lr: 1.0e-4 12 | lr_decay: 13 | method: multistep 14 | decay_step: [500] 15 | decay_rate: 0.1 16 | -------------------------------------------------------------------------------- /Train/par/pytorch/realsr.yml: -------------------------------------------------------------------------------- 1 | realsr: 2 | scale: 4 3 | channel: 3 4 | nf: 64 5 | nb: 23 6 | pixel_weight: !!float 1 7 | feature_weight: !!float 0 8 | gan_weight: !!float 0 9 | 10 | batch_shape: [16, 3, 64, 64] 11 | -------------------------------------------------------------------------------- /Train/par/pytorch/rsr.yml: -------------------------------------------------------------------------------- 1 | rsr: 2 | scale: 1 3 | 4 | batch_shape: [16, 3, 32, 32] 5 | lr: 1.0e-5 6 | lr_decay: 7 | method: multistep 8 | decay_step: [500] 9 | decay_rate: 0.1 10 | -------------------------------------------------------------------------------- /Train/par/pytorch/sofvsr.yml: -------------------------------------------------------------------------------- 1 | sofvsr: 2 | channel: 1 3 | scale: 4 4 | depth: 3 5 | 6 | batch_shape: [16, 3, 1, 32, 32] 7 | lr: 1.0e-4 8 | lr_decay: 9 | method: multistep 10 | decay_step: [250, 500, 750, 1000, 1250] 11 | decay_rate: 0.5 12 | -------------------------------------------------------------------------------- /Train/par/pytorch/spmc.yml: -------------------------------------------------------------------------------- 1 | --- 2 | spmc: 3 | scale: 4 4 | channel: 3 5 | depth: 3 6 | # stage 1: optimize ME only 7 | # stage 2: optimize SR only 8 | # stage 3: jointly optimize 9 | stage: 3 10 | lambda1: 0.01 11 | lambda2: 0.01 12 | # add bicubic to output 13 | residual: false 14 | 15 | batch_shape: [16, 3, 3, 16, 16] 16 | lr: 1.0e-4 17 | lr_decay: 18 | method: multistep 19 | decay_step: [500] 20 | decay_rate: 0.1 21 | -------------------------------------------------------------------------------- /Train/par/pytorch/srcnn.yml: -------------------------------------------------------------------------------- 1 | srcnn: 2 | scale: 4 3 | channel: 3 4 | filters: [9, 5, 5] 5 | upsample: true 6 | image_weight: 1 7 | feature_weight: 0 # perceptual loss 8 | 9 | batch_shape: [16, 3, 32, 32] 10 | lr: 1.0e-2 11 | lr_decay: 12 | method: multistep 13 | decay_step: [50, 75] 14 | decay_rate: 0.1 15 | -------------------------------------------------------------------------------- /Train/par/pytorch/srfeat.yml: -------------------------------------------------------------------------------- 1 | srfeat: 2 | channel: 3 3 | scale: 4 4 | filters: 64 5 | num_residualblocks: 16 6 | weights: [1, 0.01, 0.01] 7 | patch_size: 64 8 | 9 | batch_shape: [16, 3, 16, 16] 10 | lr: 1.0e-4 11 | lr_decay: 12 | method: multistep 13 | decay_step: [500] 14 | decay_rate: 0.1 15 | -------------------------------------------------------------------------------- /Train/par/pytorch/srmd.yml: -------------------------------------------------------------------------------- 1 | --- 2 | srmd: 3 | scale: 4 4 | channel: 3 5 | layers: 12 6 | filters: 128 7 | degradation: 8 | kernel_type: 'anisotropic' # isotropic or anisotropic 9 | l1: 0.1 # scaling of eigen values on base 0. [0.1, 10] 10 | l2: 0.1 # scaling of eigen values on base 1. [0.1, l1] 11 | theta: 0.0 # rotation angle (rad) of the kernel. [0, pi] 12 | noise: 5 # noise stddev (0, 75] 13 | 14 | batch_shape: [16, 3, 64, 64] 15 | lr: 1.0e-4 16 | -------------------------------------------------------------------------------- /Train/par/pytorch/tecogan.yml: -------------------------------------------------------------------------------- 1 | --- 2 | tecogan: 3 | scale: 4 4 | channel: 3 5 | depth: 7 6 | # L2, Flow, Ping-pong, GAN 7 | weights: [1, 1, 0.5, 5.0e-4] 8 | vgg_layers: ['block2_conv2', 'block3_conv4', 'block4_conv4', 'block5_conv4'] 9 | vgg_layer_weights: [3.0e-5, 1.4e-6, 6.0e-6, 2.0e-3] 10 | gan_layer_weights: [1.67e-3, 1.43e-3, 8.33e-4, 2.0e-5] 11 | patch_size: 64 12 | 13 | batch_shape: [16, 7, 3, 16, 16] 14 | lr: 1.0e-4 15 | lr_decay: 16 | method: multistep 17 | decay_step: [250, 500] 18 | decay_rate: 0.5 19 | -------------------------------------------------------------------------------- /Train/par/pytorch/vdsr.yml: -------------------------------------------------------------------------------- 1 | vdsr: 2 | scale: 4 3 | channel: 3 4 | layers: 20 5 | upsample: true 6 | image_weight: 1 7 | feature_weight: 0 # perceptual loss 8 | 9 | batch_shape: [16, 3, 24, 24] 10 | lr: 1.0e-2 11 | lr_decay: 12 | method: multistep 13 | decay_step: [50, 75] 14 | decay_rate: 0.1 15 | -------------------------------------------------------------------------------- /Train/par/pytorch/vespcn.yml: -------------------------------------------------------------------------------- 1 | --- 2 | vespcn: 3 | depth: 3 4 | channel: 3 5 | scale: 4 6 | 7 | batch_shape: [16,3, 3, 16, 16] 8 | lr: 1.0e-4 9 | lr_decay: 10 | method: multistep 11 | decay_step: [200, 400] 12 | decay_rate: 0.1 13 | -------------------------------------------------------------------------------- /Train/par/tensorflow/carn.yaml: -------------------------------------------------------------------------------- 1 | carn: 2 | recursive: false 3 | n_residual: 3 4 | n_blocks: 3 5 | filters: 64 6 | clip: 10 7 | weight_decay: 0 8 | scale: 4 9 | channel: 3 10 | 11 | batch_shape: [4, 16, 16, 3] 12 | lr: 1.0e-4 13 | lr_decay: 14 | method: multistep 15 | decay_step: [150000] 16 | decay_rate: 0.1 17 | -------------------------------------------------------------------------------- /Train/par/tensorflow/dbpn.yaml: -------------------------------------------------------------------------------- 1 | dbpn: 2 | bp_layers: 7 3 | use_dense: true 4 | scale: 4 5 | channel: 3 6 | 7 | patch_size: 96 8 | batch: 16 9 | lr: 1.0e-4 10 | lr_decay: 11 | method: multistep 12 | decay_step: [60000, 160000] 13 | decay_rate: 1 14 | -------------------------------------------------------------------------------- /Train/par/tensorflow/dcscn.yaml: -------------------------------------------------------------------------------- 1 | dcscn: 2 | layers: 8 3 | reconstruction_layers: 1 4 | filters: 196 5 | min_filters: 48 6 | nin_filter: [64, 32] 7 | reconst_filter: 32 8 | filters_decay_gamma: 1.5 9 | drop_out: 0.8 10 | scale: 4 11 | channel: 3 12 | 13 | batch: 16 14 | patch_size: 64 15 | lr: 1.0e-4 16 | -------------------------------------------------------------------------------- /Train/par/tensorflow/dncnn.yaml: -------------------------------------------------------------------------------- 1 | dncnn: 2 | layers: 16 3 | weight_decay: 1.0e-6 4 | scale: 1 5 | channel: 3 6 | 7 | batch: 16 8 | patch_size: 64 9 | lr: 1.0e-4 10 | -------------------------------------------------------------------------------- /Train/par/tensorflow/drcn.yaml: -------------------------------------------------------------------------------- 1 | drcn: 2 | recur: 16 3 | filters: 256 4 | alpha: 1 5 | weight_decay: 1.0e-4 6 | custom_upsample: false 7 | scale: 4 8 | channel: 3 9 | 10 | batch: 16 11 | patch_size: 64 12 | lr: 1.0e-4 13 | -------------------------------------------------------------------------------- /Train/par/tensorflow/drrn.yaml: -------------------------------------------------------------------------------- 1 | drrn: 2 | residual_unit: 9 3 | recursive_block: 1 4 | grad_clip: 0.01 5 | weight_decay: 1.0e-6 6 | custom_upsample: false 7 | scale: 4 8 | channel: 3 9 | 10 | batch: 16 11 | patch_size: 64 12 | lr: 1.0e-4 13 | -------------------------------------------------------------------------------- /Train/par/tensorflow/drsr.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | drsr: 3 | name: 'drsr' 4 | n_cb: 4 5 | n_crb: 4 6 | weights: [1, 0.5, 1.0e-6, 1.0e-3] 7 | finetune: 4000 8 | noise_config: 9 | max: 75.0 10 | scale: 1.0 11 | offset: 0.0 12 | offset2: 0.06 13 | penalty: 0.7 14 | layers: 8 15 | type: "gaussian" 16 | crf: "../Data/crf.npz" 17 | channel: 3 18 | scale: 4 19 | tfrecords: false 20 | 21 | patch_size: 64 22 | batch: 16 23 | lr: 1.0e-4 24 | lr_decay: 25 | method: multistep 26 | decay_step: [] 27 | decay_rate: 1.0 28 | -------------------------------------------------------------------------------- /Train/par/tensorflow/drsr_v2.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | drsr_v2: 3 | weights: [1, 10, 1.0e-5] 4 | # mean_shift: [0, 0, 0] 5 | level: 1 6 | arch: 'crdb' 7 | auto_shift: true 8 | noise_config: 9 | max: 75.0 10 | scale: 1.0 11 | offset: 0.0 12 | penalty: 0.7 13 | layers: 8 14 | crf: "../Data/crf.npz" 15 | valid: false 16 | channel: 3 17 | scale: 4 18 | 19 | patch_size: 128 20 | batch: 16 21 | lr: 1.0e-4 22 | lr_decay: 23 | method: multistep 24 | decay_step: [150000] 25 | decay_rate: 0.1 26 | -------------------------------------------------------------------------------- /Train/par/tensorflow/duf.yaml: -------------------------------------------------------------------------------- 1 | duf: 2 | scale: 4 3 | channel: 3 4 | layers: 16 5 | filter_size: [5, 5] 6 | depth: 7 7 | 8 | batch: 16 9 | patch_size: 64 10 | lr: 1.0e-4 11 | -------------------------------------------------------------------------------- /Train/par/tensorflow/edsr.yaml: -------------------------------------------------------------------------------- 1 | edsr: 2 | layers: 32 3 | filters: 256 4 | clip: 0.1 5 | scale: 4 6 | channel: 3 7 | 8 | batch: 16 9 | patch_size: 128 10 | lr: 1.0e-4 11 | lr_decay: 12 | method: multistep 13 | decay_step: [150000] 14 | decay_rate: 0.1 15 | -------------------------------------------------------------------------------- /Train/par/tensorflow/espcn.yaml: -------------------------------------------------------------------------------- 1 | # espcn 5-3-3 2 | --- 3 | espcn: 4 | layers: 3 5 | filters: [64, 32, 32] 6 | kernel: [5, 3, 3] 7 | scale: 4 8 | channel: 3 9 | 10 | batch_shape: [16, 16, 16, 3] 11 | lr: 1.0e-2 12 | lr_decay: 13 | method: multistep 14 | decay_step: [10000, 15000] 15 | decay_rate: 0.1 16 | -------------------------------------------------------------------------------- /Train/par/tensorflow/ffdnet.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | ffdnet: 3 | sigma: 75 4 | space_down: 2 5 | channel: 3 6 | layers: 15 7 | training: true 8 | weight_decay: 0 9 | scale: 1 10 | 11 | batch: 32 12 | patch_size: 128 13 | lr: 1.0e-4 14 | -------------------------------------------------------------------------------- /Train/par/tensorflow/frvsr.yaml: -------------------------------------------------------------------------------- 1 | frvsr: 2 | l1_weight: 1 3 | flow_weight: 10 4 | scale: 4 5 | channel: 3 6 | 7 | depth: 10 8 | batch: 16 9 | patch_size: 64 10 | lr: 1.0e-4 11 | -------------------------------------------------------------------------------- /Train/par/tensorflow/gangp.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | gangp: 3 | name: gangp 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/idn.yaml: -------------------------------------------------------------------------------- 1 | idn: 2 | blocks: 4 3 | filters: 64 4 | delta: 16 5 | slice_factor: 4 6 | leaky_slope: 0.05 7 | weight_decay: 1.0e-4 8 | fine_tune_epoch: 200 9 | scale: 4 10 | channel: 3 11 | 12 | batch: 16 13 | patch_size: 64 14 | lr: 1.0e-4 15 | -------------------------------------------------------------------------------- /Train/par/tensorflow/lapsrn.yaml: -------------------------------------------------------------------------------- 1 | lapsrn: 2 | layers: 5 3 | scale: 4 4 | channel: 3 5 | 6 | batch: 16 7 | patch_size: 64 8 | lr: 1.0e-4 9 | -------------------------------------------------------------------------------- /Train/par/tensorflow/lsgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | lsgan: 3 | name: lsgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/memnet.yaml: -------------------------------------------------------------------------------- 1 | memnet: 2 | n_memblock: 6 3 | n_recur: 6 4 | filters: 64 5 | scale: 4 6 | channel: 3 7 | 8 | batch: 16 9 | patch_size: 64 10 | lr: 1.0e-4 11 | -------------------------------------------------------------------------------- /Train/par/tensorflow/msrn.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | msrn: 3 | n_blocks: 8 4 | scale: 4 5 | channel: 3 6 | 7 | batch: 4 8 | patch_size: 64 9 | lr: 1.0e-2 10 | lr_decay: 11 | method: multistep 12 | decay_step: [10000, 15000] 13 | decay_rate: 0.1 14 | -------------------------------------------------------------------------------- /Train/par/tensorflow/nlrn.yaml: -------------------------------------------------------------------------------- 1 | nlrn: 2 | scale: 4 3 | channel: 3 4 | filters: 128 5 | recurrents: 12 6 | clip: 2.5 7 | 8 | batch: 16 9 | patch_size: 64 10 | lr: 1.0e-4 11 | -------------------------------------------------------------------------------- /Train/par/tensorflow/ragan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | ragan: 3 | name: ragan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/ragangp.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | ragangp: 3 | name: ragangp 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/ralsgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | ralsgan: 3 | name: ralsgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/rcan.yaml: -------------------------------------------------------------------------------- 1 | rcan: 2 | channel_downscaling: 16 3 | n_rcab: 4 4 | n_rg: 3 5 | filters: 64 6 | scale: 4 7 | channel: 3 8 | 9 | batch: 16 10 | patch_size: 64 11 | lr: 1.0e-4 12 | -------------------------------------------------------------------------------- /Train/par/tensorflow/rdn.yaml: -------------------------------------------------------------------------------- 1 | # note our rdn implementation has fewer layers than paper's 2 | # (hard to train...) 3 | --- 4 | rdn: 5 | rdb_blocks: 8 6 | rdb_conv: 8 7 | global_filters: 64 8 | rdb_filters: 64 9 | scale: 4 10 | channel: 3 11 | 12 | batch: 16 13 | patch_size: 64 14 | lr: 1.0e-4 15 | -------------------------------------------------------------------------------- /Train/par/tensorflow/rgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | rgan: 3 | name: rgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/rgangp.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | rgangp: 3 | name: rgangp 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/rlsgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | rlsgan: 3 | name: rlsgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/sgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | sgan: 3 | name: sgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | optimizer: 12 | name: adam 13 | beta1: 0.5 14 | channel: 3 15 | scale: 1 # keep scale is necessary though it's unused 16 | 17 | batch: 64 18 | test_batch: 100 19 | validate_every_n_epoch: 5 20 | lr: 2.0e-4 21 | -------------------------------------------------------------------------------- /Train/par/tensorflow/srcnn.yaml: -------------------------------------------------------------------------------- 1 | # srcnn 9-5-5 2 | --- 3 | srcnn: 4 | layers: 3 5 | kernel: 6 | - 9 7 | - 5 8 | - 5 9 | custom_upsample: false 10 | scale: 4 11 | channel: 1 12 | 13 | batch: 4 14 | patch_size: 64 15 | lr: 1.0e-2 16 | lr_decay: 17 | method: multistep 18 | decay_step: [10000, 15000] 19 | decay_rate: 0.1 20 | -------------------------------------------------------------------------------- /Train/par/tensorflow/srdensenet.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | srdensenet: 3 | n_blocks: 8 4 | scale: 4 5 | channel: 3 6 | 7 | patch_size: 100 8 | batch: 16 9 | lr: 1.0e-4 -------------------------------------------------------------------------------- /Train/par/tensorflow/srfeat.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | srfeat: 3 | glayers: 16 4 | dlayers: 4 5 | vgg_layer: 'block5_conv4' 6 | init_epoch: 100 7 | gan_weight: 1.0e-3 8 | vgg_weight: 0.1569 # (1 / 12.75)^2 9 | scale: 4 10 | channel: 3 11 | 12 | patch_size: 128 13 | batch: 16 14 | lr: 1.0e-4 -------------------------------------------------------------------------------- /Train/par/tensorflow/srgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | srgan: 3 | glayers: 16 4 | dlayers: 4 5 | vgg_layer: 'block2_conv2' 6 | init_epoch: 100 7 | vgg_weight: 1.0e-6 8 | use_vgg: true 9 | mse_weight: 1 10 | gan_weight: 0.001 11 | scale: 4 12 | channel: 3 13 | 14 | patch_size: 128 15 | batch: 16 16 | lr: 1.0e-4 -------------------------------------------------------------------------------- /Train/par/tensorflow/vdsr.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | vdsr: 3 | layers: 20 4 | filters: 64 5 | custom_upsample: false 6 | scale: 4 7 | channel: 1 8 | 9 | batch: 16 10 | patch_size: 64 11 | lr: 1.0e-4 12 | -------------------------------------------------------------------------------- /Train/par/tensorflow/vespcn.yaml: -------------------------------------------------------------------------------- 1 | # VESPCN need data depth to be 3, 5, or 7. (2n + 1) 2 | --- 3 | vespcn: 4 | depth: 3 # this `depth` is to inform graph builder 5 | beta: 1 6 | gamma: 0.01 7 | scale: 4 8 | channel: 3 9 | weight_decay: 0 10 | batch: 4 11 | patch_size: 96 12 | depth: 3 # must be same as the model depth 13 | lr: 1.0e-4 14 | lr_decay: 15 | method: multistep 16 | decay_step: [] 17 | decay_rate: 1 18 | -------------------------------------------------------------------------------- /Train/par/tensorflow/wgan.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | wgan: 3 | name: wgan 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | nd_iter: 5 11 | use_bias: true 12 | optimizer: 13 | name: rmsprop 14 | 15 | channel: 3 16 | scale: 1 # keep scale is necessary though it's unused 17 | 18 | batch: 64 19 | test_batch: 100 20 | validate_every_n_epoch: 5 21 | lr: 2.0e-4 22 | -------------------------------------------------------------------------------- /Train/par/tensorflow/wgangp.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | wgangp: 3 | name: wgangp 4 | patch_size: 32 5 | z_dim: 128 6 | init_filter: 512 7 | linear: true # dense layer at the entry 8 | norm_g: bn # null, bn or sn 9 | norm_d: sn 10 | use_bias: true 11 | nd_iter: 5 12 | optimizer: 13 | name: adam 14 | beta1: 0.5 15 | channel: 3 16 | scale: 1 # keep scale is necessary though it's unused 17 | 18 | batch: 64 19 | test_batch: 100 20 | validate_every_n_epoch: 5 21 | lr: 2.0e-4 22 | -------------------------------------------------------------------------------- /VSR/Backend/Keras/Framework/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 30 5 | 6 | -------------------------------------------------------------------------------- /VSR/Backend/Keras/Models/Srcnn.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 30 5 | 6 | import tensorflow as tf 7 | 8 | from .Model import SuperResolution 9 | 10 | 11 | class Srcnn(tf.keras.Model): 12 | def __init__(self, channel, filters): 13 | super(Srcnn, self).__init__() 14 | self.net = [ 15 | tf.keras.layers.Conv2D(64, filters[0], padding='same', 16 | activation=tf.nn.relu), 17 | tf.keras.layers.Conv2D(32, filters[1], padding='same', 18 | activation=tf.nn.relu), 19 | tf.keras.layers.Conv2D(channel, filters[2], padding='same')] 20 | 21 | def call(self, inputs): 22 | x = inputs 23 | for layer in self.net: 24 | x = layer(x) 25 | return x 26 | 27 | 28 | class SRCNN(SuperResolution): 29 | def __init__(self, channel, scale, **kwargs): 30 | super(SRCNN, self).__init__(scale=scale, channel=channel, name='srcnn') 31 | self.net = Srcnn(channel, kwargs.get('filters', (9, 1, 5))) 32 | self.net(tf.keras.Input([None, None, channel])) 33 | self.opt = tf.keras.optimizers.Adam(learning_rate=1e-4) 34 | 35 | def train(self, inputs, labels, learning_rate=None): 36 | lr_image = inputs[0] 37 | _, H, W, _ = lr_image.shape 38 | bi_image = tf.image.resize(lr_image, [H * self.scale, W * self.scale], 39 | tf.image.ResizeMethod.BICUBIC) 40 | with tf.GradientTape() as tape: 41 | sr = self.net(bi_image) 42 | pixel_loss = tf.reduce_mean(tf.losses.mean_squared_error(labels[0], sr)) 43 | variables = self.trainable_variables() 44 | grads = tape.gradient(pixel_loss, variables) 45 | if learning_rate: 46 | self.opt.learning_rate = learning_rate 47 | self.opt.apply_gradients(zip(grads, variables)) 48 | return { 49 | 'loss': pixel_loss.numpy() 50 | } 51 | 52 | def eval(self, inputs, labels=None, **kwargs): 53 | metrics = {} 54 | lr_image = inputs[0] 55 | _, H, W, _ = lr_image.shape 56 | bi_image = tf.image.resize(lr_image, [H * self.scale, W * self.scale], 57 | tf.image.ResizeMethod.BICUBIC) 58 | sr = self.net(bi_image) 59 | if labels is not None: 60 | metrics['psnr'] = tf.image.psnr(sr, labels[0], 1.0) 61 | step = kwargs.get('epoch') 62 | tf.summary.image('sr', sr, step=step, max_outputs=1) 63 | tf.summary.image('bicubic', bi_image, step=step, max_outputs=1) 64 | tf.summary.image('gt', labels[0], step=step, max_outputs=1) 65 | return [sr.numpy()], metrics 66 | -------------------------------------------------------------------------------- /VSR/Backend/Keras/Models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 30 5 | 6 | import importlib 7 | 8 | __all__ = ['get_model', 'list_supported_models'] 9 | 10 | models = { 11 | # alias: (file, class) 12 | 'srcnn': ('Srcnn', 'SRCNN'), 13 | } 14 | 15 | 16 | def get_model(name): 17 | module = f'.Backend.Keras.Models.{models[name][0]}' 18 | package = 'VSR' 19 | m = importlib.import_module(module, package) 20 | return m.__dict__[models[name][1]] 21 | 22 | 23 | def list_supported_models(): 24 | return models.keys() 25 | -------------------------------------------------------------------------------- /VSR/Backend/Keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 30 5 | 6 | from .. import LOG 7 | 8 | import tensorflow as tf 9 | 10 | ver_major, ver_minor, _ = [int(s) for s in tf.__version__.split('.')] 11 | if ver_major < 2: 12 | LOG.warning("legacy tensorflow 1.x is not verified in keras backend") 13 | 14 | gpus = tf.config.experimental.list_physical_devices('GPU') 15 | if gpus: 16 | # Restrict TensorFlow to only use the first GPU 17 | try: 18 | # Currently, memory growth needs to be the same across GPUs 19 | for gpu in gpus: 20 | tf.config.experimental.set_memory_growth(gpu, True) 21 | except RuntimeError as e: 22 | # Visible devices must be set before GPUs have been initialized 23 | print(e) 24 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Arch/Dense.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 17th 2018 6 | 7 | Architectures of common dense blocks used in SR researches 8 | """ 9 | 10 | from .. import tf 11 | from ..Framework.LayersHelper import Layers 12 | 13 | 14 | def dense_block(layers: Layers, inputs, depth=8, rate=16, out_dims=128, 15 | scope=None, reuse=None): 16 | filters = out_dims - rate * depth 17 | feat = [inputs] 18 | with tf.variable_scope(scope, 'DenseBlock', reuse=reuse): 19 | for _ in range(depth): 20 | filters += rate 21 | x = layers.relu_conv2d(feat[-1], filters, 3) 22 | feat.append(x) 23 | feat[-1] = tf.concat(feat[1:], axis=-1) 24 | return x 25 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Arch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/VSR/Backend/TF/Arch/__init__.py -------------------------------------------------------------------------------- /VSR/Backend/TF/Framework/Noise.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2020 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 25th 2018 6 | 7 | Utility for complicated noise model 8 | Ref: 9 | [1] https://arxiv.org/abs/1807.04686 10 | """ 11 | 12 | from .. import tf 13 | 14 | 15 | def tf_camera_response_function(inputs, crf_table, max_val=1): 16 | """TF implementation of CRF.""" 17 | 18 | with tf.name_scope('CRF'): 19 | inputs_norm = tf.clip_by_value(inputs / max_val, 0, 1) 20 | quant = int(crf_table.shape[0] - 1) 21 | inputs_index = tf.to_int32(inputs_norm * quant) 22 | return tf.nn.embedding_lookup(crf_table, inputs_index) 23 | 24 | 25 | def tf_poisson_noise(inputs, stddev=None, sigma_max=0.16): 26 | with tf.name_scope('PoissonNoise'): 27 | if stddev is None: 28 | stddev = tf.random_uniform(inputs.shape[-1:], maxval=sigma_max) 29 | stddev = tf.reshape(stddev, [1, 1, 1, -1]) 30 | sigma_map = (1 - inputs) * stddev 31 | return tf.random_normal(tf.shape(inputs)) * sigma_map 32 | 33 | 34 | def tf_gaussian_noise(inputs, stddev=None, sigma_max=0.06, channel_wise=True): 35 | with tf.name_scope('GaussianNoise'): 36 | channel = inputs.shape[-1:] if channel_wise else [1] 37 | if stddev is None: 38 | stddev = tf.random_uniform(channel, maxval=sigma_max) 39 | stddev = tf.reshape(stddev, [1, 1, 1, -1]) 40 | noise_map = tf.random_normal(tf.shape(inputs)) * stddev 41 | return noise_map 42 | 43 | 44 | def tf_gaussian_poisson_noise(inputs, stddev_s=None, stddev_c=None, 45 | max_s=0.16, max_c=0.06): 46 | with tf.name_scope('PoissonGaussianNoise'): 47 | ns = tf_poisson_noise(inputs, stddev_s, max_s) 48 | nc = tf_gaussian_noise(inputs, stddev_c, max_c) 49 | return ns + nc 50 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Framework/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/VSR/Backend/TF/Framework/__init__.py -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Crdn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Mar. 20th 2019 6 | 7 | Cascaded Residual Dense Network (NTIRE 2019) 8 | """ 9 | 10 | from .. import tf 11 | from ..Arch.Residual import cascade_rdn 12 | from ..Framework.SuperResolution import SuperResolution 13 | from ..Util import clip_image 14 | 15 | 16 | def _denormalize(inputs): 17 | return (inputs + 0) * 255 18 | 19 | 20 | def _normalize(inputs): 21 | return (inputs - 0) / 255 22 | 23 | 24 | class CRDN(SuperResolution): 25 | """A Cascaded Residual Dense Network""" 26 | 27 | def __init__(self, name='crdn', **kwargs): 28 | super(CRDN, self).__init__(**kwargs) 29 | self.name = name 30 | self.n_cb = 6 31 | 32 | def upsample(self, inputs, skips, method='nearest', scale=2): 33 | with tf.variable_scope(None, f'UpX{scale}'): 34 | c = int(inputs.shape[-1]) 35 | up0 = self.upscale(inputs, method, scale, direct_output=False) 36 | up1 = self.conv2d(up0, c, 3) 37 | fs0 = tf.concat([up1, skips], axis=-1) 38 | fs1 = self.conv2d(fs0, c // 2, 3) 39 | return fs1 40 | 41 | def build_graph(self): 42 | super(CRDN, self).build_graph() 43 | depth = self.n_cb 44 | filters = 32 45 | inputs = _normalize(self.inputs_preproc[-1]) 46 | with tf.variable_scope(self.name): 47 | x = self.conv2d(inputs, filters, 7) 48 | entry = self.conv2d(x, filters, 5) 49 | x_list = [entry] 50 | f = filters 51 | for i in range(depth // 2): 52 | x = cascade_rdn(self, x, depth=3, use_ca=True, filters=f) 53 | x_list.append(x) 54 | f *= 2 55 | x = self.conv2d(x, f, 3, strides=2, name='Down%d' % i) 56 | x = cascade_rdn(self, x, depth=3, use_ca=True, filters=f) 57 | x = cascade_rdn(self, x, depth=3, use_ca=True, filters=f) 58 | for i in range(depth // 2): 59 | f //= 2 60 | x = self.upsample(x, x_list.pop()) 61 | x = cascade_rdn(self, x, depth=3, use_ca=True, filters=f) 62 | 63 | assert len(x_list) == 1, f'length of x_list is not 1: {len(x_list)}' 64 | assert f == filters 65 | x += x_list.pop() 66 | sr = self.conv2d(x, filters, 3) 67 | sr = self.conv2d(sr, self.channel, 3) 68 | 69 | self.outputs.append(_denormalize(sr)) 70 | 71 | def build_loss(self): 72 | with tf.name_scope('Loss'): 73 | l1 = tf.losses.absolute_difference(self.label[-1], self.outputs[-1]) 74 | op = tf.train.AdamOptimizer(self.learning_rate) 75 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 76 | with tf.control_dependencies(update_ops): 77 | opt = op.minimize(l1, self.global_steps) 78 | self.loss.append(opt) 79 | 80 | self.train_metric['l1'] = l1 81 | self.metrics['psnr'] = tf.reduce_mean( 82 | tf.image.psnr(self.label[-1], self.outputs[-1], max_val=255)) 83 | self.metrics['ssim'] = tf.reduce_mean( 84 | tf.image.ssim(self.label[-1], self.outputs[-1], max_val=255)) 85 | 86 | def build_summary(self): 87 | super(CRDN, self).build_summary() 88 | tf.summary.image('sr', clip_image(self.outputs[-1])) 89 | tf.summary.image('lr', clip_image(self.inputs_preproc[-1])) 90 | tf.summary.image('hq', clip_image(self.label[-1])) 91 | 92 | def build_saver(self): 93 | var_g = tf.global_variables(self.name) 94 | misc = tf.global_variables('Loss') + [self.global_steps] 95 | self.savers['misc'] = tf.train.Saver(misc, max_to_keep=1) 96 | self.savers[self.name] = tf.train.Saver(var_g, max_to_keep=1) 97 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/DnCnn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: May 23rd 2018 6 | Updated Date: May 23rd 2018 7 | 8 | Implementing Feed-forward Denoising Convolutional Neural Network 9 | See http://ieeexplore.ieee.org/document/7839189/ 10 | """ 11 | from .. import tf 12 | from ..Framework.SuperResolution import SuperResolution 13 | 14 | 15 | class DnCNN(SuperResolution): 16 | """Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for 17 | Image Denoising 18 | 19 | Args: 20 | layers: number of layers used 21 | """ 22 | 23 | def __init__(self, layers=20, name='dncnn', **kwargs): 24 | self.name = name 25 | self.layers = layers 26 | if 'scale' in kwargs: 27 | kwargs.pop('scale') 28 | super(DnCNN, self).__init__(scale=1, **kwargs) 29 | 30 | def build_graph(self): 31 | super(DnCNN, self).build_graph() # build inputs placeholder 32 | with tf.variable_scope(self.name): 33 | # build layers 34 | x = self.inputs_preproc[-1] / 255 # use channel Y only 35 | x = self.relu_conv2d(x, 64, 3) 36 | for i in range(1, self.layers - 1): 37 | x = self.bn_relu_conv2d(x, 64, 3, use_bias=False) 38 | # the last layer w/o BN and ReLU 39 | x = self.conv2d(x, self.channel, 3) 40 | # residual training 41 | outputs = self.inputs_preproc[-1] / 255 + x 42 | self.outputs.append(outputs * 255) 43 | 44 | def build_loss(self): 45 | with tf.name_scope('loss'): 46 | mse, loss = super(DnCNN, self).build_loss() 47 | self.train_metric['loss'] = loss 48 | self.metrics['mse'] = mse 49 | self.metrics['psnr'] = tf.reduce_mean(tf.image.psnr( 50 | self.label[-1], self.outputs[-1], max_val=255)) 51 | self.metrics['ssim'] = tf.reduce_mean(tf.image.ssim( 52 | self.label[-1], self.outputs[-1], max_val=255)) 53 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Drrn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: June 8th 2018 6 | Updated Date: June 8th 2018 7 | 8 | Image Super-Resolution via Deep Recursive Residual Network (CVPR 2017) 9 | See http://cvlab.cse.msu.edu/pdfs/Tai_Yang_Liu_CVPR2017.pdf 10 | """ 11 | 12 | import logging 13 | 14 | from .. import tf 15 | from ..Framework.SuperResolution import SuperResolution 16 | from ..Util import bicubic_rescale 17 | 18 | LOG = logging.getLogger('VSR.Model.DRRN') 19 | 20 | 21 | class DRRN(SuperResolution): 22 | """Image Super-Resolution via Deep Recursive Residual Network 23 | 24 | Args: 25 | residual_unit: number of residual blocks in one recursion 26 | recursive_block: number of recursions 27 | grad_clip: gradient clip ratio according to the paper 28 | custom_upsample: use --add_custom_callbacks=upsample during fitting, or 29 | use `bicubic_rescale`. TODO: REMOVE IN FUTURE. 30 | """ 31 | 32 | def __init__(self, residual_unit=3, recursive_block=3, 33 | custom_upsample=False, 34 | grad_clip=0.01, name='drrn', **kwargs): 35 | self.ru = residual_unit 36 | self.rb = recursive_block 37 | self.grad_clip = grad_clip 38 | self.do_up = not custom_upsample 39 | self.name = name 40 | super(DRRN, self).__init__(**kwargs) 41 | 42 | def display(self): 43 | super(DRRN, self).display() 44 | LOG.info('Recursive Blocks: %d' % self.rb) 45 | LOG.info('Residual Units: %d' % self.ru) 46 | 47 | def _shared_resblock(self, inputs, **kwargs): 48 | x = self.relu_conv2d(inputs, 128, 3) 49 | for _ in range(self.ru): 50 | x = self.resblock(x, 128, 3, reuse=tf.AUTO_REUSE, name='Res') 51 | return x 52 | 53 | def build_graph(self): 54 | super(DRRN, self).build_graph() 55 | with tf.variable_scope(self.name): 56 | x = self.inputs_preproc[-1] 57 | if self.do_up: 58 | x = bicubic_rescale(self.inputs_preproc[-1], self.scale) 59 | bic = x 60 | for _ in range(self.rb): 61 | x = self._shared_resblock(x) 62 | x = self.conv2d(x, self.channel, 3) 63 | self.outputs.append(x + bic) 64 | 65 | def build_loss(self): 66 | with tf.name_scope('loss'): 67 | y_true = self.label[-1] 68 | y_pred = self.outputs[-1] 69 | mse = tf.losses.mean_squared_error(y_true, y_pred) 70 | loss = tf.add_n([mse] + tf.losses.get_regularization_losses()) 71 | opt = tf.train.AdamOptimizer(self.learning_rate) 72 | if self.grad_clip > 0: 73 | grads_and_vars = [] 74 | for grad, var in opt.compute_gradients(loss): 75 | grads_and_vars.append(( 76 | tf.clip_by_value( 77 | grad, 78 | -self.grad_clip / self.learning_rate, 79 | self.grad_clip / self.learning_rate), 80 | var)) 81 | op = opt.apply_gradients(grads_and_vars, self.global_steps) 82 | else: 83 | op = opt.minimize(loss, self.global_steps) 84 | self.loss.append(op) 85 | 86 | self.train_metric['loss'] = loss 87 | self.metrics['mse'] = mse 88 | self.metrics['psnr'] = tf.reduce_mean( 89 | tf.image.psnr(y_true, y_pred, 255)) 90 | self.metrics['ssim'] = tf.reduce_mean( 91 | tf.image.ssim(y_true, y_pred, 255)) 92 | 93 | def build_saver(self): 94 | self.savers[self.name] = tf.train.Saver(tf.global_variables(self.name), 95 | max_to_keep=1) 96 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Edsr.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: June 15th 2018 6 | Updated Date: June 15th 2018 7 | 8 | Enhanced Deep Residual Networks for Single Image Super-Resolution (CVPR 2017) 9 | See https://arxiv.org/abs/1707.02921 10 | """ 11 | 12 | from .. import tf 13 | from ..Framework.SuperResolution import SuperResolution 14 | 15 | 16 | class EDSR(SuperResolution): 17 | """Enhanced Deep Residual Networks for Single Image Super-Resolution 18 | 19 | Args: 20 | layers: number of residual blocks 21 | filters: number of filters in each conv2d 22 | clip: feature value clip ratio in each residual block 23 | """ 24 | 25 | def __init__(self, layers=32, filters=256, clip=0.1, name='edsr', **kwargs): 26 | self.layers = layers 27 | self.filters = filters 28 | self.clip = clip 29 | self.name = name 30 | super(EDSR, self).__init__(**kwargs) 31 | 32 | def build_graph(self): 33 | super(EDSR, self).build_graph() 34 | with tf.variable_scope(self.name): 35 | fe = self.conv2d(self.inputs_preproc[-1], self.filters, 3) 36 | x = fe 37 | for _ in range(self.layers): 38 | with tf.variable_scope(None, 'ResBlock'): 39 | x_old = x 40 | x = self.relu_conv2d(x, self.filters, 3) 41 | x = self.conv2d(x, self.filters, 3) 42 | x = x * self.clip + x_old 43 | x = self.conv2d(x, self.filters, 3) 44 | x += fe 45 | x = self.upscale(x, direct_output=False) 46 | x = self.conv2d(x, self.channel, 3) 47 | self.outputs.append(x) 48 | 49 | def build_loss(self): 50 | with tf.name_scope('loss'): 51 | opt = tf.train.AdamOptimizer(self.learning_rate) 52 | mse = tf.losses.mean_squared_error(self.label[-1], self.outputs[-1]) 53 | mae = tf.losses.absolute_difference(self.label[-1], self.outputs[-1]) 54 | loss = tf.add_n([mae] + tf.losses.get_regularization_losses(), 55 | name='total_loss') 56 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 57 | with tf.control_dependencies(update_ops): 58 | self.loss.append(opt.minimize(loss, self.global_steps)) 59 | 60 | self.train_metric['loss'] = loss 61 | self.metrics['mse'] = mse 62 | self.metrics['mae'] = mae 63 | self.metrics['psnr'] = tf.reduce_mean( 64 | tf.image.psnr(self.label[-1], self.outputs[-1], max_val=255)) 65 | self.metrics['ssim'] = tf.reduce_mean( 66 | tf.image.ssim(self.label[-1], self.outputs[-1], max_val=255)) 67 | 68 | def build_summary(self): 69 | super(EDSR, self).build_summary() 70 | tf.summary.image('SR', self.outputs[-1], 1) 71 | 72 | def build_saver(self): 73 | self.savers.update({ 74 | self.name: tf.train.Saver(tf.trainable_variables(self.name), 75 | max_to_keep=1) 76 | }) 77 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Espcn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: May 12th 2018 6 | Updated Date: May 25th 2018 7 | 8 | Efficient Sub-Pixel Convolutional Neural Network 9 | Ref https://arxiv.org/abs/1609.05158 10 | """ 11 | from VSR.Util import to_list 12 | from .. import tf 13 | from ..Framework.SuperResolution import SuperResolution 14 | 15 | 16 | def _normalize(x): 17 | return x / 127.5 - 1 18 | 19 | 20 | def _denormalize(x): 21 | return (x + 1) * 127.5 22 | 23 | 24 | class ESPCN(SuperResolution): 25 | """Efficient Sub-Pixel Convolutional Neural Network. 26 | 27 | Args: 28 | layers: layer number of the network 29 | filters: a tuple of integer, representing each layer's filters 30 | kernel: a tuple of integer, representing each layer's kernel size 31 | """ 32 | 33 | def __init__(self, layers=3, filters=(64, 32), kernel=(5, 3, 3), 34 | name='espcn', **kwargs): 35 | super(ESPCN, self).__init__(**kwargs) 36 | self.name = name 37 | self.layers = layers 38 | self.filters = to_list(filters, layers - 1) 39 | self.kernel_size = to_list(kernel, layers) 40 | if len(self.kernel_size) < self.layers: 41 | self.kernel_size += to_list( 42 | kernel[-1], self.layers - len(self.kernel_size)) 43 | 44 | def build_graph(self): 45 | super(ESPCN, self).build_graph() 46 | with tf.variable_scope(self.name): 47 | x = _normalize(self.inputs_preproc[-1]) 48 | for f, k in zip(self.filters, self.kernel_size): 49 | x = self.tanh_conv2d(x, f, k, kernel_initializer='torch') 50 | x = self.upscale(x, 'espcn', direct_output=True, 51 | kernel_initializer='torch') 52 | self.outputs.append(_denormalize(x)) 53 | 54 | def build_loss(self): 55 | with tf.name_scope('loss'): 56 | mse, loss = super(ESPCN, self).build_loss() 57 | self.train_metric['loss'] = loss 58 | self.metrics['mse'] = mse 59 | self.metrics['psnr'] = tf.reduce_mean( 60 | tf.image.psnr(self.label[-1], self.outputs[-1], max_val=255)) 61 | self.metrics['ssim'] = tf.reduce_mean( 62 | tf.image.ssim(self.label[-1], self.outputs[-1], max_val=255)) 63 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/FFDNet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/25 下午2:13 5 | 6 | from .. import tf 7 | from ..Framework.SuperResolution import SuperResolution 8 | 9 | 10 | class FFDNet(SuperResolution): 11 | """FFDNet: Toward a Fast and Flexible Solution for CNN-Based Image Denoising 12 | By Kai Zhang. (IEEE TIP 2018) 13 | 14 | Args: 15 | sigma: in training phase, this is the max sigma level added to clean images, 16 | in testing phase, this is input noise level, correspond to pixel [0, 255]. 17 | space_down: block size for space-to-depth (default 2, same as paper's). 18 | layers: convolutional layers used in the network. 19 | training: set to false when evaluating. 20 | """ 21 | 22 | def __init__(self, sigma, space_down=2, layers=15, training=True, 23 | name='ffdnet', **kwargs): 24 | self.name = name 25 | self.sigma = sigma 26 | self.space_down = space_down 27 | self.layers = layers 28 | self.training = training 29 | if 'scale' in kwargs: 30 | kwargs.pop('scale') 31 | super(FFDNet, self).__init__(scale=1, **kwargs) 32 | 33 | def build_graph(self): 34 | super(FFDNet, self).build_graph() # build inputs placeholder 35 | with tf.variable_scope(self.name): 36 | # build layers 37 | inputs = self.inputs_preproc[-1] / 255 38 | if self.training: 39 | sigma = tf.random_uniform((), maxval=self.sigma / 255) 40 | inputs += tf.random_normal(tf.shape(inputs)) * sigma 41 | else: 42 | sigma = self.sigma / 255 43 | inputs = tf.space_to_depth(inputs, block_size=self.space_down) 44 | noise_map = tf.ones_like(inputs)[..., 0:1] * sigma 45 | x = tf.concat([inputs, noise_map], axis=-1) 46 | x = self.relu_conv2d(x, 64, 3) 47 | for i in range(1, self.layers - 1): 48 | x = self.bn_relu_conv2d(x, 64, 3, use_bias=False) 49 | # the last layer w/o BN and ReLU 50 | x = self.conv2d(x, self.channel * self.space_down ** 2, 3) 51 | denoised = tf.depth_to_space(x, block_size=self.space_down) 52 | self.outputs.append(denoised * 255) 53 | 54 | def build_loss(self): 55 | with tf.name_scope('loss'): 56 | mse, loss = super(FFDNet, self).build_loss() 57 | self.train_metric['loss'] = loss 58 | self.metrics['mse'] = mse 59 | self.metrics['psnr'] = tf.reduce_mean(tf.image.psnr( 60 | self.label[-1], self.outputs[-1], max_val=255)) 61 | self.metrics['ssim'] = tf.reduce_mean(tf.image.ssim( 62 | self.label[-1], self.outputs[-1], max_val=255)) 63 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/LapSrn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: May 12th 2018 6 | Updated Date: May 25th 2018 7 | 8 | Deep Laplacian Pyramid Networks 9 | Ref http://vllab.ucmerced.edu/wlai24/LapSRN 10 | """ 11 | import numpy as np 12 | 13 | from VSR.Util import to_list 14 | from .. import tf 15 | from ..Framework.SuperResolution import SuperResolution 16 | from ..Util import bicubic_rescale 17 | 18 | 19 | class LapSRN(SuperResolution): 20 | """Deep Laplacian Pyramid Networks for Fast and Accurate Super-Resolution 21 | 22 | Args: 23 | layers: number of layers in each pyramid level 24 | epsilon: used in charbonnier loss function 25 | """ 26 | 27 | def __init__(self, layers, epsilon=1e-3, name='lapsrn', **kwargs): 28 | super(LapSRN, self).__init__(**kwargs) 29 | self.epsilon2 = epsilon ** 2 30 | self.name = name 31 | s0, s1 = self.scale 32 | if np.any(np.log2([s0, s1]) != np.round(np.log2([s0, s1]))): 33 | raise ValueError('Scale factor must be pow of 2.' 34 | 'Error: scale={},{}'.format(s0, s1)) 35 | assert s0 == s1 36 | self.level = int(np.log2(s0)) 37 | self.layers = to_list(layers, self.level) 38 | 39 | def build_graph(self): 40 | super(LapSRN, self).build_graph() 41 | with tf.variable_scope(self.name): 42 | x = self.inputs_preproc[-1] 43 | residual = [] 44 | with tf.variable_scope('FeatureExtraction'): 45 | for lv in range(self.level): 46 | for _ in range(self.layers[lv] - 1): 47 | x = self.leaky_conv2d(x, 64, 3) 48 | x = self.deconv2d(x, 64, 4, 2, activation='lrelu') 49 | x = self.conv2d(x, self.channel, 3) 50 | residual.append(x) 51 | with tf.name_scope('Reconstruction'): 52 | y = self.inputs_preproc[-1] 53 | _s = 2 54 | for res in residual: 55 | sr = bicubic_rescale(y, _s) + res 56 | _s *= 2 57 | self.outputs.append(sr) 58 | self.outputs.reverse() 59 | 60 | def build_loss(self): 61 | with tf.name_scope('loss'): 62 | y_true = [self.label[-1]] 63 | for _ in range(1, self.level): 64 | y_true.append(bicubic_rescale(y_true[-1], 0.5)) 65 | charbonnier = [] 66 | mse = [] 67 | for x, y in zip(self.outputs, y_true): 68 | charbonnier.append( 69 | tf.reduce_mean(tf.sqrt(tf.square(x - y) + self.epsilon2))) 70 | mse.append(tf.reduce_mean(tf.squared_difference(y, x))) 71 | charbonnier_loss = tf.reduce_mean(charbonnier) 72 | loss = tf.add_n( 73 | [charbonnier_loss] + tf.losses.get_regularization_losses()) 74 | opt = tf.train.AdamOptimizer(self.learning_rate) 75 | self.loss.append(opt.minimize(loss, self.global_steps)) 76 | 77 | self.train_metric['loss'] = loss 78 | self.train_metric['charbonnier_loss'] = charbonnier_loss 79 | for i in range(len(mse)): 80 | self.metrics['mse_x{}'.format(2 ** (i + 1))] = mse[i] 81 | self.metrics['psnr_x{}'.format(2 ** (i + 1))] = 10 * tf.log( 82 | 255 ** 2 / mse[i]) / tf.log(10.0) 83 | 84 | def build_summary(self): 85 | super(LapSRN, self).build_summary() 86 | tf.summary.image('SR', self.outputs[-1], 1) 87 | 88 | def build_saver(self): 89 | self.savers[self.name] = tf.train.Saver(tf.global_variables(self.name), 90 | max_to_keep=1) -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Msrn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Intel Corp. 2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 14th 2018 6 | 7 | Multi-scale Residual Network for Image Super-Resolution 8 | See http://openaccess.thecvf.com/content_ECCV_2018/papers/Juncheng_Li_Multi-scale_Residual_Network_ECCV_2018_paper.pdf 9 | """ 10 | 11 | from .. import tf 12 | from ..Arch.Residual import msrb 13 | from ..Framework.SuperResolution import SuperResolution 14 | 15 | 16 | def _normalize(inputs): 17 | rgb_mean = (0.4488, 0.4371, 0.4040) 18 | return inputs / 255 - rgb_mean 19 | 20 | 21 | def _denormalize(inputs): 22 | rgb_mean = (0.4488, 0.4371, 0.4040) 23 | return (inputs + rgb_mean) * 255 24 | 25 | 26 | class MSRN(SuperResolution): 27 | """Multi-scale Residual Network for Image Super-Resolution 28 | 29 | Args: 30 | n_blocks: number of MSRB blocks. 31 | """ 32 | 33 | def __init__(self, n_blocks=8, name='msrn', **kwargs): 34 | super(MSRN, self).__init__(**kwargs) 35 | self.name = name 36 | self.blocks = n_blocks 37 | 38 | def build_graph(self): 39 | super(MSRN, self).build_graph() 40 | inputs_norm = _normalize(self.inputs_preproc[-1]) 41 | with tf.variable_scope(self.name): 42 | features = [self.conv2d(inputs_norm, 64, 3)] 43 | for _ in range(self.blocks): 44 | x = features[-1] 45 | features.append(msrb(self, x)) 46 | x = self.conv2d(tf.concat(features, -1), 64, 1) 47 | x = self.upscale(x, direct_output=False) 48 | sr = self.conv2d(x, self.channel, 3) 49 | self.outputs.append(_denormalize(sr)) 50 | 51 | def build_loss(self): 52 | label_norm = _normalize(self.label[-1]) 53 | sr = _normalize(self.outputs[-1]) 54 | with tf.name_scope('Loss'): 55 | l1 = tf.losses.absolute_difference(label_norm, sr) 56 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 57 | with tf.control_dependencies(update_ops): 58 | op = tf.train.AdamOptimizer(self.learning_rate) 59 | op = op.minimize(l1, self.global_steps) 60 | self.loss.append(op) 61 | 62 | self.train_metric['l1'] = l1 63 | self.metrics['psnr'] = tf.reduce_mean(tf.image.psnr( 64 | self.label[-1], self.outputs[-1], max_val=255)) 65 | self.metrics['ssim'] = tf.reduce_mean(tf.image.ssim( 66 | self.label[-1], self.outputs[-1], max_val=255)) 67 | 68 | def build_summary(self): 69 | super(MSRN, self).build_summary() 70 | tf.summary.image('sr', self.outputs[0]) 71 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Nlrn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Intel Corp. 2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 19th 2018 6 | 7 | Non-Local Recurrent Network for Image Restoration 8 | See https://arxiv.org/abs/1806.02919 9 | """ 10 | 11 | import logging 12 | 13 | from .. import tf, tfc 14 | from ..Arch.Residual import non_local 15 | from ..Framework.SuperResolution import SuperResolution 16 | 17 | LOG = logging.getLogger('VSR.Model.NLRN') 18 | 19 | 20 | def _denormalize(inputs): 21 | return (inputs + 0) * 255 22 | 23 | 24 | def _normalize(inputs): 25 | return inputs / 255 26 | 27 | 28 | class NLRN(SuperResolution): 29 | """Non-Local Recurrent Network for Image Restoration (NIPS 2018) 30 | 31 | """ 32 | 33 | def __init__(self, recurrents=12, clip=2.5, name='nlrn', **kwargs): 34 | super(NLRN, self).__init__(**kwargs) 35 | self.name = name 36 | self.recurrents = recurrents 37 | self.clip = clip 38 | self.filters = kwargs.get('filters', 128) 39 | 40 | def display(self): 41 | LOG.info(f"Recurrents: {self.recurrents}") 42 | 43 | def rnn(self, x, y): 44 | with tf.variable_scope('RNN', reuse=tf.AUTO_REUSE): 45 | x = self.batch_norm(x, self.training_phase) 46 | x = tf.nn.relu(x) 47 | x = non_local(self, x, self.filters, scaling=2) 48 | 49 | x = self.batch_norm(x, self.training_phase) 50 | x = self.bn_relu_conv2d(x, self.filters, 3) 51 | x = self.conv2d(x, self.filters, 3) 52 | return x + y 53 | 54 | def build_graph(self): 55 | super(NLRN, self).build_graph() 56 | with tf.variable_scope(self.name): 57 | inputs_norm = _normalize(self.inputs_preproc[-1]) 58 | init_feat = self.batch_norm(inputs_norm, self.training_phase) 59 | x = init_feat = self.conv2d(init_feat, self.filters, 3) 60 | for _ in range(self.recurrents): 61 | x = self.rnn(x, init_feat) 62 | sr = self.batch_norm(x, self.training_phase) 63 | sr = tf.nn.relu(sr) 64 | sr = self.conv2d(sr, self.channel, 3) 65 | sr += inputs_norm 66 | 67 | self.outputs.append(_denormalize(sr)) 68 | 69 | def build_loss(self): 70 | with tf.name_scope('Loss'): 71 | mse = tf.losses.mean_squared_error(self.outputs[-1], self.label[-1]) 72 | 73 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 74 | with tf.control_dependencies(update_ops): 75 | opt = tf.train.AdadeltaOptimizer(self.learning_rate) 76 | grad = opt.compute_gradients(mse) 77 | grad_clip = tfc.training.clip_gradient_norms( 78 | grad, self.clip) 79 | op = opt.apply_gradients(grad_clip, self.global_steps) 80 | self.loss.append(op) 81 | 82 | self.train_metric['mse'] = mse 83 | self.metrics['psnr'] = tf.reduce_mean( 84 | tf.image.psnr(self.label[0], self.outputs[0], 255)) 85 | self.metrics['ssim'] = tf.reduce_mean( 86 | tf.image.ssim(self.label[0], self.outputs[0], 255)) 87 | 88 | def build_summary(self): 89 | super(NLRN, self).build_summary() 90 | tf.summary.image('SR', self.outputs[-1], 1) 91 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Rdn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: May 24th 2018 6 | Updated Date: May 25th 2018 7 | 8 | Architecture of Residual Dense Network (CVPR 2018) 9 | See https://arxiv.org/abs/1802.08797 10 | """ 11 | 12 | from .. import tf 13 | from ..Framework.SuperResolution import SuperResolution 14 | 15 | 16 | class ResidualDenseNetwork(SuperResolution): 17 | """Residual Dense Network for Image Super-Resolution 18 | 19 | Args: 20 | global_filters: filters used in shallow feature extraction and global fusion 21 | rdb_blocks: number of residual dense blocks 22 | rdb_conv: number of conv2d layers in each RDB 23 | rdb_filters: number of filters in RDB conv2d 24 | 25 | NOTE: total conv2d layers := `rdb_blocks` * `rdb_conv` + 5 + Upscale 26 | """ 27 | 28 | def __init__(self, global_filters=64, rdb_blocks=10, rdb_conv=6, 29 | rdb_filters=64, name='rdn', **kwargs): 30 | super(ResidualDenseNetwork, self).__init__(**kwargs) 31 | self.name = name 32 | self.gfilter = global_filters 33 | self.block = rdb_blocks 34 | self.conv = rdb_conv 35 | self.growth_rate = rdb_filters 36 | 37 | def _rdb(self, inputs, **kwargs): 38 | """Make Residual Dense Block 39 | 40 | Args: 41 | inputs: input features 42 | """ 43 | with tf.variable_scope(kwargs.get('name'), 'ResDenseBlock'): 44 | filters, conv = self.growth_rate, self.conv 45 | x = [inputs] 46 | x += [self.relu_conv2d(x[-1], filters, 3)] 47 | for i in range(1, conv): 48 | x += [self.relu_conv2d(tf.concat(x, axis=-1), filters, 3)] 49 | # 1x1 conv 50 | local_fusion = self.conv2d(tf.concat(x, axis=-1), filters, 1) 51 | # local residual learning 52 | outputs = inputs + local_fusion 53 | return outputs 54 | 55 | def build_graph(self): 56 | super(ResidualDenseNetwork, self).build_graph() 57 | with tf.variable_scope(self.name): 58 | x = self.inputs_preproc[-1] 59 | # shallow feature extraction 60 | # NOTE: no activation 61 | with tf.variable_scope('ShallowFeature'): 62 | sf0 = self.conv2d(x, self.gfilter, 3) 63 | sf1 = self.conv2d(sf0, self.gfilter, 3) 64 | with tf.variable_scope('ResBlocks'): 65 | F = [sf1] 66 | for i in range(self.block): 67 | F += [self._rdb(F[-1])] 68 | with tf.variable_scope('GlobalFusion'): 69 | gf0 = self.conv2d(tf.concat(F[1:], axis=-1), self.gfilter, 1) 70 | gf1 = self.conv2d(gf0, self.gfilter, 3) 71 | dense_feature = sf0 + gf1 72 | # use pixel shift in ESPCN to upscale 73 | upscaled = self.upscale(dense_feature, direct_output=False) 74 | hr = self.conv2d(upscaled, self.channel, 3) 75 | self.outputs.append(hr) 76 | 77 | def build_loss(self): 78 | """In paper, authors use L1 loss instead of MSE error. Claimed a better perf.""" 79 | with tf.name_scope('loss'): 80 | y_true = self.label[-1] 81 | y_pred = self.outputs[-1] 82 | mae = tf.losses.absolute_difference(y_true, y_pred) 83 | mse = tf.losses.mean_squared_error(y_true, y_pred) 84 | opt = tf.train.AdamOptimizer(self.learning_rate) 85 | loss = tf.add_n([mae] + tf.losses.get_regularization_losses()) 86 | self.loss.append(opt.minimize(loss, self.global_steps)) 87 | self.train_metric['loss'] = loss 88 | self.metrics['mse'] = mse 89 | self.metrics['mae'] = mae 90 | self.metrics['psnr'] = tf.reduce_mean( 91 | tf.image.psnr(y_true, self.outputs[-1], 255)) 92 | self.metrics['ssim'] = tf.reduce_mean( 93 | tf.image.ssim(y_true, self.outputs[-1], 255)) 94 | 95 | def build_saver(self): 96 | self.savers[self.name] = tf.train.Saver(tf.global_variables(self.name), 97 | max_to_keep=1) 98 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/SRDenseNet.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Intel Corp. 2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 19th 2018 6 | 7 | Image Super-Resolution Using Dense Skip Connections 8 | See http://openaccess.thecvf.com/content_ICCV_2017/papers/Tong_Image_Super-Resolution_Using_ICCV_2017_paper.pdf 9 | """ 10 | 11 | from .. import tf 12 | from ..Arch import Dense 13 | from ..Framework.SuperResolution import SuperResolution 14 | 15 | 16 | def _denormalize(inputs): 17 | return (inputs + 0) * 255 18 | 19 | 20 | def _normalize(inputs): 21 | return inputs / 255 22 | 23 | 24 | class SRDenseNet(SuperResolution): 25 | """Image Super-Resolution Using Dense Skip Connections. 26 | Args: 27 | n_blocks: number of dense blocks. 28 | """ 29 | 30 | def __init__(self, name='srdensenet', n_blocks=8, **kwargs): 31 | super(SRDenseNet, self).__init__(**kwargs) 32 | self.name = name 33 | self.n_blocks = n_blocks 34 | 35 | def build_graph(self): 36 | super(SRDenseNet, self).build_graph() 37 | with tf.variable_scope(self.name): 38 | inputs_norm = _normalize(self.inputs_preproc[-1]) 39 | feat = [self.conv2d(inputs_norm, 64, 3)] 40 | for i in range(self.n_blocks): 41 | feat.append(Dense.dense_block(self, feat[-1])) 42 | bottleneck = self.conv2d(tf.concat(feat, -1), 256, 1) 43 | sr = self.upscale(bottleneck, 'deconv', direct_output=False) 44 | sr = self.conv2d(sr, self.channel, 3) 45 | self.outputs.append(_denormalize(sr)) 46 | 47 | def build_loss(self): 48 | mse, loss = super(SRDenseNet, self).build_loss() 49 | self.train_metric['mse'] = mse 50 | self.train_metric['loss'] = loss 51 | self.metrics['psnr'] = tf.reduce_mean(tf.image.psnr( 52 | self.label[-1], self.outputs[-1], max_val=255)) 53 | self.metrics['ssim'] = tf.reduce_mean(tf.image.ssim( 54 | self.label[-1], self.outputs[-1], max_val=255)) 55 | 56 | def build_summary(self): 57 | super(SRDenseNet, self).build_summary() 58 | tf.summary.image('sr', self.outputs[-1]) 59 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Srcnn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: May 8th 2018 6 | Updated Date: May 25th 2018 7 | 8 | SRCNN mainly for framework tests (ECCV 2014) 9 | Ref https://arxiv.org/abs/1501.00092 10 | """ 11 | 12 | from VSR.Util import to_list 13 | from .. import tf 14 | from ..Framework.SuperResolution import SuperResolution 15 | from ..Util import bicubic_rescale 16 | 17 | 18 | class SRCNN(SuperResolution): 19 | """Image Super-Resolution Using Deep Convolutional Networks 20 | 21 | Args: 22 | layers: number layers to use 23 | filters: number of filters of conv2d(s) 24 | kernel: a tuple of integer, representing kernel size of each layer, 25 | can also be one integer to specify the same size 26 | custom_upsample: use --add_custom_callbacks=upsample during fitting, or 27 | use `bicubic_rescale`. TODO: REMOVE IN FUTURE. 28 | """ 29 | 30 | def __init__(self, layers=3, filters=64, kernel=(9, 5, 5), 31 | custom_upsample=False, 32 | name='srcnn', **kwargs): 33 | super(SRCNN, self).__init__(**kwargs) 34 | self.name = name 35 | self.do_up = not custom_upsample 36 | self.layers = layers 37 | self.filters = filters 38 | self.kernel_size = to_list(kernel) 39 | if len(self.kernel_size) < self.layers: 40 | self.kernel_size += to_list(kernel[-1], 41 | self.layers - len(self.kernel_size)) 42 | 43 | def build_graph(self): 44 | super(SRCNN, self).build_graph() 45 | with tf.variable_scope(self.name): 46 | x = self.inputs_preproc[-1] 47 | if self.do_up: 48 | x = bicubic_rescale(x, self.scale) 49 | f = self.filters 50 | ks = self.kernel_size 51 | x = self.relu_conv2d(x, f, ks[0]) 52 | for i in range(1, self.layers - 1): 53 | x = self.relu_conv2d(x, f, ks[i]) 54 | x = self.conv2d(x, self.channel, ks[-1]) 55 | self.outputs.append(x) 56 | 57 | def build_loss(self): 58 | with tf.name_scope('loss'): 59 | y_pred = self.outputs[-1] 60 | y_true = self.label[-1] 61 | mse = tf.losses.mean_squared_error(y_true, y_pred) 62 | loss = tf.losses.get_total_loss() 63 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 64 | with tf.control_dependencies(update_ops): 65 | opt = tf.train.AdamOptimizer(self.learning_rate) 66 | self.loss.append(opt.minimize(loss, self.global_steps)) 67 | 68 | self.train_metric['loss'] = loss 69 | self.metrics['mse'] = mse 70 | self.metrics['psnr'] = tf.reduce_mean( 71 | tf.image.psnr(self.label[-1], self.outputs[-1], max_val=255)) 72 | self.metrics['ssim'] = tf.reduce_mean( 73 | tf.image.ssim(self.label[-1], self.outputs[-1], max_val=255)) 74 | 75 | def build_summary(self): 76 | super(SRCNN, self).build_summary() 77 | tf.summary.image('SR', self.outputs[-1], 1) 78 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/Vdsr.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: June 5th 2018 6 | Updated Date: June 15th 2018 7 | 8 | Accurate Image Super-Resolution Using Very Deep Convolutional Networks 9 | See https://arxiv.org/abs/1511.04587 10 | """ 11 | 12 | from .. import tf 13 | from ..Framework.SuperResolution import SuperResolution 14 | from ..Util import bicubic_rescale 15 | 16 | 17 | class VDSR(SuperResolution): 18 | """Accurate Image Super-Resolution Using Very Deep Convolutional Networks 19 | 20 | Args: 21 | layers: number of conv2d layers 22 | filters: number of filters in conv2d(s) 23 | custom_upsample: use --add_custom_callbacks=upsample during fitting, or 24 | use `bicubic_rescale`. TODO: REMOVE IN FUTURE. 25 | """ 26 | 27 | def __init__(self, layers=20, filters=64, custom_upsample=False, 28 | name='vdsr', **kwargs): 29 | self.layers = layers 30 | self.filters = filters 31 | self.do_up = not custom_upsample 32 | self.name = name 33 | super(VDSR, self).__init__(**kwargs) 34 | 35 | def build_graph(self): 36 | super(VDSR, self).build_graph() 37 | with tf.variable_scope(self.name): 38 | # bicubic upscale 39 | x = bic = self.inputs_preproc[-1] 40 | if self.do_up: 41 | x = bic = bicubic_rescale(self.inputs_preproc[-1], self.scale) 42 | for _ in range(self.layers - 1): 43 | x = self.relu_conv2d(x, self.filters, 3) 44 | x = self.conv2d(x, self.channel, 3) 45 | self.outputs.append(x + bic) 46 | 47 | def build_loss(self): 48 | with tf.name_scope('loss'): 49 | mae = tf.losses.absolute_difference(self.label[-1], self.outputs[-1]) 50 | loss = tf.losses.get_total_loss() 51 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 52 | with tf.control_dependencies(update_ops): 53 | opt = tf.train.AdamOptimizer(self.learning_rate) 54 | self.loss.append(opt.minimize(loss, self.global_steps)) 55 | 56 | self.train_metric['loss'] = loss 57 | self.metrics['mae'] = mae 58 | self.metrics['psnr'] = tf.reduce_mean( 59 | tf.image.psnr(self.label[-1], self.outputs[-1], max_val=255)) 60 | self.metrics['ssim'] = tf.reduce_mean( 61 | tf.image.ssim(self.label[-1], self.outputs[-1], max_val=255)) 62 | -------------------------------------------------------------------------------- /VSR/Backend/TF/Models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2018 - 8 - 1 5 | 6 | import importlib 7 | import re 8 | from pathlib import Path 9 | 10 | __all__ = ['get_model', 'list_supported_models'] 11 | 12 | 13 | def auto_search(root): 14 | def _parse_class(file): 15 | obj = [] 16 | key_to_remove = set() 17 | file = Path(file) 18 | with file.open('r', encoding='utf-8') as fd: 19 | line = fd.readline() 20 | while line: 21 | if line.startswith('class'): 22 | if '(SuperResolution)' in line: 23 | try: 24 | classname = re.compile("(?<=class\s)\w+\\b").findall(line)[0] 25 | obj.append(classname) 26 | except IndexError: 27 | print(" [!] class: " + line) 28 | else: 29 | for cls in obj: 30 | if f'({cls})' in line: 31 | try: 32 | classname = re.compile("(?<=class\s)\w+\\b").findall(line)[0] 33 | obj.append(classname) 34 | key_to_remove.add(cls) 35 | except IndexError: 36 | print(" [!] class: " + line) 37 | line = fd.readline() 38 | for key in key_to_remove: 39 | obj.remove(key) 40 | return {file.stem: obj} 41 | 42 | mods = sorted(filter( 43 | lambda x: x.is_file() and not x.stem.startswith('__'), 44 | Path(root).glob('*.py'))) 45 | for _m in mods: 46 | cls = _parse_class(_m) 47 | for k in cls: 48 | if k.lower() in models: 49 | print(" [!] duplicated model names found: " + k) 50 | continue 51 | if len(cls[k]) == 1: 52 | models[k.lower()] = (k, cls[k][0]) 53 | elif len(cls[k]) > 1: 54 | for i in cls[k]: 55 | models[f'{k.lower()}.{i.lower()}'] = (k, i) 56 | 57 | 58 | models = {} 59 | auto_search(Path(__file__).parent) 60 | 61 | 62 | def get_model(name): 63 | module = f'.Backend.TF.Models.{models[name][0]}' 64 | package = 'VSR' 65 | m = importlib.import_module(module, package) 66 | return m.__dict__[models[name][1]] 67 | 68 | 69 | def list_supported_models(): 70 | return models.keys() 71 | -------------------------------------------------------------------------------- /VSR/Backend/TF/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 30 5 | 6 | import tensorflow as tf 7 | 8 | ver_major, ver_minor, _ = [int(s) for s in tf.__version__.split('.')] 9 | if ver_major >= 2: 10 | import tensorflow.compat.v1 as tf 11 | 12 | tf.disable_v2_behavior() 13 | else: 14 | tfc = tf.contrib 15 | if ver_minor >= 15: 16 | import tensorflow.compat.v1 as tf 17 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Framework/Summary.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | try: 7 | # torch >= 1.1.0 8 | from torch.utils.tensorboard import SummaryWriter 9 | except ImportError: 10 | from tensorboardX import SummaryWriter 11 | 12 | _writer_container = {} 13 | 14 | 15 | class Summarizer: 16 | def __init__(self, log_path, key=None): 17 | if key is not None: 18 | self.key = hash(key) 19 | else: 20 | self.key = hash(str(log_path)) 21 | self._logd = log_path 22 | self.writer = SummaryWriter(str(log_path)) 23 | _writer_container[self.key] = self 24 | 25 | def close(self): 26 | self.writer.close() 27 | 28 | def scalar(self, name, x, step=None, collection=None): 29 | if collection is not None: 30 | name = f'{collection}/{name}' 31 | self.writer.add_scalar(name, x, step) 32 | 33 | def image(self, name, image, max=3, step=None, collection=None): 34 | if image.ndimension() == 4: 35 | images = image.split(1, dim=0)[:max] 36 | else: 37 | assert image.ndimension() == 3, \ 38 | f'Dim of image is not 3, which is {image.ndimension()}' 39 | images = [image] 40 | if collection is not None: 41 | name = f'{collection}/{name}' 42 | for i, img in enumerate(images): 43 | self.writer.add_image(f'{name}_{i}', img.squeeze(0), step) 44 | 45 | def tensor(self, name, tensor, max=3, step=None, reshape=None): 46 | assert tensor.ndimension() == 4, \ 47 | f"Support 4-D tensor only! {tensor.ndimension()}" 48 | shape = tensor.shape 49 | 50 | def _placement(t): 51 | if t <= 16: 52 | return 4, t // 4 53 | elif t <= 64: 54 | return 8, t // 8 55 | elif t <= 256: 56 | return 16, t // 16 57 | else: 58 | return 32, t // 32 59 | 60 | if reshape: 61 | col, row = reshape 62 | else: 63 | col, row = _placement(shape[1]) 64 | tensor = tensor.view([shape[0], row, col, shape[2], shape[3]]) 65 | tensor = tensor.transpose(2, 3) 66 | tensor = tensor.view([shape[0], row * shape[2], 1, col * shape[3], 1]) 67 | tensor = tensor.squeeze([2, 4]) 68 | tensor = tensor.unsqueeze(1) 69 | self.image(name, tensor, step, max, collection='features') 70 | 71 | def graph(self, model, *args, **kwargs): 72 | self.writer.add_graph(model, args, **kwargs) 73 | 74 | 75 | def get_writer(key) -> Summarizer: 76 | return _writer_container.get(hash(key)) 77 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Framework/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | __all__ = [ 7 | 'Environment', 8 | 'Summary', 9 | 'Trainer' 10 | ] 11 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Bicubic.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 17 5 | 6 | # Non-trainable bicubic, for performance benchmarking and debugging 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torchvision as tv 11 | 12 | from .Model import SuperResolution 13 | from ..Util.Metrics import psnr 14 | 15 | 16 | class Cubic(nn.Module): 17 | def __init__(self, scale): 18 | super(Cubic, self).__init__() 19 | self.to_pil = tv.transforms.ToPILImage() 20 | self.to_tensor = tv.transforms.ToTensor() 21 | self.scale = scale 22 | 23 | def forward(self, x): 24 | if self.scale == 1: 25 | return x 26 | ret = [] 27 | for img in [i[0] for i in x.split(1, dim=0)]: 28 | img = self.to_pil(img.cpu()) 29 | w = img.width 30 | h = img.height 31 | img = img.resize([w * self.scale, h * self.scale], 3) 32 | img = self.to_tensor(img) 33 | ret.append(img) 34 | return torch.stack(ret).to(x.device) 35 | 36 | 37 | class BICUBIC(SuperResolution): 38 | def __init__(self, scale=4, channel=3, **kwargs): 39 | super(BICUBIC, self).__init__(scale, channel, **kwargs) 40 | self.cubic = Cubic(scale) 41 | self.cri = nn.L1Loss() 42 | 43 | def train(self, inputs, labels, learning_rate=None): 44 | sr = self.cubic(inputs[0]) 45 | loss = self.cri(sr, labels[0]) 46 | return {'l1': loss.detach().cpu().numpy()} 47 | 48 | def eval(self, inputs, labels=None, **kwargs): 49 | metrics = {} 50 | sr = self.cubic(inputs[0]).cpu().detach() 51 | if labels is not None: 52 | metrics['psnr'] = psnr(sr.numpy(), labels[0].cpu().numpy()) 53 | return [sr.numpy()], metrics 54 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 16 5 | 6 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire19/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/20 上午11:47 5 | 6 | import logging 7 | _logger = logging.getLogger("VSR.NTIRE2019") 8 | _logger.info("Top rank models in NTIRE 2019." 9 | "Denoise: sRGB and Real Super-Resolution") 10 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire20/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 28 5 | 6 | import logging 7 | _logger = logging.getLogger("VSR.NTIRE2020") 8 | _logger.info("Top rank models in NTIRE 2020." 9 | "Real World Super-Resolution") 10 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire20/xiaozhong/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 28 5 | 6 | import logging 7 | _logger = logging.getLogger("VSR.RWSR") 8 | _logger.info("LICENSE: RealSR is implemented by Xiaozhong Ji. " 9 | "@xiaozhongji https://github.com/jixiaozhong/RealSR") 10 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire20/xiaozhong/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Apache 2.0. 2 | # Author: Xiaozhong Ji 3 | # Update: 2020 - 5 - 28 4 | 5 | from .discriminator import ( 6 | Discriminator_VGG_128, Discriminator_VGG_256, Discriminator_VGG_512, 7 | NLayerDiscriminator, VGGFeatureExtractor 8 | ) 9 | from .network import RRDBNet 10 | 11 | 12 | #################### 13 | # define network 14 | #################### 15 | 16 | def define_G(which_model='RRDBNet', **opt): 17 | """ 18 | Generator 19 | :param which_model: 20 | :param opt: 21 | :return: 22 | """ 23 | 24 | if which_model == 'RRDBNet': 25 | return RRDBNet(in_nc=opt['in_nc'], out_nc=opt['out_nc'], nf=opt['nf'], 26 | nb=opt['nb']) 27 | else: 28 | raise NotImplementedError(f'Generator model [{which_model}] not recognized') 29 | 30 | 31 | def define_D(which_model='NLayerDiscriminator', **opt): 32 | """ 33 | Discriminator 34 | :param which_model: 35 | :param opt: 36 | :return: 37 | """ 38 | 39 | if which_model == 'discriminator_vgg_128': 40 | netD = Discriminator_VGG_128(in_nc=opt['in_nc'], nf=opt['nf']) 41 | elif which_model == 'discriminator_vgg_256': 42 | netD = Discriminator_VGG_256(in_nc=opt['in_nc'], nf=opt['nf']) 43 | elif which_model == 'discriminator_vgg_512': 44 | netD = Discriminator_VGG_512(in_nc=opt['in_nc'], nf=opt['nf']) 45 | elif which_model == 'NLayerDiscriminator': 46 | netD = NLayerDiscriminator(input_nc=opt['in_nc'], ndf=opt['nf'], 47 | n_layers=opt['nlayer']) 48 | else: 49 | raise NotImplementedError( 50 | f'Discriminator model [{which_model}] not recognized') 51 | return netD 52 | 53 | 54 | def define_F(use_bn=False): 55 | """ 56 | Define Network used for Perceptual Loss 57 | PyTorch pre-trained VGG19-54, before ReLU. 58 | :param use_bn: 59 | :return: 60 | """ 61 | 62 | feature_layer = 49 if use_bn else 34 63 | netF = VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn, 64 | use_input_norm=True) 65 | netF.eval() # No need to train 66 | return netF 67 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire20/xiaozhong/ops/loss.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Apache 2.0. 2 | # Author: Xiaozhong Ji 3 | # Update: 2020 - 5 - 28 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | class CharbonnierLoss(nn.Module): 9 | """Charbonnier Loss (L1)""" 10 | 11 | def __init__(self, eps=1e-6): 12 | super(CharbonnierLoss, self).__init__() 13 | self.eps = eps 14 | 15 | def forward(self, x, y): 16 | diff = x - y 17 | loss = torch.sum(torch.sqrt(diff * diff + self.eps)) 18 | return loss 19 | 20 | 21 | # Define GAN loss: [vanilla | lsgan | wgan-gp] 22 | class GANLoss(nn.Module): 23 | def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0): 24 | super(GANLoss, self).__init__() 25 | self.gan_type = gan_type.lower() 26 | self.real_label_val = real_label_val 27 | self.fake_label_val = fake_label_val 28 | 29 | if self.gan_type == 'gan' or self.gan_type == 'ragan': 30 | self.loss = nn.BCEWithLogitsLoss() 31 | elif self.gan_type == 'lsgan': 32 | self.loss = nn.MSELoss() 33 | elif self.gan_type == 'wgan-gp': 34 | 35 | def wgan_loss(input, target): 36 | # target is boolean 37 | return -1 * input.mean() if target else input.mean() 38 | 39 | self.loss = wgan_loss 40 | else: 41 | raise NotImplementedError( 42 | 'GAN type [{:s}] is not found'.format(self.gan_type)) 43 | 44 | def get_target_label(self, input, target_is_real): 45 | if self.gan_type == 'wgan-gp': 46 | return target_is_real 47 | if target_is_real: 48 | return torch.empty_like(input).fill_(self.real_label_val) 49 | else: 50 | return torch.empty_like(input).fill_(self.fake_label_val) 51 | 52 | def forward(self, input, target_is_real): 53 | target_label = self.get_target_label(input, target_is_real) 54 | loss = self.loss(input, target_label) 55 | return loss 56 | 57 | 58 | class GradientPenaltyLoss(nn.Module): 59 | def __init__(self, device=torch.device('cpu')): 60 | super(GradientPenaltyLoss, self).__init__() 61 | self.register_buffer('grad_outputs', torch.Tensor()) 62 | self.grad_outputs = self.grad_outputs.to(device) 63 | 64 | def get_grad_outputs(self, input): 65 | if self.grad_outputs.size() != input.size(): 66 | self.grad_outputs.resize_(input.size()).fill_(1.0) 67 | return self.grad_outputs 68 | 69 | def forward(self, interp, interp_crit): 70 | grad_outputs = self.get_grad_outputs(interp_crit) 71 | grad_interp = torch.autograd.grad(outputs=interp_crit, inputs=interp, 72 | grad_outputs=grad_outputs, 73 | create_graph=True, 74 | retain_graph=True, only_inputs=True)[0] 75 | grad_interp = grad_interp.view(grad_interp.size(0), -1) 76 | grad_interp_norm = grad_interp.norm(2, dim=1) 77 | 78 | loss = ((grad_interp_norm - 1) ** 2).mean() 79 | return loss 80 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Contrib/ntire20/xiaozhong/ops/network.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Apache 2.0. 2 | # Author: Xiaozhong Ji 3 | # Update: 2020 - 5 - 28 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | 10 | def initialize_weights(net_l, scale=1): 11 | if not isinstance(net_l, list): 12 | net_l = [net_l] 13 | for net in net_l: 14 | for m in net.modules(): 15 | if isinstance(m, nn.Conv2d): 16 | nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') 17 | m.weight.data *= scale # for residual block 18 | if m.bias is not None: 19 | m.bias.data.zero_() 20 | elif isinstance(m, nn.Linear): 21 | nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') 22 | m.weight.data *= scale 23 | if m.bias is not None: 24 | m.bias.data.zero_() 25 | elif isinstance(m, nn.BatchNorm2d): 26 | nn.init.constant_(m.weight, 1) 27 | nn.init.constant_(m.bias.data, 0.0) 28 | 29 | 30 | class ResidualDenseBlock_5C(nn.Module): 31 | def __init__(self, nf=64, gc=32, bias=True): 32 | super(ResidualDenseBlock_5C, self).__init__() 33 | # gc: growth channel, i.e. intermediate channels 34 | self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias) 35 | self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias) 36 | self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias) 37 | self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias) 38 | self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias) 39 | self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) 40 | 41 | # initialization 42 | initialize_weights( 43 | [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1) 44 | 45 | def forward(self, x): 46 | x1 = self.lrelu(self.conv1(x)) 47 | x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) 48 | x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) 49 | x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) 50 | x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) 51 | return x5 * 0.2 + x 52 | 53 | 54 | class RRDB(nn.Module): 55 | '''Residual in Residual Dense Block''' 56 | 57 | def __init__(self, nf, gc=32): 58 | super(RRDB, self).__init__() 59 | self.RDB1 = ResidualDenseBlock_5C(nf, gc) 60 | self.RDB2 = ResidualDenseBlock_5C(nf, gc) 61 | self.RDB3 = ResidualDenseBlock_5C(nf, gc) 62 | 63 | def forward(self, x): 64 | out = self.RDB1(x) 65 | out = self.RDB2(out) 66 | out = self.RDB3(out) 67 | return out * 0.2 + x 68 | 69 | 70 | class RRDBNet(nn.Module): 71 | def __init__(self, in_nc, out_nc, nf, nb, gc=32): 72 | super(RRDBNet, self).__init__() 73 | self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) 74 | self.RRDB_trunk = nn.Sequential(*[RRDB(nf=nf, gc=gc) for _ in range(nb)]) 75 | self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 76 | #### upsampling 77 | self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 78 | self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 79 | self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) 80 | self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True) 81 | self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) 82 | 83 | def forward(self, x): 84 | fea = self.conv_first(x) 85 | trunk = self.trunk_conv(self.RRDB_trunk(fea)) 86 | fea = fea + trunk 87 | fea = self.lrelu( 88 | self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest'))) 89 | fea = self.lrelu( 90 | self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest'))) 91 | out = self.conv_last(self.lrelu(self.HRconv(fea))) 92 | return out 93 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Crdn.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019 - 3 - 13 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | from VSR.Util.Utility import to_list 11 | from .Ops.Blocks import CascadeRdn 12 | from .Optim.SISR import L1Optimizer 13 | 14 | 15 | class Upsample(nn.Module): 16 | def __init__(self, channels): 17 | super(Upsample, self).__init__() 18 | in_c, out_c = to_list(channels, 2) 19 | self.c1 = nn.Conv2d(in_c, out_c, 3, 1, 1) 20 | self.c2 = nn.Conv2d(in_c, out_c, 3, 1, 1) 21 | 22 | def forward(self, inputs, skips, scale=2): 23 | up = F.interpolate(inputs, scale_factor=scale) 24 | up = self.c1(up) 25 | con = torch.cat([up, skips], dim=1) 26 | return self.c2(con) 27 | 28 | 29 | class Crdn(nn.Module): 30 | def __init__(self, blocks=(4, 4), **kwargs): 31 | super(Crdn, self).__init__() 32 | self.blocks = to_list(blocks, 2) 33 | 34 | self.entry = nn.Sequential( 35 | nn.Conv2d(3, 32, 7, 1, 3), 36 | nn.Conv2d(32, 32, 5, 1, 2)) 37 | self.exit = nn.Sequential( 38 | nn.Conv2d(32, 32, 3, 1, 1), 39 | nn.Conv2d(32, 3, 3, 1, 1)) 40 | self.down1 = nn.Conv2d(32, 64, 3, 2, 1) 41 | self.down2 = nn.Conv2d(64, 128, 3, 2, 1) 42 | self.up1 = Upsample([128, 64]) 43 | self.up2 = Upsample([64, 32]) 44 | self.cb1 = CascadeRdn(32, 32, 3, True) 45 | self.cb2 = CascadeRdn(64, 64, 3, True) 46 | self.cb3 = CascadeRdn(128, 128, 3, True) 47 | self.cb4 = CascadeRdn(128, 128, 3, True) 48 | self.cb5 = CascadeRdn(64, 64, 3, True) 49 | self.cb6 = CascadeRdn(32, 32, 3, True) 50 | 51 | def forward(self, inputs): 52 | entry = self.entry(inputs) 53 | x1 = self.cb1(entry) 54 | x = self.down1(x1) 55 | x2 = self.cb2(x) 56 | x = self.down2(x2) 57 | x = self.cb3(x) 58 | x = self.cb4(x) 59 | x = self.up1(x, x2) 60 | x = self.cb5(x) 61 | x = self.up2(x, x1) 62 | x = self.cb6(x) 63 | x += entry 64 | out = self.exit(x) 65 | return out 66 | 67 | 68 | class CRDN(L1Optimizer): 69 | def __init__(self, channel=3, scale=1, **kwargs): 70 | self.rsr = Crdn() 71 | super(CRDN, self).__init__(scale=scale, channel=channel, **kwargs) 72 | 73 | def fn(self, x): 74 | return self.rsr(x) 75 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Esrgan.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019 - 3 - 15 5 | 6 | import logging 7 | 8 | import numpy as np 9 | import torch.nn as nn 10 | 11 | from .Ops.Blocks import Activation, EasyConv2d, Rrdb 12 | from .Ops.Discriminator import DCGAN 13 | from .Ops.Scale import Upsample 14 | from .Optim.SISR import PerceptualOptimizer 15 | 16 | _logger = logging.getLogger("VSR.ESRGAN") 17 | _logger.info("LICENSE: ESRGAN is implemented by Xintao Wang. " 18 | "@xinntao https://github.com/xinntao/ESRGAN") 19 | 20 | 21 | class RRDB_Net(nn.Module): 22 | def __init__(self, channel, scale, nf, nb, gc=32): 23 | super(RRDB_Net, self).__init__() 24 | self.head = EasyConv2d(channel, nf, kernel_size=3) 25 | rb_blocks = [ 26 | Rrdb(nf, gc, 5, 0.2, kernel_size=3, 27 | activation=Activation('lrelu', negative_slope=0.2)) 28 | for _ in range(nb)] 29 | LR_conv = EasyConv2d(nf, nf, kernel_size=3) 30 | upsampler = [Upsample(nf, scale, 'nearest', 31 | activation=Activation('lrelu', negative_slope=0.2))] 32 | HR_conv0 = EasyConv2d(nf, nf, kernel_size=3, activation='lrelu', 33 | negative_slope=0.2) 34 | HR_conv1 = EasyConv2d(nf, channel, kernel_size=3) 35 | self.body = nn.Sequential(*rb_blocks, LR_conv) 36 | self.tail = nn.Sequential(*upsampler, HR_conv0, HR_conv1) 37 | 38 | def forward(self, x): 39 | x = self.head(x) 40 | x = self.body(x) + x 41 | x = self.tail(x) 42 | return x 43 | 44 | 45 | class ESRGAN(PerceptualOptimizer): 46 | def __init__(self, channel, scale, patch_size=128, weights=(0.01, 1, 5e-3), 47 | nf=64, nb=23, gc=32, **kwargs): 48 | self.rrdb = RRDB_Net(channel, scale, nf, nb, gc) 49 | super(ESRGAN, self).__init__(scale, channel, 50 | discriminator=DCGAN, 51 | discriminator_kwargs={ 52 | 'channel': channel, 53 | 'scale': scale, 54 | 'num_layers': np.log2(patch_size // 4) * 2, 55 | 'norm': 'BN' 56 | }, 57 | image_weight=weights[0], 58 | feature_weight=weights[1], 59 | gan_weight=weights[2], **kwargs) 60 | 61 | def fn(self, x): 62 | return self.rrdb(x) 63 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Ffdnet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/5/27 下午5:22 5 | 6 | import torch 7 | import torch.nn.functional as F 8 | from torch import nn 9 | 10 | from .Model import SuperResolution 11 | from .Ops.Blocks import EasyConv2d 12 | from .Ops.Scale import SpaceToDepth, Upsample 13 | from ..Framework.Summary import get_writer 14 | from ..Util import Metrics 15 | 16 | 17 | class Net(nn.Module): 18 | def __init__(self, channel, layers, bn, filters=64): 19 | super(Net, self).__init__() 20 | self.spd = SpaceToDepth(2) 21 | body = [EasyConv2d(channel * 4 + 1, filters, 3, activation='relu')] 22 | for i in range(1, layers): 23 | body.append(EasyConv2d(filters, filters, 3, activation='relu', use_bn=bn)) 24 | body += [ 25 | Upsample(filters, 2), 26 | EasyConv2d(filters, channel, 3) 27 | ] 28 | self.body = nn.Sequential(*body) 29 | 30 | def forward(self, x, sigma): 31 | x = self.spd(x) 32 | sig = torch.ones_like(x)[:, 0:1] * sigma 33 | return self.body(torch.cat((x, sig), dim=1)) 34 | 35 | 36 | class FFDNET(SuperResolution): 37 | def __init__(self, scale, channel, n_layers, level, training, **kwargs): 38 | super(FFDNET, self).__init__(scale, channel) 39 | self.ffdnet = Net(channel, n_layers, True) 40 | self.opt = torch.optim.Adam(self.trainable_variables(), 1e-4) 41 | self.level = level / 255 42 | self.is_training = training 43 | 44 | def train(self, inputs, labels, learning_rate=None): 45 | for opt in self.opts.values(): 46 | if learning_rate: 47 | for param_group in opt.param_groups: 48 | param_group["lr"] = learning_rate 49 | lr = inputs[0] 50 | sigma = torch.rand(1, device=lr.device) * 75 / 255 51 | noise = torch.randn_like(lr) * sigma 52 | hr = self.ffdnet((lr + noise).clamp(0, 1), sigma) 53 | loss = F.l1_loss(hr, labels[0]) 54 | self.opt.zero_grad() 55 | loss.backward() 56 | self.opt.step() 57 | return { 58 | 'loss': loss.detach().cpu().numpy() 59 | } 60 | 61 | def eval(self, inputs, labels=None, **kwargs): 62 | metrics = {} 63 | lr = inputs[0] 64 | if self.is_training: 65 | sigma = torch.rand(1, device=lr.device) * 75 / 255 66 | noise = torch.randn_like(lr) * sigma 67 | else: 68 | sigma = self.level 69 | noise = torch.zeros_like(lr) 70 | hr = self.ffdnet((lr + noise).clamp(0, 1), sigma).detach().cpu() 71 | if labels is not None: 72 | metrics['psnr'] = Metrics.psnr(hr, labels[0]) 73 | writer = get_writer(self.name) 74 | if writer is not None: 75 | step = kwargs.get('epoch') 76 | writer.image('gt', labels[0], step=step) 77 | writer.image('clean', hr.clamp(0, 1), step=step) 78 | return [hr.numpy()], metrics 79 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Msrn.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019 - 3 - 15 5 | 6 | import logging 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | from .Ops.Blocks import EasyConv2d, MeanShift 13 | from .Ops.Scale import Upsample 14 | from .Optim.SISR import L1Optimizer 15 | 16 | _logger = logging.getLogger("VSR.MSRN") 17 | _logger.info("LICENSE: MSRN is implemented by Juncheng Li. " 18 | "@MIVRC https://github.com/MIVRC/MSRN-PyTorch") 19 | 20 | 21 | class MSRB(nn.Module): 22 | def __init__(self, n_feats=64): 23 | super(MSRB, self).__init__() 24 | self.conv_3_1 = EasyConv2d(n_feats, n_feats, 3) 25 | self.conv_3_2 = EasyConv2d(n_feats * 2, n_feats * 2, 3) 26 | self.conv_5_1 = EasyConv2d(n_feats, n_feats, 5) 27 | self.conv_5_2 = EasyConv2d(n_feats * 2, n_feats * 2, 5) 28 | self.confusion = nn.Conv2d(n_feats * 4, n_feats, 1, padding=0, stride=1) 29 | 30 | def forward(self, x): 31 | input_1 = x 32 | output_3_1 = F.relu(self.conv_3_1(input_1)) 33 | output_5_1 = F.relu(self.conv_5_1(input_1)) 34 | input_2 = torch.cat([output_3_1, output_5_1], 1) 35 | output_3_2 = F.relu(self.conv_3_2(input_2)) 36 | output_5_2 = F.relu(self.conv_5_2(input_2)) 37 | input_3 = torch.cat([output_3_2, output_5_2], 1) 38 | output = self.confusion(input_3) 39 | output += x 40 | return output 41 | 42 | 43 | class Msrn(nn.Module): 44 | def __init__(self, channel, scale, n_feats, n_blocks, rgb_range): 45 | super(Msrn, self).__init__() 46 | self.n_blocks = n_blocks 47 | # RGB mean for DIV2K 48 | rgb_mean = (0.4488, 0.4371, 0.4040) 49 | self.sub_mean = MeanShift(rgb_mean, True, rgb_range) 50 | # define head module 51 | modules_head = [EasyConv2d(channel, n_feats, 3)] 52 | # define body module 53 | modules_body = nn.ModuleList() 54 | for i in range(n_blocks): 55 | modules_body.append(MSRB(n_feats=n_feats)) 56 | # define tail module 57 | modules_tail = [ 58 | EasyConv2d(n_feats * (self.n_blocks + 1), n_feats, 1), 59 | EasyConv2d(n_feats, n_feats, 3), 60 | Upsample(n_feats, scale), 61 | EasyConv2d(n_feats, channel, 3)] 62 | 63 | self.add_mean = MeanShift(rgb_mean, False, rgb_range) 64 | self.head = nn.Sequential(*modules_head) 65 | self.body = nn.Sequential(*modules_body) 66 | self.tail = nn.Sequential(*modules_tail) 67 | 68 | def forward(self, x): 69 | x = self.sub_mean(x) 70 | x = self.head(x) 71 | res = x 72 | 73 | MSRB_out = [] 74 | for i in range(self.n_blocks): 75 | x = self.body[i](x) 76 | MSRB_out.append(x) 77 | MSRB_out.append(res) 78 | 79 | res = torch.cat(MSRB_out, 1) 80 | x = self.tail(res) 81 | x = self.add_mean(x) 82 | return x 83 | 84 | 85 | class MSRN(L1Optimizer): 86 | def __init__(self, scale, channel, n_feats=64, n_blocks=8, rgb_range=255, 87 | **kwargs): 88 | self.rgb_range = rgb_range 89 | self.msrn = Msrn(channel, scale, n_feats, n_blocks, rgb_range) 90 | super(MSRN, self).__init__(scale, channel, **kwargs) 91 | 92 | def fn(self, x): 93 | return self.msrn(x * self.rgb_range) / self.rgb_range 94 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/NTIRE19.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/16 5 | 6 | from VSR.Util.Config import Config 7 | from .Contrib.ntire19 import denoise, edrn, frn, ran2 8 | from .Optim.SISR import L1Optimizer 9 | 10 | 11 | class EDRN(L1Optimizer): 12 | """EDRN is one candidate of NTIRE19 RSR""" 13 | 14 | def __init__(self, scale, channel, **kwargs): 15 | args = Config(kwargs) 16 | args.scale = [scale] 17 | args.n_colors = channel 18 | self.rgb_range = args.rgb_range 19 | self.edrn = edrn.EDRN(args) 20 | super(EDRN, self).__init__(channel=channel, scale=scale, **kwargs) 21 | 22 | def fn(self, x): 23 | return self.edrn(x * self.rgb_range) / self.rgb_range 24 | 25 | 26 | class FRN(L1Optimizer): 27 | def __init__(self, scale, channel, **kwargs): 28 | args = Config(kwargs) 29 | args.scale = [scale] 30 | args.n_colors = channel 31 | self.rgb_range = args.rgb_range 32 | self.frn = frn.FRN_UPDOWN(args) 33 | super(FRN, self).__init__(channel=channel, scale=scale, **kwargs) 34 | 35 | def fn(self, x): 36 | return self.frn(x * self.rgb_range) / self.rgb_range 37 | 38 | 39 | class RAN(L1Optimizer): 40 | def __init__(self, scale, channel, **kwargs): 41 | args = Config(kwargs) 42 | args.scale = [scale] 43 | args.n_colors = channel 44 | self.rgb_range = args.rgb_range 45 | self.ran = ran2.RAN(args) 46 | super(RAN, self).__init__(channel=channel, scale=scale, **kwargs) 47 | 48 | def fn(self, x): 49 | return self.ran(x * self.rgb_range) / self.rgb_range 50 | 51 | 52 | class DIDN(L1Optimizer): 53 | def __init__(self, channel, filters, umodule, **kwargs): 54 | self.didn = denoise.EraserTeam.DIDN(channel, filters, umodule) 55 | super(DIDN, self).__init__(channel=channel, **kwargs) 56 | 57 | def fn(self, x): 58 | return self.didn(x) 59 | 60 | 61 | class DHDN(L1Optimizer): 62 | def __init__(self, channel, filters, **kwargs): 63 | self.dhdn = denoise.EraserTeam.DHDN(channel, filters) 64 | super(DHDN, self).__init__(channel=channel, **kwargs) 65 | 66 | def fn(self, x): 67 | return self.dhdn(x) 68 | 69 | 70 | class GRDN(L1Optimizer): 71 | def __init__(self, channel, filters, grdb, rdb, **kwargs): 72 | self.grdn = denoise.DGUTeam.GRDN(channel, filters, grdb, rdb) 73 | super(GRDN, self).__init__(channel=channel, **kwargs) 74 | 75 | def fn(self, x): 76 | return self.grdn(x) 77 | 78 | 79 | class ResUNet(L1Optimizer): 80 | def __init__(self, channel, filters, rb, **kwargs): 81 | self.resunet = denoise.HITVPCTeam.ResUNet(channel, filters, rb) 82 | super(ResUNet, self).__init__(channel=channel, **kwargs) 83 | 84 | def fn(self, x): 85 | return self.resunet(x) 86 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/NTIRE20.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 5 - 28 5 | 6 | import torch 7 | import torch.nn.functional as F 8 | 9 | from .Contrib.ntire20.xiaozhong.ops import define_D, define_F, define_G 10 | from .Model import SuperResolution 11 | from ..Util import Metrics 12 | 13 | 14 | class RealSR(SuperResolution): 15 | """ 16 | RealSR proposed by Xiaozhong Ji 17 | See (NTIRE report, not full paper): https://arxiv.org/pdf/2005.01996.pdf 18 | """ 19 | 20 | def __init__(self, channel=3, scale=4, nf=64, nb=23, **kwargs): 21 | super(RealSR, self).__init__(channel=channel, scale=scale) 22 | self.weights = [ 23 | kwargs.get('pixel_weight', 1), 24 | kwargs.get('feature_weight', 0), 25 | kwargs.get('gan_weight', 0) 26 | ] 27 | self.realsr_g = define_G(in_nc=channel, out_nc=channel, nf=nf, nb=nb) 28 | self.opt_g = torch.optim.Adam(self.trainable_variables('realsr_g'), 1e-4, 29 | betas=(0.5, 0.999)) 30 | if self.weights[1] > 0: 31 | self.feature_net = define_F() # for feature loss 32 | if self.weights[2] > 0: 33 | self.realsr_d = define_D(in_nc=channel, nf=nf, nlayers=3) 34 | self.opt_d = torch.optim.Adam(self.trainable_variables('realsr_d'), 1e-4, 35 | betas=(0.5, 0.999)) 36 | 37 | def train(self, inputs, labels, learning_rate=None): 38 | sr = self.realsr_g(inputs[0]) 39 | pixel_loss = F.l1_loss(sr, labels[0]) * self.weights[0] 40 | loss = pixel_loss 41 | if learning_rate: 42 | for param_group in self.opt.param_groups: 43 | param_group["lr"] = learning_rate 44 | self.opt_g.zero_grad() 45 | loss.backward() 46 | self.opt_g.step() 47 | return {'l1': loss.detach().cpu().numpy()} 48 | 49 | def eval(self, inputs, labels=None, **kwargs): 50 | metrics = {} 51 | sr = self.realsr_g(inputs[0]).cpu().detach() 52 | if labels is not None: 53 | metrics['psnr'] = Metrics.psnr(sr.numpy(), labels[0].cpu().numpy()) 54 | return [sr.numpy()], metrics 55 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Ops/Distortion.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 15 5 | 6 | import random 7 | 8 | import torch 9 | import torch.nn as nn 10 | 11 | from VSR.Util.Math import gaussian_kernel 12 | from ...Util.Utility import gaussian_noise, imfilter, poisson_noise 13 | 14 | 15 | class Distorter(nn.Module): 16 | """Randomly add the noise and blur of an image. 17 | 18 | Args: 19 | gaussian_noise_std (float or tuple of float (min, max)): How much to 20 | additive gaussian white noise. gaussian_noise_std is chosen uniformly 21 | from [0, std] or the given [min, max]. Should be non negative numbers. 22 | poisson_noise_std (float or tuple of float (min, max)): How much to 23 | poisson noise. poisson_noise_std is chosen uniformly from [0, std] or 24 | the given [min, max]. Should be non negative numbers. 25 | gaussian_blur_std (float or tuple of float (min, max)): How much to 26 | blur kernel. gaussian_blur_std is chosen uniformly from [0, std] or 27 | the given [min, max]. Should be non negative numbers. 28 | """ 29 | 30 | def __init__(self, 31 | gaussian_noise_std=0, 32 | poisson_noise_std=0, 33 | gaussian_blur_std=0): 34 | super(Distorter, self).__init__() 35 | self.awgn = self._check_input(gaussian_noise_std, 'awgn', center=0, 36 | bound=(0, 75 / 255), clip_first_on_zero=True) 37 | self.poisson = self._check_input(poisson_noise_std, 'poisson', center=0, 38 | bound=(0, 50 / 255), 39 | clip_first_on_zero=True) 40 | self.blur = self._check_input(gaussian_blur_std, 'blur', center=0) 41 | self.blur_padding = nn.ReflectionPad2d(7) 42 | 43 | def _check_input(self, value, name, center=1, bound=(0, float('inf')), 44 | clip_first_on_zero=True): 45 | if isinstance(value, (tuple, list)) and len(value) == 2: 46 | if not bound[0] <= value[0] <= value[1] <= bound[1]: 47 | raise ValueError("{} values should be between {}".format(name, bound)) 48 | else: 49 | if value < 0: 50 | raise ValueError( 51 | "If {} is a single number, it must be non negative.".format(name)) 52 | value = [center - value, center + value] 53 | if clip_first_on_zero: 54 | value[0] = max(value[0], 0) 55 | # if value is 0 or (1., 1.) for brightness/contrast/saturation 56 | # or (0., 0.) for hue, do nothing 57 | if value[0] == value[1] == center: 58 | value = None 59 | return value 60 | 61 | def forward(self, img): 62 | factors = [] 63 | # noise & blur 64 | blur_factor = 0 65 | if self.blur is not None: 66 | blur_factor = random.uniform(*self.blur) 67 | img = imfilter( 68 | img, 69 | torch.tensor(gaussian_kernel(15, blur_factor), 70 | device=img.device), 71 | self.blur_padding) 72 | awgn_factor = (0, 0, 0) 73 | if self.awgn is not None: 74 | _r = random.uniform(*self.awgn) 75 | _g = random.uniform(*self.awgn) 76 | _b = random.uniform(*self.awgn) 77 | img += gaussian_noise(img, stddev=(_r, _g, _b)) 78 | awgn_factor = (_r, _g, _b) 79 | poisson_factor = (_r, _g, _b) 80 | if self.poisson is not None: 81 | _r = random.uniform(*self.poisson) 82 | _g = random.uniform(*self.poisson) 83 | _b = random.uniform(*self.poisson) 84 | img += poisson_noise(img, stddev=(_r, _g, _b)) 85 | poisson_factor = (_r, _g, _b) 86 | fac = [blur_factor, *awgn_factor, *poisson_factor] 87 | factors.append(torch.tensor(fac)) 88 | img = img.clamp(0, 1) 89 | return img, torch.stack(factors).to(img.device) 90 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Ops/Initializer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 15 5 | 6 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 15 5 | 6 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/Rcan.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019 - 3 - 15 5 | 6 | import logging 7 | 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | from .Model import SuperResolution 13 | from .Optim.SISR import L1Optimizer 14 | from .Ops.Blocks import EasyConv2d, MeanShift, Rcab 15 | from .Ops.Scale import Upsample 16 | from ..Util import Metrics 17 | 18 | _logger = logging.getLogger("VSR.RCAN") 19 | _logger.info("LICENSE: RCAN is implemented by Yulun Zhang. " 20 | "@yulunzhang https://github.com/yulunzhang/RCAN.") 21 | 22 | 23 | class ResidualGroup(nn.Module): 24 | def __init__(self, n_feat, kernel_size, reduction, n_resblocks): 25 | super(ResidualGroup, self).__init__() 26 | modules_body = [Rcab(n_feat, reduction, kernel_size=kernel_size) for _ in 27 | range(n_resblocks)] 28 | modules_body.append(EasyConv2d(n_feat, n_feat, kernel_size)) 29 | self.body = nn.Sequential(*modules_body) 30 | 31 | def forward(self, x): 32 | res = self.body(x) 33 | res += x 34 | return res 35 | 36 | 37 | class Rcan(nn.Module): 38 | def __init__(self, channel, scale, n_resgroups, n_resblocks, n_feats, 39 | reduction, rgb_range): 40 | super(Rcan, self).__init__() 41 | # RGB mean for DIV2K 42 | rgb_mean = (0.4488, 0.4371, 0.4040) 43 | self.sub_mean = MeanShift(rgb_mean, True, rgb_range) 44 | # define head module 45 | modules_head = [EasyConv2d(channel, n_feats, 3)] 46 | # define body module 47 | modules_body = [ 48 | ResidualGroup(n_feats, 3, reduction, n_resblocks) for _ in 49 | range(n_resgroups)] 50 | modules_body.append(EasyConv2d(n_feats, n_feats, 3)) 51 | # define tail module 52 | modules_tail = [Upsample(n_feats, scale), 53 | EasyConv2d(n_feats, channel, 3)] 54 | self.add_mean = MeanShift(rgb_mean, False, rgb_range) 55 | self.head = nn.Sequential(*modules_head) 56 | self.body = nn.Sequential(*modules_body) 57 | self.tail = nn.Sequential(*modules_tail) 58 | 59 | def forward(self, x): 60 | x = self.sub_mean(x) 61 | x = self.head(x) 62 | res = self.body(x) + x 63 | x = self.tail(res) 64 | x = self.add_mean(x) 65 | return x 66 | 67 | 68 | class RCAN(L1Optimizer): 69 | def __init__(self, channel, scale, n_resgroups, n_resblocks, n_feats, 70 | reduction, rgb_range=255, **kwargs): 71 | self.rgb_range = rgb_range 72 | self.rcan = Rcan(channel, scale, n_resgroups, n_resblocks, n_feats, 73 | reduction, rgb_range) 74 | super(RCAN, self).__init__(scale, channel, **kwargs) 75 | 76 | def fn(self, x): 77 | return self.rcan(x * self.rgb_range) / self.rgb_range 78 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c): Wenyi Tang 2017-2019. 2 | # Author: Wenyi Tang 3 | # Email: wenyi.tang@intel.com 4 | # Update Date: 2019/4/3 下午5:10 5 | 6 | import importlib 7 | 8 | __all__ = ['get_model', 'list_supported_models'] 9 | 10 | models = { 11 | 'cubic': ('Bicubic', 'BICUBIC'), 12 | # alias: (file, class) 13 | 'espcn': ('Classic', 'ESPCN'), 14 | 'srcnn': ('Classic', 'SRCNN'), 15 | 'vdsr': ('Classic', 'VDSR'), 16 | 'dncnn': ('Classic', 'DNCNN'), 17 | 'drcn': ('Classic', 'DRCN'), 18 | 'drrn': ('Classic', 'DRRN'), 19 | 'ffdnet': ('Ffdnet', 'FFDNET'), 20 | 'edsr': ('Edsr', 'EDSR'), 21 | 'carn': ('Carn', 'CARN'), 22 | 'dbpn': ('Dbpn', 'DBPN'), 23 | 'rcan': ('Rcan', 'RCAN'), 24 | 'srfeat': ('SRFeat', 'SRFEAT'), 25 | 'esrgan': ('Esrgan', 'ESRGAN'), 26 | 'msrn': ('Msrn', 'MSRN'), 27 | 'crdn': ('Crdn', 'CRDN'), 28 | 'mldn': ('Mldn', 'MLDN'), 29 | 'drn': ('Drn', 'DRN'), 30 | 'sofvsr': ('Sofvsr', 'SOFVSR'), 31 | 'vespcn': ('Vespcn', 'VESPCN'), 32 | 'frvsr': ('Frvsr', 'FRVSR'), 33 | 'qprn': ('Qprn', 'QPRN'), 34 | 'ufvsr': ('Ufvsr', 'UFVSR'), 35 | 'yovsr': ('Yovsr', 'YOVSR'), 36 | 'tecogan': ('TecoGAN', 'TeCoGAN'), 37 | 'spmc': ('Spmc', 'SPMC'), 38 | 'rbpn': ('Rbpn', 'RBPN'), 39 | 'srmd': ('Srmd', 'SRMD'), 40 | # NTIRE 2019 Collections 41 | 'didn': ('NTIRE19', 'DIDN'), 42 | 'dhdn': ('NTIRE19', 'DHDN'), 43 | 'grdn': ('NTIRE19', 'GRDN'), 44 | 'resunet': ('NTIRE19', 'ResUNet'), 45 | 'edrn': ('NTIRE19', 'EDRN'), 46 | 'frn': ('NTIRE19', 'FRN'), 47 | 'ran': ('NTIRE19', 'RAN'), 48 | # NTIRE 2020 49 | 'realsr': ('NTIRE20', 'RealSR'), 50 | 'esr': ('EfficientSR', 'ESR') 51 | } 52 | 53 | 54 | def get_model(name): 55 | module = f'.Backend.Torch.Models.{models[name][0]}' 56 | package = 'VSR' 57 | m = importlib.import_module(module, package) 58 | return m.__dict__[models[name][1]] 59 | 60 | 61 | def list_supported_models(): 62 | return models.keys() 63 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/README.md: -------------------------------------------------------------------------------- 1 | ## VSR PyTorch Implementation 2 | VSR is a collection of recent SR/VSR models, implemented in tensorflow. 3 | 4 | This is an **experimental** extension for VSR, porting to [**pytorch**](https://pytorch.org) framework. 5 | 6 | ## PyTorch v.s. Tensorflow 7 | Why re-invent wheels? This is from my own experience that torch's implementation is always better than tensorflow's (**IMPORTANT: only confirmed on SISR models**). 8 | This extension is used for cross-framework comparison. They share the same data loader, model architecture, benchmark script and even random seeds. 9 | 10 | Besides, most works are implemented in pytorch now, it's quite easy to integrate their pre-trained models into this framework. 11 | 12 | ## Getting Started 13 | Similar to original VSR, there are two main entries "`train.py`" and "`eval.py`" 14 | 15 | #### Train models 16 | ```bash 17 | python train.py [--cuda] [--dataset name] [--epochs num] 18 | ``` 19 | 20 | #### Test models 21 | ```bash 22 | python eval.py [--cuda] [--test path/or/name] [--pth model.pth/path] 23 | ``` 24 | 25 | For more information about dataset name and more advanced options, please refer to documents of original VSR, [here](../Data/README.md) and [here](../Train/README.md). 26 | 27 | #### Sample Code: CARN 28 | ```bash 29 | # training 30 | python train.py carn --cuda --dataset div2k --epochs 1000 31 | # testing 32 | python eval.py carn --cuda --test set5 set14 bsd100 urban100 33 | # predicting own data 34 | python eval.py carn --cuda --test /tmp/myphotos/*.png 35 | ``` 36 | 37 | 38 | ## NTIRE 19 Reproduce 39 | Refer to [here](../Docs/README_NTIRE19.md). -------------------------------------------------------------------------------- /VSR/Backend/Torch/Util/Metrics.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | import numpy as np 7 | import torch 8 | 9 | 10 | def psnr(x, y, max_val=1.0): 11 | if isinstance(x, torch.Tensor): 12 | x = x.detach().cpu().numpy() 13 | if isinstance(y, torch.Tensor): 14 | y = y.detach().cpu().numpy() 15 | mse = np.square(x - y).mean() 16 | return 10 * np.log10(max_val ** 2 / mse) 17 | 18 | 19 | def total_variance(x, reduction='mean'): 20 | hor = x[..., :-1, :] - x[..., 1:, :] 21 | ver = x[..., :-1] - x[..., 1:] 22 | if isinstance(x, torch.Tensor): 23 | tot_var = torch.sum(torch.abs(hor)) + torch.sum(torch.abs(ver)) 24 | elif isinstance(x, np.ndarray): 25 | tot_var = np.abs(hor).sum() + np.abs(ver).sum() 26 | if reduction == 'mean': 27 | reduce = x.shape[-1] * x.shape[-2] 28 | else: 29 | reduce = 1 30 | return tot_var / reduce 31 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/Util/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | -------------------------------------------------------------------------------- /VSR/Backend/Torch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LoSealL/VideoSuperResolution/4c86e49d81c7a9bea1fe0780d651afc126768df3/VSR/Backend/Torch/__init__.py -------------------------------------------------------------------------------- /VSR/Backend/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | import os 7 | import logging 8 | from importlib import import_module 9 | from pathlib import Path 10 | 11 | import yaml 12 | 13 | try: 14 | from yaml import FullLoader as _Loader 15 | except ImportError: 16 | # For older versions 17 | from yaml import Loader as _Loader 18 | 19 | LOG = logging.getLogger('VSR') 20 | HOME = os.environ.get('VSR_HOME') 21 | if not HOME: 22 | HOME = Path('~').expanduser() / '.vsr' 23 | CONFIG = { 24 | 'backend': os.environ.get('VSR_BACKEND', 'pytorch'), 25 | 'verbose': os.environ.get('VSR_VERBOSE', 'info'), 26 | } 27 | if Path(HOME / 'config.yml').exists(): 28 | with open(HOME / 'config.yml', encoding='utf8') as fd: 29 | CONFIG = yaml.load(fd.read(), Loader=_Loader) 30 | 31 | LOG.setLevel(CONFIG['verbose'].upper()) 32 | hdl = logging.StreamHandler() 33 | hdl.setFormatter(logging.Formatter("%(asctime)s %(levelname)s: %(message)s")) 34 | LOG.addHandler(hdl) 35 | 36 | BACKEND = CONFIG['backend'].lower() 37 | if BACKEND == 'auto': 38 | BACKEND = 'tensorflow' 39 | if BACKEND not in ('tensorflow', 'keras', 'pytorch'): 40 | BACKEND = 'pytorch' 41 | 42 | if BACKEND in ('tensorflow', 'keras'): 43 | try: 44 | tf = import_module('tensorflow') 45 | CONFIG['data_format'] = 'channels_last' 46 | tf_ver_major, tf_ver_minor, _ = [int(s) for s in tf.__version__.split('.')] 47 | if BACKEND == 'keras' and tf_ver_major < 2: 48 | LOG.warning(f"[!] Current tensorflow version is {tf.__version__}") 49 | LOG.info("[*] Fallback to use legacy tensorflow v1.x") 50 | BACKEND = 'tensorflow' 51 | if tf_ver_major == 1 and tf_ver_minor < 15: 52 | LOG.warning("[!!] VSR does not support TF < 1.15.0 any longer.") 53 | LOG.warning("[!] Considering use an old version of VSR, " 54 | "or update your tensorflow version.") 55 | raise ImportError 56 | except ImportError: 57 | LOG.warning("[!] Tensorflow package not found in your system.") 58 | LOG.info("[*] Fallback to use PyTorch...") 59 | BACKEND = 'pytorch' 60 | 61 | if BACKEND == 'pytorch': 62 | try: 63 | torch = import_module('torch') 64 | CONFIG['data_format'] = 'channels_first' 65 | _ver = torch.__version__.split('.') 66 | if _ver[0] != '1' or _ver[1] <= '1': 67 | LOG.warning( 68 | f"[!] PyTorch version too low: {torch.__version__}, recommended 1.2.0") 69 | except ImportError: 70 | LOG.fatal("[!] PyTorch package not found in your system.") 71 | raise ImportError("Not an available backend found! Check your environment.") 72 | 73 | DATA_FORMAT = CONFIG['data_format'].lower() 74 | -------------------------------------------------------------------------------- /VSR/DataLoader/Crop.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | import numpy as np 7 | from ..Backend import DATA_FORMAT 8 | 9 | 10 | class Cropper(object): 11 | def __init__(self, scale): 12 | self.scale = scale 13 | 14 | def __call__(self, img_pair: tuple, shape: list): 15 | assert len(img_pair) >= 2, \ 16 | f"Pair must contain more than 2 elements, which is {img_pair}" 17 | for img in img_pair: 18 | assert img.ndim == len(shape), \ 19 | f"Shape mis-match: {img.ndim} != {len(shape)}" 20 | 21 | return self.call(img_pair, shape) 22 | 23 | def call(self, img: tuple, shape: (list, tuple)) -> tuple: 24 | raise NotImplementedError 25 | 26 | 27 | class RandomCrop(Cropper): 28 | def call(self, img: tuple, shape: (list, tuple)) -> tuple: 29 | hr, lr = img 30 | if lr.shape[-2] < shape[-2]: 31 | raise ValueError( 32 | f"Batch shape is larger than data: {lr.shape} vs {shape}") 33 | ind = [np.random.randint(nd + 1) for nd in lr.shape - np.array(shape)] 34 | slc1 = [slice(n, n + s) for n, s in zip(ind, shape)] 35 | slc2 = slc1.copy() 36 | if DATA_FORMAT == 'channels_last': 37 | slc2[-2] = slice(ind[-2] * self.scale, 38 | (ind[-2] + shape[-2]) * self.scale) 39 | slc2[-3] = slice(ind[-3] * self.scale, 40 | (ind[-3] + shape[-3]) * self.scale) 41 | else: 42 | slc2[-1] = slice(ind[-1] * self.scale, 43 | (ind[-1] + shape[-1]) * self.scale) 44 | slc2[-2] = slice(ind[-2] * self.scale, 45 | (ind[-2] + shape[-2]) * self.scale) 46 | return hr[tuple(slc2)], lr[tuple(slc1)] 47 | 48 | 49 | class CenterCrop(Cropper): 50 | def call(self, img: tuple, shape: (list, tuple)) -> tuple: 51 | hr, lr = img 52 | ind = [nd // 2 for nd in hr.shape - np.array(shape)] 53 | slc1 = [slice(n, n + s) for n, s in zip(ind, shape)] 54 | slc2 = slc1.copy() 55 | if DATA_FORMAT == 'channels_last': 56 | slc2[-2] = slice(ind[-2] * self.scale, 57 | (ind[-2] + shape[-2]) * self.scale) 58 | slc2[-3] = slice(ind[-3] * self.scale, 59 | (ind[-3] + shape[-3]) * self.scale) 60 | else: 61 | slc2[-1] = slice(ind[-1] * self.scale, 62 | (ind[-1] + shape[-1]) * self.scale) 63 | slc2[-2] = slice(ind[-2] * self.scale, 64 | (ind[-2] + shape[-2]) * self.scale) 65 | return hr[tuple(slc2)], lr[tuple(slc1)] 66 | -------------------------------------------------------------------------------- /VSR/DataLoader/FloDecoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | import logging 6 | from importlib import import_module 7 | 8 | import numpy as np 9 | 10 | 11 | def open_flo(fn): 12 | """ Read .flo file in Middlebury format 13 | # Code adapted from: 14 | # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy 15 | 16 | # WARNING: this will work on little-endian architectures (eg Intel x86) only! 17 | # print 'fn = %s'%(fn) 18 | """ 19 | with open(fn, 'rb') as f: 20 | magic = np.fromfile(f, np.float32, count=1) 21 | if 202021.25 != magic: 22 | logging.error('Magic number incorrect. Invalid .flo file') 23 | return None 24 | else: 25 | w = np.fromfile(f, np.int32, count=1)[0] 26 | h = np.fromfile(f, np.int32, count=1)[0] 27 | # print 'Reading %d x %d flo file\n' % (w, h) 28 | data = np.fromfile(f, np.float32, count=2 * w * h) 29 | # Reshape data into 3D array (columns, rows, bands) 30 | # The reshape here is for visualization, the original code is (w,h,2) 31 | return np.resize(data, (int(h), int(w), 2)) 32 | 33 | 34 | def write_flo(filename, uv, v=None): 35 | """ Write optical flow to file. 36 | 37 | Original code by Deqing Sun, adapted from Daniel Scharstein. 38 | """ 39 | n_bands = 2 40 | _TAG_CHAR = np.array([202021.25], np.float32) 41 | 42 | if v is None: 43 | u = uv[..., 0] 44 | v = uv[..., 1] 45 | else: 46 | u = uv 47 | height, width = u.shape 48 | with open(filename, 'wb') as f: 49 | # write the header 50 | f.write(_TAG_CHAR.tobytes()) 51 | np.array(width).astype(np.int32).tofile(f) 52 | np.array(height).astype(np.int32).tofile(f) 53 | # arrange into matrix form 54 | tmp = np.zeros((height, width * n_bands)) 55 | tmp[:, np.arange(width) * 2] = u 56 | tmp[:, np.arange(width) * 2 + 1] = v 57 | tmp.astype(np.float32).tofile(f) 58 | 59 | 60 | class KITTI: 61 | @staticmethod 62 | def open_png16(fn): 63 | """Read 16bit png file""" 64 | 65 | png = import_module('png') 66 | reader = png.Reader(fn) 67 | data = reader.asDirect() 68 | pixels = [] 69 | for row in data[2]: 70 | row = np.reshape(np.asarray(row), [-1, 3]) 71 | pixels += [row] 72 | return np.stack(pixels, 0) 73 | 74 | @staticmethod 75 | def open_flow(fn): 76 | flow = KITTI.open_png16(fn) 77 | valid = flow[..., -1] 78 | u = flow[..., 0].astype('float32') 79 | v = flow[..., 1].astype('float32') 80 | u = (u - 2 ** 15) / 64 * valid 81 | v = (v - 2 ** 15) / 64 * valid 82 | return np.stack([u, v], -1) 83 | -------------------------------------------------------------------------------- /VSR/DataLoader/NVDecoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | # Image customized decoder for NV12([Y][UV/4]), NV21([Y][VU/4]) 7 | # NOTE: [Y] means Y channel is a planar channel, [UV] means UV 8 | # channels together is planar, but U and V are packed. [UV/4] 9 | # means U and V are sub-sampled by a factor of [2, 2] 10 | 11 | import numpy as np 12 | from PIL import ImageFile 13 | 14 | 15 | class NV12Decoder(ImageFile.PyDecoder): 16 | """PIL.Image.DECODERS for NV12 format raw bytes 17 | 18 | Registered in `Image.DECODERS`, don't use this class directly! 19 | """ 20 | 21 | def __init__(self, mode, *args): 22 | super(NV12Decoder, self).__init__(mode, *args) 23 | 24 | def decode(self, buffer): 25 | if self.mode == 'L': 26 | # discard UV channel 27 | self.set_as_raw(buffer, 'L') 28 | else: 29 | w, h = self.im.size 30 | y = np.frombuffer(buffer, 'uint8', count=w * h) 31 | uv = np.frombuffer(buffer, 'uint8', count=w * h // 2, offset=w * h) 32 | y = np.reshape(y, [h, w, 1]) 33 | uv = np.reshape(uv, [h // 2, w // 2, 2]) 34 | uv = uv[np.arange(h) // 2][:, np.arange(w) // 2] 35 | yuv = np.concatenate([y, uv], axis=-1) 36 | self.set_as_raw(yuv.flatten().tobytes()) 37 | return -1, 0 38 | 39 | 40 | class NV21Decoder(ImageFile.PyDecoder): 41 | """PIL.Image.DECODERS for NV21 format raw bytes 42 | 43 | Registered in `Image.DECODERS`, don't use this class directly! 44 | """ 45 | 46 | def __init__(self, mode, *args): 47 | super(NV21Decoder, self).__init__(mode, *args) 48 | 49 | def decode(self, buffer): 50 | if self.mode == 'L': 51 | # discard UV channel 52 | self.set_as_raw(buffer, 'L') 53 | else: 54 | w, h = self.im.size 55 | y = np.frombuffer(buffer, 'uint8', count=w * h) 56 | vu = np.frombuffer(buffer, 'uint8', count=w * h // 2, offset=w * h) 57 | y = np.reshape(y, [h, w, 1]) 58 | vu = np.reshape(vu, [h // 2, w // 2, 2]) 59 | vu = vu[np.arange(h) // 2][:, np.arange(w) // 2] 60 | uv = vu[:, :, ::-1] 61 | yuv = np.concatenate([y, uv], axis=-1) 62 | self.set_as_raw(yuv.flatten().tobytes()) 63 | return -1, 0 64 | -------------------------------------------------------------------------------- /VSR/DataLoader/Transform.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | from PIL import Image, ImageFilter, ImageEnhance 7 | import numpy as np 8 | 9 | 10 | class Transformer(object): 11 | """Image transformer. 12 | 13 | Args: 14 | value: the parameter for each transform function. 15 | random: if specify 'uniform', generate value sampled from 0 to `+value`; 16 | if specify 'normal', generate value N~(mean=0, std=value) 17 | """ 18 | 19 | def __init__(self, value=1, random=None): 20 | self._v = value 21 | self._r = random 22 | 23 | @property 24 | def value(self): 25 | if self._r == 'uniform': 26 | return np.random.uniform(0, self._v) 27 | elif self._r == 'normal': 28 | return np.random.normal(0, self._v) 29 | else: 30 | return self._v 31 | 32 | 33 | class _Transformer1(Transformer): 34 | def __call__(self, img: Image.Image): 35 | assert isinstance(img, Image.Image) 36 | return self.call(img) 37 | 38 | def call(self, img): 39 | raise NotImplementedError 40 | 41 | 42 | class Tidy(_Transformer1): 43 | def call(self, img: Image.Image): 44 | scale = self.value 45 | shape = np.array((img.width, img.height)) 46 | shape -= shape % scale 47 | return img.crop([0, 0, *shape.tolist()]) 48 | 49 | 50 | class Bicubic(_Transformer1): 51 | def call(self, img: Image.Image): 52 | scale = self.value 53 | shape = np.array((img.width, img.height)) 54 | if scale < 1: 55 | rscale = int(1 / scale) 56 | if np.any(shape % rscale): 57 | raise ValueError(f"Image size is not divisible by {rscale}.") 58 | return img.resize(shape // rscale, resample=Image.BICUBIC) 59 | else: 60 | return img.resize((shape * scale).astype('int32'), resample=Image.BICUBIC) 61 | 62 | 63 | class Brightness(_Transformer1): 64 | def call(self, img: Image.Image): 65 | brightness = max(0, self.value) 66 | return ImageEnhance.Brightness(img).enhance(brightness) 67 | 68 | 69 | class Contrast(_Transformer1): 70 | def call(self, img: Image.Image): 71 | contrast = self.value 72 | return ImageEnhance.Contrast(img).enhance(contrast) 73 | 74 | 75 | class Sharpness(_Transformer1): 76 | def call(self, img): 77 | sharp = min(max(0, self.value), 2) 78 | return ImageEnhance.Sharpness(img).enhance(sharp) 79 | 80 | 81 | class GaussianBlur(_Transformer1): 82 | def call(self, img): 83 | radius = self.value 84 | return ImageFilter.GaussianBlur(radius).filter(img) 85 | 86 | 87 | class _Transformer2(Transformer): 88 | def __call__(self, img: np.ndarray): 89 | assert isinstance(img, np.ndarray) 90 | return self.call(img) 91 | 92 | def call(self, img): 93 | raise NotImplementedError 94 | 95 | 96 | class GaussianWhiteNoise(_Transformer2): 97 | def call(self, img): 98 | shape = img.shape 99 | noise = np.random.normal(0, self.value, shape) 100 | noise += img.astype('float32') 101 | return np.clip(np.round(noise), 0, 255).astype('uint8') 102 | 103 | 104 | class FixedVideoLengthBatch(_Transformer2): 105 | def call(self, img): 106 | assert img.ndim == 5, f"img is not 5D, which is {img.ndim}" 107 | depth = int(self.value) 108 | shape = img.shape 109 | if shape[1] <= depth: 110 | return img 111 | ret = [] 112 | for i in range(shape[1] - depth + 1): 113 | ret.append(img[:, i * depth: (i + 1) * depth]) 114 | return np.stack(ret, 1).reshape([-1, depth, *shape[-3:]]) 115 | -------------------------------------------------------------------------------- /VSR/DataLoader/YVDecoder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | # Image customized decoder for YV12([Y][U/4][V/4]), YV21([Y][V/4][U/4]) 7 | # NOTE: [Y][U][V] means Y/U/V channel is a planar channel, [U/4] means 8 | # U channel is sub-sampled by a factor of [2, 2] 9 | 10 | import numpy as np 11 | from PIL import ImageFile 12 | 13 | 14 | class YV12Decoder(ImageFile.PyDecoder): 15 | """PIL.Image.DECODERS for YV12 format raw bytes 16 | 17 | Registered in `Image.DECODERS`, don't use this class directly! 18 | """ 19 | 20 | def __init__(self, mode, *args): 21 | super(YV12Decoder, self).__init__(mode, *args) 22 | 23 | def decode(self, buffer): 24 | if self.mode == 'L': 25 | # discard UV channel 26 | self.set_as_raw(buffer, 'L') 27 | else: 28 | w, h = self.im.size 29 | y = np.frombuffer(buffer, 'uint8', count=w * h) 30 | u = np.frombuffer(buffer, 'uint8', count=w * h // 4, offset=w * h) 31 | v = np.frombuffer( 32 | buffer, 'uint8', count=w * h // 4, offset=w * h + w * h // 4) 33 | y = np.reshape(y, [h, w]) 34 | u = np.reshape(u, [h // 2, w // 2]) 35 | v = np.reshape(v, [h // 2, w // 2]) 36 | u = u[np.arange(h) // 2][:, np.arange(w) // 2] 37 | v = v[np.arange(h) // 2][:, np.arange(w) // 2] 38 | yuv = np.stack([y, u, v], axis=-1) 39 | self.set_as_raw(yuv.flatten().tobytes()) 40 | return -1, 0 41 | 42 | 43 | class YV21Decoder(ImageFile.PyDecoder): 44 | """PIL.Image.DECODERS for YV21 format raw bytes 45 | 46 | Registered in `Image.DECODERS`, don't use this class directly! 47 | """ 48 | 49 | def __init__(self, mode, *args): 50 | super(YV21Decoder, self).__init__(mode, *args) 51 | 52 | def decode(self, buffer): 53 | if self.mode == 'L': 54 | # discard UV channel 55 | self.set_as_raw(buffer, 'L') 56 | else: 57 | w, h = self.im.size 58 | y = np.frombuffer(buffer, 'uint8', count=w * h) 59 | v = np.frombuffer(buffer, 'uint8', count=w * h // 4, offset=w * h) 60 | u = np.frombuffer( 61 | buffer, 'uint8', count=w * h // 4, offset=w * h + w * h // 4) 62 | y = np.reshape(y, [h, w]) 63 | u = np.reshape(u, [h // 2, w // 2]) 64 | v = np.reshape(v, [h // 2, w // 2]) 65 | u = u[np.arange(h) // 2][:, np.arange(w) // 2] 66 | v = v[np.arange(h) // 2][:, np.arange(w) // 2] 67 | yuv = np.stack([y, u, v], axis=-1) 68 | self.set_as_raw(yuv.flatten().tobytes()) 69 | return -1, 0 70 | -------------------------------------------------------------------------------- /VSR/DataLoader/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | from .Crop import CenterCrop, RandomCrop 7 | from .Dataset import Container, Dataset, load_datasets 8 | from .Loader import Loader 9 | from .Transform import ( 10 | Bicubic, Brightness, Contrast, FixedVideoLengthBatch, GaussianBlur, 11 | GaussianWhiteNoise, Sharpness 12 | ) 13 | 14 | __all__ = [ 15 | 'load_datasets', 16 | 'Dataset', 17 | 'Loader', 18 | 'CenterCrop', 19 | 'RandomCrop', 20 | 'Bicubic', 21 | 'Brightness', 22 | 'Contrast', 23 | 'FixedVideoLengthBatch', 24 | 'GaussianWhiteNoise', 25 | 'GaussianBlur', 26 | 'Sharpness' 27 | ] 28 | -------------------------------------------------------------------------------- /VSR/Model/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | from importlib import import_module 7 | 8 | from ..Backend import BACKEND 9 | 10 | __all__ = [ 11 | 'get_model', 12 | 'list_supported_models' 13 | ] 14 | 15 | 16 | def get_model(name: str): 17 | name = name.lower() 18 | try: 19 | if BACKEND == 'pytorch': 20 | return import_module('.Models', 'VSR.Backend.Torch').get_model(name) 21 | elif BACKEND == 'tensorflow': 22 | return import_module('.Models', 'VSR.Backend.TF').get_model(name) 23 | elif BACKEND == 'keras': 24 | return import_module('.Models', 'VSR.Backend.Keras').get_model(name) 25 | except (KeyError, ImportError): 26 | raise ImportError(f"Using {BACKEND}, can't find model {name}.") 27 | 28 | 29 | def list_supported_models(): 30 | if BACKEND == 'pytorch': 31 | return import_module('.Models', 'VSR.Backend.Torch').list_supported_models() 32 | elif BACKEND == 'tensorflow': 33 | return import_module('.Models', 'VSR.Backend.TF').list_supported_models() 34 | elif BACKEND == 'keras': 35 | return import_module('.Models', 'VSR.Backend.Keras').list_supported_models() 36 | -------------------------------------------------------------------------------- /VSR/Util/Config.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | import easydict 7 | import yaml 8 | 9 | try: 10 | from yaml import FullLoader as _Loader 11 | except ImportError: 12 | from yaml import Loader as _Loader 13 | 14 | 15 | class Config(easydict.EasyDict): 16 | def __init__(self, obj=None, **kwargs): 17 | super(Config, self).__init__(**kwargs) 18 | if obj is not None: 19 | assert isinstance(obj, (dict, str)) 20 | if isinstance(obj, str): 21 | with open(obj, 'r') as fd: 22 | obj = yaml.load(fd, Loader=_Loader) 23 | self.update(**obj) 24 | 25 | def __getattr__(self, item): 26 | return self.get(item) 27 | -------------------------------------------------------------------------------- /VSR/Util/Ensemble.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 6 - 17 5 | 6 | import numpy as np 7 | 8 | 9 | class Ensembler: 10 | @staticmethod 11 | def expand(feature: np.ndarray): 12 | r0 = feature.copy() 13 | r1 = np.rot90(feature, 1, axes=[-3, -2]) 14 | r2 = np.rot90(feature, 2, axes=[-3, -2]) 15 | r3 = np.rot90(feature, 3, axes=[-3, -2]) 16 | r4 = np.flip(feature, axis=-2) 17 | r5 = np.rot90(r4, 1, axes=[-3, -2]) 18 | r6 = np.rot90(r4, 2, axes=[-3, -2]) 19 | r7 = np.rot90(r4, 3, axes=[-3, -2]) 20 | return r0, r1, r2, r3, r4, r5, r6, r7 21 | 22 | @staticmethod 23 | def merge(outputs: [np.ndarray]): 24 | results = [] 25 | for i in outputs: 26 | outputs_ensemble = [ 27 | i[0], 28 | np.rot90(i[1], 3, axes=[-3, -2]), 29 | np.rot90(i[2], 2, axes=[-3, -2]), 30 | np.rot90(i[3], 1, axes=[-3, -2]), 31 | np.flip(i[4], axis=-2), 32 | np.flip(np.rot90(i[5], 3, axes=[-3, -2]), axis=-2), 33 | np.flip(np.rot90(i[6], 2, axes=[-3, -2]), axis=-2), 34 | np.flip(np.rot90(i[7], 1, axes=[-3, -2]), axis=-2), 35 | ] 36 | results.append( 37 | np.concatenate(outputs_ensemble).mean(axis=0, keepdims=True)) 38 | return results 39 | -------------------------------------------------------------------------------- /VSR/Util/GoogleDriveDownloader.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright: Wenyi Tang 2017-2018 3 | Author: Wenyi Tang 4 | Email: wenyi.tang@intel.com 5 | Created Date: Dec 21st 2018 6 | 7 | Download binary files shared on google drive 8 | """ 9 | 10 | import io 11 | import sys 12 | from pathlib import Path 13 | 14 | try: 15 | from googleapiclient.discovery import build 16 | from googleapiclient.http import MediaIoBaseDownload 17 | from httplib2 import Http 18 | from oauth2client import file, client, tools 19 | except ImportError as ex: 20 | raise ImportError( 21 | "To download shared google drive file via python," 22 | "google-api-python-client, oauth2client is required." 23 | "Please use pip install google-api-python-client oauth2client.") 24 | 25 | SCOPES = 'https://www.googleapis.com/auth/drive.readonly' 26 | CREDENTIALS = './Data/credentials.json' 27 | 28 | 29 | def require_authorize(store, credentials, scope): 30 | _argv = sys.argv.copy() 31 | sys.argv = _argv[:1] 32 | if '--noauth_local_webserver' in _argv: 33 | sys.argv.append('--noauth_local_webserver') 34 | flow = client.flow_from_clientsecrets(credentials, scope) 35 | creds = tools.run_flow(flow, store) 36 | sys.argv = _argv 37 | return creds 38 | 39 | 40 | def drive_download(name, fileid, path): 41 | store_path = Path(path) / name 42 | if store_path.exists(): 43 | print("{} exists, skip download.".format(name)) 44 | return store_path 45 | # The file token.json stores the user's access and refresh tokens, and is 46 | # created automatically when the authorization flow completes for the first 47 | # time. 48 | store = file.Storage('/tmp/token.json') 49 | creds = store.get() 50 | if not creds or creds.invalid: 51 | creds = require_authorize(store, CREDENTIALS, SCOPES) 52 | service = build('drive', 'v3', http=creds.authorize(Http())) 53 | 54 | request = service.files().get_media(fileId=fileid) 55 | 56 | fh = io.FileIO(store_path.resolve(), 'wb') 57 | downloader = MediaIoBaseDownload(fh, request) 58 | done = False 59 | while not done: 60 | status, done = downloader.next_chunk() 61 | print("\rDownload {}%.".format(int(status.progress() * 100))) 62 | print('\n', flush=True) 63 | if done: 64 | return store_path 65 | -------------------------------------------------------------------------------- /VSR/Util/Hook.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | from functools import partial 7 | from pathlib import Path 8 | 9 | from .ImageProcess import array_to_img 10 | 11 | 12 | def _str_to_slice(index: str): 13 | index = index.split(':') 14 | if len(index) == 1: 15 | ind = int(index[0]) 16 | if ind < 0: 17 | sl = slice(ind, None, None) 18 | else: 19 | sl = slice(ind, ind + 1) 20 | else: 21 | def _maybe_int(x): 22 | try: 23 | return int(x) 24 | except ValueError: 25 | return None 26 | 27 | sl = slice(*(_maybe_int(i) for i in index)) 28 | return sl 29 | 30 | 31 | def _save_model_predicted_images(output, names, save_dir, index, auto_rename): 32 | assert len(names) == 1, f"Name list exceeds 1, which is {names}" 33 | name = names[0] 34 | for img in output[_str_to_slice(index)]: 35 | shape = img.shape 36 | path = Path(save_dir) 37 | if shape[0] > 1 or auto_rename: 38 | path /= name 39 | path.mkdir(exist_ok=True, parents=True) 40 | for i, n in enumerate(img): 41 | rep = 0 42 | if auto_rename: 43 | while (path / f"{name}_id{i:04d}_{rep:04d}.png").exists(): 44 | rep += 1 45 | path /= f"{name}_id{i:04d}_{rep:04d}.png" 46 | array_to_img(n).convert('RGB').save(str(path)) 47 | return output 48 | 49 | 50 | def save_inference_images(save_dir, multi_output_index=-1, auto_rename=None): 51 | return partial(_save_model_predicted_images, save_dir=save_dir, 52 | index=multi_output_index, auto_rename=auto_rename) 53 | -------------------------------------------------------------------------------- /VSR/Util/LearningRateScheduler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | from functools import partial 6 | 7 | 8 | def _exponential_decay(start_lr, steps, decay_step, decay_rate, **kwargs): 9 | return start_lr * decay_rate ** (steps / decay_step) 10 | 11 | 12 | def _poly_decay(start_lr, end_lr, steps, decay_step, power, **kwargs): 13 | return (start_lr - end_lr) * (1 - steps / decay_step) ** power + end_lr 14 | 15 | 16 | def _stair_decay(start_lr, steps, decay_step, decay_rate, **kwargs): 17 | return start_lr * decay_rate ** (steps // decay_step) 18 | 19 | 20 | def _multistep_decay(start_lr, steps, decay_step, decay_rate, **kwargs): 21 | if not decay_step: 22 | return start_lr 23 | for n, s in enumerate(decay_step): 24 | if steps <= s: 25 | return start_lr * (decay_rate ** n) 26 | if steps > decay_step[-1]: 27 | return start_lr * (decay_rate ** len(decay_step)) 28 | 29 | 30 | def lr_decay(method, lr, **kwargs): 31 | if method == 'exp': 32 | return partial(_exponential_decay, start_lr=lr, **kwargs) 33 | elif method == 'poly': 34 | return partial(_poly_decay, start_lr=lr, **kwargs) 35 | elif method == 'stair': 36 | return partial(_stair_decay, start_lr=lr, **kwargs) 37 | elif method == 'multistep': 38 | return partial(_multistep_decay, start_lr=lr, **kwargs) 39 | else: 40 | print('invalid decay method!') 41 | return None 42 | -------------------------------------------------------------------------------- /VSR/Util/VisualizeOpticalFlow.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | import numpy as np 7 | from .ImageProcess import array_to_img 8 | from ..Backend import DATA_FORMAT 9 | 10 | 11 | def _color_wheel(): 12 | red_yellow, yellow_green, green_cyan = 15, 6, 4 13 | cyan_blue, blue_magenta, magenta_red = 11, 13, 6 14 | colors = [red_yellow, yellow_green, green_cyan, 15 | cyan_blue, blue_magenta, magenta_red] 16 | color = np.zeros([np.sum(colors), 3]) 17 | for i in range(red_yellow): 18 | color[i] = [255, 255 * i / red_yellow, 0] 19 | for i in range(yellow_green): 20 | color[i + np.sum(colors[:1])] = [255 - 255 * i / yellow_green, 255, 0] 21 | for i in range(green_cyan): 22 | color[i + np.sum(colors[:2])] = [0, 255, 255 * i / green_cyan] 23 | for i in range(cyan_blue): 24 | color[i + np.sum(colors[:3])] = [0, 255 - 255 * i / cyan_blue, 255] 25 | for i in range(blue_magenta): 26 | color[i + np.sum(colors[:4])] = [255 * i / blue_magenta, 0, 255] 27 | for i in range(magenta_red): 28 | color[i + np.sum(colors[:5])] = [255, 0, 255 - 255 * i / magenta_red] 29 | return color / 255 30 | 31 | 32 | def _viz_flow(u, v, logscale=True, scaledown=6): 33 | """ 34 | Copied from @jswulff: 35 | https://github.com/jswulff/pcaflow/blob/master/pcaflow/utils/viz_flow.py 36 | 37 | top_left is zero, u is horizon, v is vertical 38 | red is 3 o'clock, yellow is 6, light blue is 9, blue/purple is 12 39 | """ 40 | color_wheel = _color_wheel() 41 | n_cols = color_wheel.shape[0] 42 | 43 | radius = np.sqrt(u ** 2 + v ** 2) 44 | if logscale: 45 | radius = np.log(radius + 1) 46 | radius = radius / scaledown 47 | rot = np.arctan2(-v, -u) / np.pi 48 | 49 | fk = (rot + 1) / 2 * (n_cols - 1) # -1~1 mapped to 0~n_cols 50 | k0 = fk.astype(np.uint8) # 0, 1, 2, ..., n_cols 51 | 52 | k1 = k0 + 1 53 | k1[k1 == n_cols] = 0 54 | 55 | f = fk - k0 56 | 57 | n_colors = color_wheel.shape[1] 58 | img = np.zeros(u.shape + (n_colors,)) 59 | for i in range(n_colors): 60 | tmp = color_wheel[:, i] 61 | col0 = tmp[k0] 62 | col1 = tmp[k1] 63 | col = (1 - f) * col0 + f * col1 64 | 65 | idx = radius <= 1 66 | # increase saturation with radius 67 | col[idx] = 1 - radius[idx] * (1 - col[idx]) 68 | # out of range 69 | col[~idx] *= 0.75 70 | img[:, :, i] = np.floor(255 * col).astype(np.uint8) 71 | 72 | return img.astype(np.uint8) 73 | 74 | 75 | def visualize_flow(flow, v=None): 76 | if DATA_FORMAT == 'channels_last': 77 | u = flow[..., 0] if v is None else flow 78 | v = flow[..., 1] if v is None else v 79 | else: 80 | u = flow[0, :, :] if v is None else flow 81 | v = flow[1, :, :] if v is None else v 82 | viz = _viz_flow(u, v) 83 | return array_to_img(viz, 'RGB') 84 | -------------------------------------------------------------------------------- /VSR/Util/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | 6 | from .Config import Config 7 | from .Hook import save_inference_images 8 | from .ImageProcess import ( 9 | array_to_img, img_to_array, imread, imresize, rgb_to_yuv 10 | ) 11 | from .LearningRateScheduler import lr_decay 12 | from .Utility import (str_to_bytes, suppress_opt_by_args, to_list, compat_param) 13 | 14 | 15 | __all__ = [ 16 | 'Config', 17 | 'lr_decay', 18 | 'str_to_bytes', 19 | 'suppress_opt_by_args', 20 | 'to_list', 21 | 'array_to_img', 22 | 'imresize', 23 | 'imread', 24 | 'img_to_array', 25 | 'rgb_to_yuv', 26 | 'save_inference_images', 27 | 'compat_param', 28 | ] 29 | -------------------------------------------------------------------------------- /VSR/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 7 5 | # Collect and reproduce Image/Video Super-Resolution Algorithms. 6 | 7 | from __future__ import absolute_import 8 | from __future__ import print_function 9 | 10 | import VSR.Backend 11 | import VSR.DataLoader 12 | 13 | __all__ = [ 14 | 'Backend', 15 | 'DataLoader', 16 | 'Model', 17 | 'Util', 18 | ] 19 | -------------------------------------------------------------------------------- /azure-pipelines.yml: -------------------------------------------------------------------------------- 1 | # Python package 2 | # Create and test a Python package on multiple Python versions. 3 | # Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more: 4 | # https://docs.microsoft.com/azure/devops/pipelines/languages/python 5 | 6 | trigger: 7 | - master 8 | 9 | jobs: 10 | 11 | - job: 'Test' 12 | pool: 13 | vmImage: 'Ubuntu-16.04' 14 | 15 | steps: 16 | - task: UsePythonVersion@0 17 | inputs: 18 | versionSpec: '3.6' 19 | architecture: 'x64' 20 | 21 | - script: | 22 | python -m pip install --upgrade pip 23 | pip install -q tensorflow==1.15.0 google-api-python-client oauth2client 24 | pip install -q torch==1.4.0+cpu torchvision==0.5.0+cpu -f https://download.pytorch.org/whl/torch_stable.html 25 | pip install -U -q -e . 26 | displayName: 'Install dependencies' 27 | 28 | - script: | 29 | pytest -v Tests/training_test.py --disable-warnings 30 | displayName: 'test with pytest' 31 | 32 | - task: PublishTestResults@2 33 | inputs: 34 | testResultsFiles: '**/test-results.xml' 35 | testRunTitle: 'Python $(python.version)' 36 | condition: succeededOrFailed() 37 | enabled: false 38 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2017-2020 Wenyi Tang. 2 | # Author: Wenyi Tang 3 | # Email: wenyitang@outlook.com 4 | # Update: 2020 - 2 - 16 5 | 6 | from setuptools import find_packages 7 | from setuptools import setup 8 | 9 | # Get version from CHANGELOG 10 | try: 11 | with open('CHANGELOG.md') as fd: 12 | VERSION = fd.readline()[:-1] 13 | except IOError: 14 | VERSION = '0.0.0' 15 | 16 | REQUIRED_PACKAGES = [ 17 | 'numpy', 18 | 'scipy', 19 | 'scikit-image', 20 | 'matplotlib', 21 | 'pillow', 22 | 'pypng', 23 | 'pytest', 24 | 'PyYAML', 25 | 'psutil', 26 | 'tqdm', 27 | 'h5py', 28 | 'easydict >= 1.9', 29 | ] 30 | 31 | with open('README.md', 'r', encoding='utf-8') as fd: 32 | long_desp = fd.read() 33 | 34 | setup( 35 | name='VSR', 36 | version=VERSION, 37 | description='Video Super-Resolution Framework', 38 | long_description=long_desp, 39 | long_description_content_type="text/markdown", 40 | url='https://github.com/LoSealL/VideoSuperResolution', 41 | packages=find_packages(), 42 | install_requires=REQUIRED_PACKAGES, 43 | license='MIT', 44 | author='Wenyi Tang', 45 | author_email='wenyitang@outlook.com', 46 | keywords="super-resolution sr vsr cnn srcnn vespcn", 47 | classifiers=[ 48 | "Programming Language :: Python :: 3", 49 | "License :: OSI Approved :: MIT License", 50 | "Operating System :: OS Independent", 51 | ], 52 | python_requires='>=3.6', 53 | ) 54 | --------------------------------------------------------------------------------