Add MangaScaleV3 on ESRGAN+ arch

.gitignore

@@ -0,0 +1,26 @@
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk

converter.py

@@ -28,7 +28,7 @@ parser = argparse.ArgumentParser(
 )
 parser.add_argument('filename')
 required_args = parser.add_argument_group('required')
-required_args.add_argument('--type', choices=['esrgan_old', 'esrgan_old_lite', 'real_esrgan', 'real_esrgan_compact'], required=True, help='Type of the model')
+required_args.add_argument('--type', choices=['esrgan_old', 'esrgan_old_lite', 'real_esrgan', 'real_esrgan_compact', 'esrgan_plus'], required=True, help='Type of the model')
 required_args.add_argument('--name', type=str, required=True, help='Name of the model')
 required_args.add_argument('--scale', type=int, required=True, help='Scale factor of the model')
 required_args.add_argument('--out-dir', type=str, required=True, help='Output directory')
@@ -118,6 +118,9 @@ elif args.type == 'real_esrgan':
 elif args.type == 'real_esrgan_compact':
     from basicsr.archs.srvgg_arch import SRVGGNetCompact
     torch_model = SRVGGNetCompact(num_in_ch=channels, num_out_ch=channels, num_feat=num_features, num_conv=num_convs, upscale=args.scale, act_type='prelu')
+elif args.type == 'esrgan_plus':
+    from esrgan_plus.codes.models.modules.architecture import RRDBNet
+    torch_model = RRDBNet(in_nc=channels, out_nc=channels, nf=num_features, nb=num_blocks, gc=32, upscale=args.scale)
 else:
     logger.fatal('Unknown model type: %s', args.type)
     sys.exit(-1)

esrgan_plus/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

esrgan_plus/README.md

@@ -0,0 +1,48 @@
+# ESRGAN+ nESRGAN+ Tarsier
+## ICASSP 2020 - ESRGAN+ : Further Improving Enhanced Super-Resolution Generative Adversarial Network
+### [Paper arXiv](https://arxiv.org/abs/2001.08073)
+### [Paper IEEE Xplore](https://ieeexplore.ieee.org/document/9054071)
+## ICPR 2020 - Tarsier: Evolving Noise Injection in Super-Resolution GANs
+### [Paper arXiv](https://arxiv.org/abs/2009.12177)
+
+<p align="center">
+  <img height="250" src="./figures/noise_per_residual_dense_block.PNG">
+</p>
+
+<p align="center">
+  <img src="./figures/qualitative_result.PNG">
+</p>
+
+### Dependencies
+
+- Python 3 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux))
+- [PyTorch >= 1.0.0](https://pytorch.org/)
+- NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads)
+- Python packages: `pip install numpy opencv-python lmdb tensorboardX`
+
+## How to test
+1. Place your low-resolution images in `test_image/LR` folder.
+2. Download pretrained models from [Google Drive](https://drive.google.com/drive/folders/1lNky9afqEP-qdxrAwDFPJ1g0ui4x7Sin?usp=sharing) and place them in `test_image/pretrained_models`.
+2. Run the command: `python test_image/test.py test_image/pretrained_models/nESRGANplus.pth` (or any other models).
+3. The results are in `test_image/results` folder.
+
+
+## How to train
+1. Prepare the datasets which can be downloaded from [Google Drive](https://drive.google.com/drive/folders/1pRmhEmmY-tPF7uH8DuVthfHoApZWJ1QU).
+2. Prepare the PSNR-oriented pretrained model (all pretrained models can be downloaded from [Google Drive](https://drive.google.com/drive/folders/1lNky9afqEP-qdxrAwDFPJ1g0ui4x7Sin?usp=sharing)).
+2. Modify the configuration file `codes/options/train/train_ESRGANplus.json`.
+3. Run the command `python train.py -opt codes/options/train/train_ESRGANplus.json`. 
+
+## Acknowledgement
+- This code is based on [BasicSR](https://github.com/xinntao/BasicSR).
+
+## Citation
+
+    @INPROCEEDINGS{9054071,
+        author = {N. C. {Rakotonirina} and A. {Rasoanaivo}},  
+        booktitle={ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},   
+        title={ESRGAN+ : Further Improving Enhanced Super-Resolution Generative Adversarial Network},   
+        year={2020},  
+        volume={},  
+        number={},  
+        pages={3637-3641},}

esrgan_plus/codes/auto_test.py

@@ -0,0 +1,32 @@
+'''auto test several models.'''
+
+import json
+import os
+
+test_json_path = 'options/test/test_esrgan_auto.json'
+
+
+def modify_json(json_path, model_name, iteration):
+    with open(json_path, 'r+') as json_file:
+        config = json.load(json_file)
+
+        config['name'] = model_name
+        config['datasets']['test_1']['name'] = 'pirm_test_{:d}k'.format(iteration)
+        # config['datasets']['test_1']['dataroot_LR'] = \
+        #   '/home/carraz/datasets/PIRM/PIRM_Test_set/LR'
+        config['path']['pretrain_model_G'] = \
+            '../experiments/{:s}/models/{:d}_G.pth'.format(model_name, iteration*1000)
+        json_file.seek(0)  # rewind
+        json.dump(config, json_file)
+        json_file.truncate()  # if the new data is smaller than the previous
+
+
+model_iter_dict = {}
+model_iter_dict['100_ESRGAN_SRResNet_pristine_pixel10_minc'] = [80, 85, 90, 95]
+
+for model_name, iter_list in model_iter_dict.items():
+    for iteration in iter_list:
+        modify_json(test_json_path, model_name, iteration)
+        # run test scripts
+        print('\n\nTesting {:s} {:d}k...'.format(model_name, iteration))
+        os.system('python test.py -opt ' + test_json_path)

esrgan_plus/codes/data/LRHR_dataset.py

@@ -0,0 +1,128 @@
+import os.path
+import random
+import numpy as np
+import cv2
+import torch
+import torch.utils.data as data
+import data.util as util
+
+
+class LRHRDataset(data.Dataset):
+    '''
+    Read LR and HR image pairs.
+    If only HR image is provided, generate LR image on-the-fly.
+    The pair is ensured by 'sorted' function, so please check the name convention.
+    '''
+
+    def __init__(self, opt):
+        super(LRHRDataset, self).__init__()
+        self.opt = opt
+        self.paths_LR = None
+        self.paths_HR = None
+        self.LR_env = None  # environment for lmdb
+        self.HR_env = None
+
+        # read image list from subset list txt
+        if opt['subset_file'] is not None and opt['phase'] == 'train':
+            with open(opt['subset_file']) as f:
+                self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \
+                        for line in f])
+            if opt['dataroot_LR'] is not None:
+                raise NotImplementedError('Now subset only supports generating LR on-the-fly.')
+        else:  # read image list from lmdb or image files
+            self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR'])
+            self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
+
+        assert self.paths_HR, 'Error: HR path is empty.'
+        if self.paths_LR and self.paths_HR:
+            assert len(self.paths_LR) == len(self.paths_HR), \
+                'HR and LR datasets have different number of images - {}, {}.'.format(\
+                len(self.paths_LR), len(self.paths_HR))
+
+        self.random_scale_list = [1]
+
+    def __getitem__(self, index):
+        HR_path, LR_path = None, None
+        scale = self.opt['scale']
+        HR_size = self.opt['HR_size']
+
+        # get HR image
+        HR_path = self.paths_HR[index]
+        img_HR = util.read_img(self.HR_env, HR_path)
+        # modcrop in the validation / test phase
+        if self.opt['phase'] != 'train':
+            img_HR = util.modcrop(img_HR, scale)
+        # change color space if necessary
+        if self.opt['color']:
+            img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]
+
+        # get LR image
+        if self.paths_LR:
+            LR_path = self.paths_LR[index]
+            img_LR = util.read_img(self.LR_env, LR_path)
+        else:  # down-sampling on-the-fly
+            # randomly scale during training
+            if self.opt['phase'] == 'train':
+                random_scale = random.choice(self.random_scale_list)
+                H_s, W_s, _ = img_HR.shape
+
+                def _mod(n, random_scale, scale, thres):
+                    rlt = int(n * random_scale)
+                    rlt = (rlt // scale) * scale
+                    return thres if rlt < thres else rlt
+
+                H_s = _mod(H_s, random_scale, scale, HR_size)
+                W_s = _mod(W_s, random_scale, scale, HR_size)
+                img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
+                # force to 3 channels
+                if img_HR.ndim == 2:
+                    img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
+
+            H, W, _ = img_HR.shape
+            # using matlab imresize
+            img_LR = util.imresize_np(img_HR, 1 / scale, True)
+            if img_LR.ndim == 2:
+                img_LR = np.expand_dims(img_LR, axis=2)
+
+        if self.opt['phase'] == 'train':
+            # if the image size is too small
+            H, W, _ = img_HR.shape
+            if H < HR_size or W < HR_size:
+                img_HR = cv2.resize(
+                    np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
+                # using matlab imresize
+                img_LR = util.imresize_np(img_HR, 1 / scale, True)
+                if img_LR.ndim == 2:
+                    img_LR = np.expand_dims(img_LR, axis=2)
+
+            H, W, C = img_LR.shape
+            LR_size = HR_size // scale
+
+            # randomly crop
+            rnd_h = random.randint(0, max(0, H - LR_size))
+            rnd_w = random.randint(0, max(0, W - LR_size))
+            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
+            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
+            img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
+
+            # augmentation - flip, rotate
+            img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
+                self.opt['use_rot'])
+
+        # change color space if necessary
+        if self.opt['color']:
+            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        if img_HR.shape[2] == 3:
+            img_HR = img_HR[:, :, [2, 1, 0]]
+            img_LR = img_LR[:, :, [2, 1, 0]]
+        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
+        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+
+        if LR_path is None:
+            LR_path = HR_path
+        return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path}
+
+    def __len__(self):
+        return len(self.paths_HR)

esrgan_plus/codes/data/LRHR_seg_bg_dataset.py

@@ -0,0 +1,149 @@
+import os.path
+import random
+import numpy as np
+import cv2
+import torch
+import torch.utils.data as data
+import data.util as util
+
+
+class LRHRSeg_BG_Dataset(data.Dataset):
+    '''
+    Read HR image, segmentation probability map; generate LR image, category for SFTGAN
+    also sample general scenes for background
+    need to generate LR images on-the-fly
+    '''
+
+    def __init__(self, opt):
+        super(LRHRSeg_BG_Dataset, self).__init__()
+        self.opt = opt
+        self.paths_LR = None
+        self.paths_HR = None
+        self.paths_HR_bg = None  # HR images for background scenes
+        self.LR_env = None  # environment for lmdb
+        self.HR_env = None
+        self.HR_env_bg = None
+
+        # read image list from lmdb or image files
+        self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR'])
+        self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
+        self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'], \
+            opt['dataroot_HR_bg'])
+
+        assert self.paths_HR, 'Error: HR path is empty.'
+        if self.paths_LR and self.paths_HR:
+            assert len(self.paths_LR) == len(self.paths_HR), \
+                'HR and LR datasets have different number of images - {}, {}.'.format(\
+                len(self.paths_LR), len(self.paths_HR))
+
+        self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
+        self.ratio = 10  # 10 OST data samples and 1 DIV2K general data samples(background)
+
+    def __getitem__(self, index):
+        HR_path, LR_path = None, None
+        scale = self.opt['scale']
+        HR_size = self.opt['HR_size']
+
+        # get HR image
+        if self.opt['phase'] == 'train' and \
+                random.choice(list(range(self.ratio))) == 0:  # read background images
+            bg_index = random.randint(0, len(self.paths_HR_bg) - 1)
+            HR_path = self.paths_HR_bg[bg_index]
+            img_HR = util.read_img(self.HR_env_bg, HR_path)
+            seg = torch.FloatTensor(8, img_HR.shape[0], img_HR.shape[1]).fill_(0)
+            seg[0, :, :] = 1  # background
+        else:
+            HR_path = self.paths_HR[index]
+            img_HR = util.read_img(self.HR_env, HR_path)
+            seg = torch.load(HR_path.replace('/img/', '/bicseg/').replace('.png', '.pth'))
+            # read segmentatin files, you should change it to your settings.
+
+        # modcrop in the validation / test phase
+        if self.opt['phase'] != 'train':
+            img_HR = util.modcrop(img_HR, 8)
+
+        seg = np.transpose(seg.numpy(), (1, 2, 0))
+
+        # get LR image
+        if self.paths_LR:
+            LR_path = self.paths_LR[index]
+            img_LR = util.read_img(self.LR_env, LR_path)
+        else:  # down-sampling on-the-fly
+            # randomly scale during training
+            if self.opt['phase'] == 'train':
+                random_scale = random.choice(self.random_scale_list)
+                H_s, W_s, _ = seg.shape
+
+                def _mod(n, random_scale, scale, thres):
+                    rlt = int(n * random_scale)
+                    rlt = (rlt // scale) * scale
+                    return thres if rlt < thres else rlt
+
+                H_s = _mod(H_s, random_scale, scale, HR_size)
+                W_s = _mod(W_s, random_scale, scale, HR_size)
+                img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
+                seg = cv2.resize(np.copy(seg), (W_s, H_s), interpolation=cv2.INTER_NEAREST)
+
+            H, W, _ = img_HR.shape
+            # using matlab imresize
+            img_LR = util.imresize_np(img_HR, 1 / scale, True)
+            if img_LR.ndim == 2:
+                img_LR = np.expand_dims(img_LR, axis=2)
+
+        H, W, C = img_LR.shape
+        if self.opt['phase'] == 'train':
+            LR_size = HR_size // scale
+
+            # randomly crop
+            rnd_h = random.randint(0, max(0, H - LR_size))
+            rnd_w = random.randint(0, max(0, W - LR_size))
+            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
+            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
+            img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
+            seg = seg[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
+
+            # augmentation - flip, rotate
+            img_LR, img_HR, seg = util.augment([img_LR, img_HR, seg], self.opt['use_flip'],
+                                               self.opt['use_rot'])
+
+            # category
+            if 'building' in HR_path:
+                category = 1
+            elif 'plant' in HR_path:
+                category = 2
+            elif 'mountain' in HR_path:
+                category = 3
+            elif 'water' in HR_path:
+                category = 4
+            elif 'sky' in HR_path:
+                category = 5
+            elif 'grass' in HR_path:
+                category = 6
+            elif 'animal' in HR_path:
+                category = 7
+            else:
+                category = 0  # background
+        else:
+            category = -1  # during val, useless
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        if img_HR.shape[2] == 3:
+            img_HR = img_HR[:, :, [2, 1, 0]]
+            img_LR = img_LR[:, :, [2, 1, 0]]
+        img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
+        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+        seg = torch.from_numpy(np.ascontiguousarray(np.transpose(seg, (2, 0, 1)))).float()
+
+        if LR_path is None:
+            LR_path = HR_path
+        return {
+            'LR': img_LR,
+            'HR': img_HR,
+            'seg': seg,
+            'category': category,
+            'LR_path': LR_path,
+            'HR_path': HR_path
+        }
+
+    def __len__(self):
+        return len(self.paths_HR)

esrgan_plus/codes/data/LR_dataset.py

@@ -0,0 +1,40 @@
+import numpy as np
+import torch
+import torch.utils.data as data
+import data.util as util
+
+
+class LRDataset(data.Dataset):
+    '''Read LR images only in the test phase.'''
+
+    def __init__(self, opt):
+        super(LRDataset, self).__init__()
+        self.opt = opt
+        self.paths_LR = None
+        self.LR_env = None  # environment for lmdb
+
+        # read image list from lmdb or image files
+        self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
+        assert self.paths_LR, 'Error: LR paths are empty.'
+
+    def __getitem__(self, index):
+        LR_path = None
+
+        # get LR image
+        LR_path = self.paths_LR[index]
+        img_LR = util.read_img(self.LR_env, LR_path)
+        H, W, C = img_LR.shape
+
+        # change color space if necessary
+        if self.opt['color']:
+            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        if img_LR.shape[2] == 3:
+            img_LR = img_LR[:, :, [2, 1, 0]]
+        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
+
+        return {'LR': img_LR, 'LR_path': LR_path}
+
+    def __len__(self):
+        return len(self.paths_LR)

esrgan_plus/codes/data/__init__.py

@@ -0,0 +1,37 @@
+'''create dataset and dataloader'''
+import logging
+import torch.utils.data
+
+
+def create_dataloader(dataset, dataset_opt):
+    '''create dataloader '''
+    phase = dataset_opt['phase']
+    if phase == 'train':
+        return torch.utils.data.DataLoader(
+            dataset,
+            batch_size=dataset_opt['batch_size'],
+            shuffle=dataset_opt['use_shuffle'],
+            num_workers=dataset_opt['n_workers'],
+            drop_last=True,
+            pin_memory=True)
+    else:
+        return torch.utils.data.DataLoader(
+            dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=True)
+
+
+def create_dataset(dataset_opt):
+    '''create dataset'''
+    mode = dataset_opt['mode']
+    if mode == 'LR':
+        from data.LR_dataset import LRDataset as D
+    elif mode == 'LRHR':
+        from data.LRHR_dataset import LRHRDataset as D
+    elif mode == 'LRHRseg_bg':
+        from data.LRHR_seg_bg_dataset import LRHRSeg_BG_Dataset as D
+    else:
+        raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
+    dataset = D(dataset_opt)
+    logger = logging.getLogger('base')
+    logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
+                                                           dataset_opt['name']))
+    return dataset

esrgan_plus/codes/data/util.py

@@ -0,0 +1,434 @@
+import os
+import math
+import pickle
+import random
+import numpy as np
+import lmdb
+import torch
+import cv2
+import logging
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']
+
+####################
+# Files & IO
+####################
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def _get_paths_from_images(path):
+    assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+    images = []
+    for dirpath, _, fnames in sorted(os.walk(path)):
+        for fname in sorted(fnames):
+            if is_image_file(fname):
+                img_path = os.path.join(dirpath, fname)
+                images.append(img_path)
+    assert images, '{:s} has no valid image file'.format(path)
+    return images
+
+
+def _get_paths_from_lmdb(dataroot):
+    env = lmdb.open(dataroot, readonly=True, lock=False, readahead=False, meminit=False)
+    keys_cache_file = os.path.join(dataroot, '_keys_cache.p')
+    logger = logging.getLogger('base')
+    if os.path.isfile(keys_cache_file):
+        logger.info('Read lmdb keys from cache: {}'.format(keys_cache_file))
+        keys = pickle.load(open(keys_cache_file, "rb"))
+    else:
+        with env.begin(write=False) as txn:
+            logger.info('Creating lmdb keys cache: {}'.format(keys_cache_file))
+            keys = [key.decode('ascii') for key, _ in txn.cursor()]
+        pickle.dump(keys, open(keys_cache_file, 'wb'))
+    paths = sorted([key for key in keys if not key.endswith('.meta')])
+    return env, paths
+
+
+def get_image_paths(data_type, dataroot):
+    env, paths = None, None
+    if dataroot is not None:
+        if data_type == 'lmdb':
+            env, paths = _get_paths_from_lmdb(dataroot)
+        elif data_type == 'img':
+            paths = sorted(_get_paths_from_images(dataroot))
+        else:
+            raise NotImplementedError('data_type [{:s}] is not recognized.'.format(data_type))
+    return env, paths
+
+
+def _read_lmdb_img(env, path):
+    with env.begin(write=False) as txn:
+        buf = txn.get(path.encode('ascii'))
+        buf_meta = txn.get((path + '.meta').encode('ascii')).decode('ascii')
+    img_flat = np.frombuffer(buf, dtype=np.uint8)
+    H, W, C = [int(s) for s in buf_meta.split(',')]
+    img = img_flat.reshape(H, W, C)
+    return img
+
+
+def read_img(env, path):
+    # read image by cv2 or from lmdb
+    # return: Numpy float32, HWC, BGR, [0,1]
+    if env is None:  # img
+        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    else:
+        img = _read_lmdb_img(env, path)
+    img = img.astype(np.float32) / 255.
+    if img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    # some images have 4 channels
+    if img.shape[2] > 3:
+        img = img[:, :, :3]
+    return img
+
+
+####################
+# image processing
+# process on numpy image
+####################
+
+
+def augment(img_list, hflip=True, rot=True):
+    # horizontal flip OR rotate
+    hflip = hflip and random.random() < 0.5
+    vflip = rot and random.random() < 0.5
+    rot90 = rot and random.random() < 0.5
+
+    def _augment(img):
+        if hflip: img = img[:, ::-1, :]
+        if vflip: img = img[::-1, :, :]
+        if rot90: img = img.transpose(1, 0, 2)
+        return img
+
+    return [_augment(img) for img in img_list]
+
+
+def channel_convert(in_c, tar_type, img_list):
+    # conversion among BGR, gray and y
+    if in_c == 3 and tar_type == 'gray':  # BGR to gray
+        gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in gray_list]
+    elif in_c == 3 and tar_type == 'y':  # BGR to y
+        y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in y_list]
+    elif in_c == 1 and tar_type == 'RGB':  # gray/y to BGR
+        return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
+    else:
+        return img_list
+
+
+def rgb2ycbcr(img, only_y=True):
+    '''same as matlab rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+                              [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def bgr2ycbcr(img, only_y=True):
+    '''bgr version of rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+                              [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def ycbcr2rgb(img):
+    '''same as matlab ycbcr2rgb
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                          [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def modcrop(img_in, scale):
+    # img_in: Numpy, HWC or HW
+    img = np.copy(img_in)
+    if img.ndim == 2:
+        H, W = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r]
+    elif img.ndim == 3:
+        H, W, C = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r, :]
+    else:
+        raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
+    return img
+
+
+####################
+# Functions
+####################
+
+
+# matlab 'imresize' function, now only support 'bicubic'
+def cubic(x):
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
+        (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    if (scale < 1) and (antialiasing):
+        # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5+scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    P = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
+        1, P).expand(out_length, P)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
+    # apply cubic kernel
+    if (scale < 1) and (antialiasing):
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, P)
+
+    # If a column in weights is all zero, get rid of it. only consider the first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, P - 2)
+        weights = weights.narrow(1, 1, P - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, P - 2)
+        weights = weights.narrow(1, 0, P - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+def imresize(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: CHW RGB [0,1]
+    # output: CHW RGB [0,1] w/o round
+
+    in_C, in_H, in_W = img.size()
+    out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
+    img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:, :sym_len_Hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_He:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_C, out_H, in_W)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        out_1[0, i, :] = img_aug[0, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+        out_1[1, i, :] = img_aug[1, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+        out_1[2, i, :] = img_aug[2, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
+    out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_Ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_We:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_C, out_H, out_W)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        out_2[0, :, i] = out_1_aug[0, :, idx:idx + kernel_width].mv(weights_W[i])
+        out_2[1, :, i] = out_1_aug[1, :, idx:idx + kernel_width].mv(weights_W[i])
+        out_2[2, :, i] = out_1_aug[2, :, idx:idx + kernel_width].mv(weights_W[i])
+
+    return out_2
+
+
+def imresize_np(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: Numpy, HWC BGR [0,1]
+    # output: HWC BGR [0,1] w/o round
+    img = torch.from_numpy(img)
+
+    in_H, in_W, in_C = img.size()
+    out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
+    img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:sym_len_Hs, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[-sym_len_He:, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(out_H, in_W, in_C)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        out_1[i, :, 0] = img_aug[idx:idx + kernel_width, :, 0].transpose(0, 1).mv(weights_H[i])
+        out_1[i, :, 1] = img_aug[idx:idx + kernel_width, :, 1].transpose(0, 1).mv(weights_H[i])
+        out_1[i, :, 2] = img_aug[idx:idx + kernel_width, :, 2].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
+    out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :sym_len_Ws, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, -sym_len_We:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(out_H, out_W, in_C)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        out_2[:, i, 0] = out_1_aug[:, idx:idx + kernel_width, 0].mv(weights_W[i])
+        out_2[:, i, 1] = out_1_aug[:, idx:idx + kernel_width, 1].mv(weights_W[i])
+        out_2[:, i, 2] = out_1_aug[:, idx:idx + kernel_width, 2].mv(weights_W[i])
+
+    return out_2.numpy()
+
+
+if __name__ == '__main__':
+    # test imresize function
+    # read images
+    img = cv2.imread('test.png')
+    img = img * 1.0 / 255
+    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+    # imresize
+    scale = 1 / 4
+    import time
+    total_time = 0
+    for i in range(10):
+        start_time = time.time()
+        rlt = imresize(img, scale, antialiasing=True)
+        use_time = time.time() - start_time
+        total_time += use_time
+    print('average time: {}'.format(total_time / 10))
+
+    import torchvision.utils
+    torchvision.utils.save_image(
+        (rlt * 255).round() / 255, 'rlt.png', nrow=1, padding=0, normalize=False)

esrgan_plus/codes/models/SFTGAN_ACD_model.py

@@ -0,0 +1,261 @@
+import os
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.optim import lr_scheduler
+
+import models.networks as networks
+from .base_model import BaseModel
+from models.modules.loss import GANLoss, GradientPenaltyLoss
+
+logger = logging.getLogger('base')
+
+
+class SFTGAN_ACD_Model(BaseModel):
+    def __init__(self, opt):
+        super(SFTGAN_ACD_Model, self).__init__(opt)
+        train_opt = opt['train']
+
+        # define networks and load pretrained models
+        self.netG = networks.define_G(opt).to(self.device)  # G
+        if self.is_train:
+            self.netD = networks.define_D(opt).to(self.device)  # D
+            self.netG.train()
+            self.netD.train()
+        self.load()  # load G and D if needed
+
+        # define losses, optimizer and scheduler
+        if self.is_train:
+            # G pixel loss
+            if train_opt['pixel_weight'] > 0:
+                l_pix_type = train_opt['pixel_criterion']
+                if l_pix_type == 'l1':
+                    self.cri_pix = nn.L1Loss().to(self.device)
+                elif l_pix_type == 'l2':
+                    self.cri_pix = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_pix_type))
+                self.l_pix_w = train_opt['pixel_weight']
+            else:
+                logging.info('Remove pixel loss.')
+                self.cri_pix = None
+
+            # G feature loss
+            if train_opt['feature_weight'] > 0:
+                l_fea_type = train_opt['feature_criterion']
+                if l_fea_type == 'l1':
+                    self.cri_fea = nn.L1Loss().to(self.device)
+                elif l_fea_type == 'l2':
+                    self.cri_fea = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
+                self.l_fea_w = train_opt['feature_weight']
+            else:
+                logging.info('Remove feature loss.')
+                self.cri_fea = None
+            if self.cri_fea:  # load VGG perceptual loss
+                self.netF = networks.define_F(opt, use_bn=False).to(self.device)
+
+            # GD gan loss
+            self.cri_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
+            self.l_gan_w = train_opt['gan_weight']
+            # D_update_ratio and D_init_iters are for WGAN
+            self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1
+            self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0
+
+            if train_opt['gan_type'] == 'wgan-gp':
+                self.random_pt = torch.Tensor(1, 1, 1, 1).to(self.device)
+                # gradient penalty loss
+                self.cri_gp = GradientPenaltyLoss(device=self.device).to(self.device)
+                self.l_gp_w = train_opt['gp_weigth']
+
+            # D cls loss
+            self.cri_ce = nn.CrossEntropyLoss(ignore_index=0).to(self.device)
+            # ignore background, since bg images may conflict with other classes
+
+            # optimizers
+            # G
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params_SFT = []
+            optim_params_other = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if 'SFT' in k or 'Cond' in k:
+                    optim_params_SFT.append(v)
+                else:
+                    optim_params_other.append(v)
+            self.optimizer_G_SFT = torch.optim.Adam(optim_params_SFT, lr=train_opt['lr_G']*5, \
+                weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999))
+            self.optimizer_G_other = torch.optim.Adam(optim_params_other, lr=train_opt['lr_G'], \
+                weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999))
+            self.optimizers.append(self.optimizer_G_SFT)
+            self.optimizers.append(self.optimizer_G_other)
+            # D
+            wd_D = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
+            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=train_opt['lr_D'], \
+                weight_decay=wd_D, betas=(train_opt['beta1_D'], 0.999))
+            self.optimizers.append(self.optimizer_D)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
+                        train_opt['lr_steps'], train_opt['lr_gamma']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+        # print network
+        self.print_network()
+
+    def feed_data(self, data, need_HR=True):
+        # LR
+        self.var_L = data['LR'].to(self.device)
+        # seg
+        self.var_seg = data['seg'].to(self.device)
+        # category
+        self.var_cat = data['category'].long().to(self.device)
+
+        if need_HR:  # train or val
+            self.var_H = data['HR'].to(self.device)
+
+    def optimize_parameters(self, step):
+        # G
+        self.optimizer_G_SFT.zero_grad()
+        self.optimizer_G_other.zero_grad()
+        self.fake_H = self.netG((self.var_L, self.var_seg))
+
+        l_g_total = 0
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            if self.cri_pix:  # pixel loss
+                l_g_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.var_H)
+                l_g_total += l_g_pix
+            if self.cri_fea:  # feature loss
+                real_fea = self.netF(self.var_H).detach()
+                fake_fea = self.netF(self.fake_H)
+                l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
+                l_g_total += l_g_fea
+            # G gan + cls loss
+            pred_g_fake, cls_g_fake = self.netD(self.fake_H)
+            l_g_gan = self.l_gan_w * self.cri_gan(pred_g_fake, True)
+            l_g_cls = self.l_gan_w * self.cri_ce(cls_g_fake, self.var_cat)
+            l_g_total += l_g_gan
+            l_g_total += l_g_cls
+
+            l_g_total.backward()
+            self.optimizer_G_SFT.step()
+        if step > 20000:
+            self.optimizer_G_other.step()
+
+        # D
+        self.optimizer_D.zero_grad()
+        l_d_total = 0
+        # real data
+        pred_d_real, cls_d_real = self.netD(self.var_H)
+        l_d_real = self.cri_gan(pred_d_real, True)
+        l_d_cls_real = self.cri_ce(cls_d_real, self.var_cat)
+        # fake data
+        pred_d_fake, cls_d_fake = self.netD(self.fake_H.detach())  # detach to avoid BP to G
+        l_d_fake = self.cri_gan(pred_d_fake, False)
+        l_d_cls_fake = self.cri_ce(cls_d_fake, self.var_cat)
+
+        l_d_total = l_d_real + l_d_cls_real + l_d_fake + l_d_cls_fake
+
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            batch_size = self.var_H.size(0)
+            if self.random_pt.size(0) != batch_size:
+                self.random_pt.resize_(batch_size, 1, 1, 1)
+            self.random_pt.uniform_()  # Draw random interpolation points
+            interp = self.random_pt * self.fake_H.detach() + (1 - self.random_pt) * self.var_H
+            interp.requires_grad = True
+            interp_crit, _ = self.netD(interp)
+            l_d_gp = self.l_gp_w * self.cri_gp(interp, interp_crit)  # maybe wrong in cls?
+            l_d_total += l_d_gp
+
+        l_d_total.backward()
+        self.optimizer_D.step()
+
+        # set log
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            # G
+            if self.cri_pix:
+                self.log_dict['l_g_pix'] = l_g_pix.item()
+            if self.cri_fea:
+                self.log_dict['l_g_fea'] = l_g_fea.item()
+            self.log_dict['l_g_gan'] = l_g_gan.item()
+        # D
+        self.log_dict['l_d_real'] = l_d_real.item()
+        self.log_dict['l_d_fake'] = l_d_fake.item()
+        self.log_dict['l_d_cls_real'] = l_d_cls_real.item()
+        self.log_dict['l_d_cls_fake'] = l_d_cls_fake.item()
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            self.log_dict['l_d_gp'] = l_d_gp.item()
+        # D outputs
+        self.log_dict['D_real'] = torch.mean(pred_d_real.detach())
+        self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG((self.var_L, self.var_seg))
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_HR=True):
+        out_dict = OrderedDict()
+        out_dict['LR'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_HR:
+            out_dict['HR'] = self.var_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        # G
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+
+        logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+        logger.info(s)
+        if self.is_train:
+            # D
+            s, n = self.get_network_description(self.netD)
+            if isinstance(self.netD, nn.DataParallel):
+                net_struc_str = '{} - {}'.format(self.netD.__class__.__name__,
+                                                self.netD.module.__class__.__name__)
+            else:
+                net_struc_str = '{}'.format(self.netD.__class__.__name__)
+
+            logger.info('Network D structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+            if self.cri_fea:  # F, Perceptual Network
+                s, n = self.get_network_description(self.netF)
+                if isinstance(self.netF, nn.DataParallel):
+                    net_struc_str = '{} - {}'.format(self.netF.__class__.__name__,
+                                                    self.netF.module.__class__.__name__)
+                else:
+                    net_struc_str = '{}'.format(self.netF.__class__.__name__)
+
+                logger.info('Network F structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+                logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG)
+        load_path_D = self.opt['path']['pretrain_model_D']
+        if self.opt['is_train'] and load_path_D is not None:
+            logger.info('Loading pretrained model for D [{:s}] ...'.format(load_path_D))
+            self.load_network(load_path_D, self.netD)
+
+    def save(self, iter_step):
+        self.save_network(self.netG, 'G', iter_step)
+        self.save_network(self.netD, 'D', iter_step)

esrgan_plus/codes/models/SRGAN_model.py

@@ -0,0 +1,240 @@
+import os
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.optim import lr_scheduler
+
+import models.networks as networks
+from .base_model import BaseModel
+from models.modules.loss import GANLoss, GradientPenaltyLoss
+
+logger = logging.getLogger('base')
+
+
+class SRGANModel(BaseModel):
+    def __init__(self, opt):
+        super(SRGANModel, self).__init__(opt)
+        train_opt = opt['train']
+
+        # define networks and load pretrained models
+        self.netG = networks.define_G(opt).to(self.device)  # G
+        if self.is_train:
+            self.netD = networks.define_D(opt).to(self.device)  # D
+            self.netG.train()
+            self.netD.train()
+        self.load()  # load G and D if needed
+
+        # define losses, optimizer and scheduler
+        if self.is_train:
+            # G pixel loss
+            if train_opt['pixel_weight'] > 0:
+                l_pix_type = train_opt['pixel_criterion']
+                if l_pix_type == 'l1':
+                    self.cri_pix = nn.L1Loss().to(self.device)
+                elif l_pix_type == 'l2':
+                    self.cri_pix = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_pix_type))
+                self.l_pix_w = train_opt['pixel_weight']
+            else:
+                logger.info('Remove pixel loss.')
+                self.cri_pix = None
+
+            # G feature loss
+            if train_opt['feature_weight'] > 0:
+                l_fea_type = train_opt['feature_criterion']
+                if l_fea_type == 'l1':
+                    self.cri_fea = nn.L1Loss().to(self.device)
+                elif l_fea_type == 'l2':
+                    self.cri_fea = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
+                self.l_fea_w = train_opt['feature_weight']
+            else:
+                logger.info('Remove feature loss.')
+                self.cri_fea = None
+            if self.cri_fea:  # load VGG perceptual loss
+                self.netF = networks.define_F(opt, use_bn=False).to(self.device)
+
+            # GD gan loss
+            self.cri_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
+            self.l_gan_w = train_opt['gan_weight']
+            # D_update_ratio and D_init_iters are for WGAN
+            self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1
+            self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0
+
+            if train_opt['gan_type'] == 'wgan-gp':
+                self.random_pt = torch.Tensor(1, 1, 1, 1).to(self.device)
+                # gradient penalty loss
+                self.cri_gp = GradientPenaltyLoss(device=self.device).to(self.device)
+                self.l_gp_w = train_opt['gp_weigth']
+
+            # optimizers
+            # G
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if v.requires_grad:
+                    optim_params.append(v)
+                else:
+                    logger.warning('Params [{:s}] will not optimize.'.format(k))
+            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], \
+                weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999))
+            self.optimizers.append(self.optimizer_G)
+            # D
+            wd_D = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
+            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=train_opt['lr_D'], \
+                weight_decay=wd_D, betas=(train_opt['beta1_D'], 0.999))
+            self.optimizers.append(self.optimizer_D)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
+                        train_opt['lr_steps'], train_opt['lr_gamma']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+        # print network
+        self.print_network()
+
+    def feed_data(self, data, need_HR=True):
+        # LR
+        self.var_L = data['LR'].to(self.device)
+        if need_HR:  # train or val
+            self.var_H = data['HR'].to(self.device)
+
+            input_ref = data['ref'] if 'ref' in data else data['HR']
+            self.var_ref = input_ref.to(self.device)
+
+    def optimize_parameters(self, step):
+        # G
+        self.optimizer_G.zero_grad()
+        self.fake_H = self.netG(self.var_L)
+
+        l_g_total = 0
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            if self.cri_pix:  # pixel loss
+                l_g_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.var_H)
+                l_g_total += l_g_pix
+            if self.cri_fea:  # feature loss
+                real_fea = self.netF(self.var_H).detach()
+                fake_fea = self.netF(self.fake_H)
+                l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
+                l_g_total += l_g_fea
+            # G gan + cls loss
+            pred_g_fake = self.netD(self.fake_H)
+            l_g_gan = self.l_gan_w * self.cri_gan(pred_g_fake, True)
+            l_g_total += l_g_gan
+
+            l_g_total.backward()
+            self.optimizer_G.step()
+
+        # D
+        self.optimizer_D.zero_grad()
+        l_d_total = 0
+        # real data
+        pred_d_real = self.netD(self.var_ref)
+        l_d_real = self.cri_gan(pred_d_real, True)
+        # fake data
+        pred_d_fake = self.netD(self.fake_H.detach())  # detach to avoid BP to G
+        l_d_fake = self.cri_gan(pred_d_fake, False)
+
+        l_d_total = l_d_real + l_d_fake
+
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            batch_size = self.var_ref.size(0)
+            if self.random_pt.size(0) != batch_size:
+                self.random_pt.resize_(batch_size, 1, 1, 1)
+            self.random_pt.uniform_()  # Draw random interpolation points
+            interp = self.random_pt * self.fake_H.detach() + (1 - self.random_pt) * self.var_ref
+            interp.requires_grad = True
+            interp_crit = self.netD(interp)
+            l_d_gp = self.l_gp_w * self.cri_gp(interp, interp_crit)
+            l_d_total += l_d_gp
+
+        l_d_total.backward()
+        self.optimizer_D.step()
+
+        # set log
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            # G
+            if self.cri_pix:
+                self.log_dict['l_g_pix'] = l_g_pix.item()
+            if self.cri_fea:
+                self.log_dict['l_g_fea'] = l_g_fea.item()
+            self.log_dict['l_g_gan'] = l_g_gan.item()
+        # D
+        self.log_dict['l_d_real'] = l_d_real.item()
+        self.log_dict['l_d_fake'] = l_d_fake.item()
+
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            self.log_dict['l_d_gp'] = l_d_gp.item()
+        # D outputs
+        self.log_dict['D_real'] = torch.mean(pred_d_real.detach())
+        self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG(self.var_L)
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_HR=True):
+        out_dict = OrderedDict()
+        out_dict['LR'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_HR:
+            out_dict['HR'] = self.var_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        # Generator
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+        logger.info(s)
+        if self.is_train:
+            # Discriminator
+            s, n = self.get_network_description(self.netD)
+            if isinstance(self.netD, nn.DataParallel):
+                net_struc_str = '{} - {}'.format(self.netD.__class__.__name__,
+                                                self.netD.module.__class__.__name__)
+            else:
+                net_struc_str = '{}'.format(self.netD.__class__.__name__)
+            logger.info('Network D structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+            if self.cri_fea:  # F, Perceptual Network
+                s, n = self.get_network_description(self.netF)
+                if isinstance(self.netF, nn.DataParallel):
+                    net_struc_str = '{} - {}'.format(self.netF.__class__.__name__,
+                                                    self.netF.module.__class__.__name__)
+                else:
+                    net_struc_str = '{}'.format(self.netF.__class__.__name__)
+                logger.info('Network F structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+                logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG)
+        load_path_D = self.opt['path']['pretrain_model_D']
+        if self.opt['is_train'] and load_path_D is not None:
+            logger.info('Loading pretrained model for D [{:s}] ...'.format(load_path_D))
+            self.load_network(load_path_D, self.netD)
+
+    def save(self, iter_step):
+        self.save_network(self.netG, 'G', iter_step)
+        self.save_network(self.netD, 'D', iter_step)

esrgan_plus/codes/models/SRRaGAN_model.py

@@ -0,0 +1,251 @@
+import os
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.optim import lr_scheduler
+
+import models.networks as networks
+from .base_model import BaseModel
+from models.modules.loss import GANLoss, GradientPenaltyLoss
+logger = logging.getLogger('base')
+
+
+class SRRaGANModel(BaseModel):
+    def __init__(self, opt):
+        super(SRRaGANModel, self).__init__(opt)
+        train_opt = opt['train']
+
+        # define networks and load pretrained models
+        self.netG = networks.define_G(opt).to(self.device)  # G
+        if self.is_train:
+            self.netD = networks.define_D(opt).to(self.device)  # D
+            self.netG.train()
+            self.netD.train()
+        self.load()  # load G and D if needed
+
+        # define losses, optimizer and scheduler
+        if self.is_train:
+            # G pixel loss
+            if train_opt['pixel_weight'] > 0:
+                l_pix_type = train_opt['pixel_criterion']
+                if l_pix_type == 'l1':
+                    self.cri_pix = nn.L1Loss().to(self.device)
+                elif l_pix_type == 'l2':
+                    self.cri_pix = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_pix_type))
+                self.l_pix_w = train_opt['pixel_weight']
+            else:
+                logger.info('Remove pixel loss.')
+                self.cri_pix = None
+
+            # G feature loss
+            if train_opt['feature_weight'] > 0:
+                l_fea_type = train_opt['feature_criterion']
+                if l_fea_type == 'l1':
+                    self.cri_fea = nn.L1Loss().to(self.device)
+                elif l_fea_type == 'l2':
+                    self.cri_fea = nn.MSELoss().to(self.device)
+                else:
+                    raise NotImplementedError('Loss type [{:s}] not recognized.'.format(l_fea_type))
+                self.l_fea_w = train_opt['feature_weight']
+            else:
+                logger.info('Remove feature loss.')
+                self.cri_fea = None
+            if self.cri_fea:  # load VGG perceptual loss
+                self.netF = networks.define_F(opt, use_bn=False).to(self.device)
+
+            # GD gan loss
+            self.cri_gan = GANLoss(train_opt['gan_type'], 1.0, 0.0).to(self.device)
+            self.l_gan_w = train_opt['gan_weight']
+            # D_update_ratio and D_init_iters are for WGAN
+            self.D_update_ratio = train_opt['D_update_ratio'] if train_opt['D_update_ratio'] else 1
+            self.D_init_iters = train_opt['D_init_iters'] if train_opt['D_init_iters'] else 0
+
+            if train_opt['gan_type'] == 'wgan-gp':
+                self.random_pt = torch.Tensor(1, 1, 1, 1).to(self.device)
+                # gradient penalty loss
+                self.cri_gp = GradientPenaltyLoss(device=self.device).to(self.device)
+                self.l_gp_w = train_opt['gp_weigth']
+
+            # optimizers
+            # G
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if v.requires_grad:
+                    optim_params.append(v)
+                else:
+                    logger.warning('Params [{:s}] will not optimize.'.format(k))
+            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'], \
+                weight_decay=wd_G, betas=(train_opt['beta1_G'], 0.999))
+            self.optimizers.append(self.optimizer_G)
+            # D
+            wd_D = train_opt['weight_decay_D'] if train_opt['weight_decay_D'] else 0
+            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=train_opt['lr_D'], \
+                weight_decay=wd_D, betas=(train_opt['beta1_D'], 0.999))
+            self.optimizers.append(self.optimizer_D)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
+                        train_opt['lr_steps'], train_opt['lr_gamma']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+        # print network
+        self.print_network()
+
+    def feed_data(self, data, need_HR=True):
+        # LR
+        self.var_L = data['LR'].to(self.device)
+
+        if need_HR:  # train or val
+            self.var_H = data['HR'].to(self.device)
+
+            input_ref = data['ref'] if 'ref' in data else data['HR']
+            self.var_ref = input_ref.to(self.device)
+
+    def optimize_parameters(self, step):
+        # G
+        for p in self.netD.parameters():
+            p.requires_grad = False
+
+        self.optimizer_G.zero_grad()
+
+        self.fake_H = self.netG(self.var_L)
+
+        l_g_total = 0
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            if self.cri_pix:  # pixel loss
+                l_g_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.var_H)
+                l_g_total += l_g_pix
+            if self.cri_fea:  # feature loss
+                real_fea = self.netF(self.var_H).detach()
+                fake_fea = self.netF(self.fake_H)
+                l_g_fea = self.l_fea_w * self.cri_fea(fake_fea, real_fea)
+                l_g_total += l_g_fea
+            # G gan + cls loss
+            pred_g_fake = self.netD(self.fake_H)
+            pred_d_real = self.netD(self.var_ref).detach()
+
+            l_g_gan = self.l_gan_w * (self.cri_gan(pred_d_real - torch.mean(pred_g_fake), False) +
+                                      self.cri_gan(pred_g_fake - torch.mean(pred_d_real), True)) / 2
+            l_g_total += l_g_gan
+
+            l_g_total.backward()
+            self.optimizer_G.step()
+
+        # D
+        for p in self.netD.parameters():
+            p.requires_grad = True
+
+        self.optimizer_D.zero_grad()
+        l_d_total = 0
+        pred_d_real = self.netD(self.var_ref)
+        pred_d_fake = self.netD(self.fake_H.detach())  # detach to avoid BP to G
+        l_d_real = self.cri_gan(pred_d_real - torch.mean(pred_d_fake), True)
+        l_d_fake = self.cri_gan(pred_d_fake - torch.mean(pred_d_real), False)
+
+        l_d_total = (l_d_real + l_d_fake) / 2
+
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            batch_size = self.var_ref.size(0)
+            if self.random_pt.size(0) != batch_size:
+                self.random_pt.resize_(batch_size, 1, 1, 1)
+            self.random_pt.uniform_()  # Draw random interpolation points
+            interp = self.random_pt * self.fake_H.detach() + (1 - self.random_pt) * self.var_ref
+            interp.requires_grad = True
+            interp_crit = self.netD(interp)
+            l_d_gp = self.l_gp_w * self.cri_gp(interp, interp_crit)
+            l_d_total += l_d_gp
+
+        l_d_total.backward()
+        self.optimizer_D.step()
+
+        # set log
+        if step % self.D_update_ratio == 0 and step > self.D_init_iters:
+            # G
+            if self.cri_pix:
+                self.log_dict['l_g_pix'] = l_g_pix.item()
+            if self.cri_fea:
+                self.log_dict['l_g_fea'] = l_g_fea.item()
+            self.log_dict['l_g_gan'] = l_g_gan.item()
+        # D
+        self.log_dict['l_d_real'] = l_d_real.item()
+        self.log_dict['l_d_fake'] = l_d_fake.item()
+
+        if self.opt['train']['gan_type'] == 'wgan-gp':
+            self.log_dict['l_d_gp'] = l_d_gp.item()
+        # D outputs
+        self.log_dict['D_real'] = torch.mean(pred_d_real.detach())
+        self.log_dict['D_fake'] = torch.mean(pred_d_fake.detach())
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG(self.var_L)
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_HR=True):
+        out_dict = OrderedDict()
+        out_dict['LR'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_HR:
+            out_dict['HR'] = self.var_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        # Generator
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+
+        logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+        logger.info(s)
+        if self.is_train:
+            # Discriminator
+            s, n = self.get_network_description(self.netD)
+            if isinstance(self.netD, nn.DataParallel):
+                net_struc_str = '{} - {}'.format(self.netD.__class__.__name__,
+                                                self.netD.module.__class__.__name__)
+            else:
+                net_struc_str = '{}'.format(self.netD.__class__.__name__)
+
+            logger.info('Network D structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+            if self.cri_fea:  # F, Perceptual Network
+                s, n = self.get_network_description(self.netF)
+                if isinstance(self.netF, nn.DataParallel):
+                    net_struc_str = '{} - {}'.format(self.netF.__class__.__name__,
+                                                    self.netF.module.__class__.__name__)
+                else:
+                    net_struc_str = '{}'.format(self.netF.__class__.__name__)
+
+                logger.info('Network F structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+                logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG)
+        load_path_D = self.opt['path']['pretrain_model_D']
+        if self.opt['is_train'] and load_path_D is not None:
+            logger.info('Loading pretrained model for D [{:s}] ...'.format(load_path_D))
+            self.load_network(load_path_D, self.netD)
+
+    def save(self, iter_step):
+        self.save_network(self.netG, 'G', iter_step)
+        self.save_network(self.netD, 'D', iter_step)

esrgan_plus/codes/models/SR_model.py

@@ -0,0 +1,151 @@
+import os
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.optim import lr_scheduler
+
+import models.networks as networks
+from .base_model import BaseModel
+
+logger = logging.getLogger('base')
+
+
+class SRModel(BaseModel):
+    def __init__(self, opt):
+        super(SRModel, self).__init__(opt)
+        train_opt = opt['train']
+
+        # define network and load pretrained models
+        self.netG = networks.define_G(opt).to(self.device)
+        self.load()
+
+        if self.is_train:
+            self.netG.train()
+
+            # loss
+            loss_type = train_opt['pixel_criterion']
+            if loss_type == 'l1':
+                self.cri_pix = nn.L1Loss().to(self.device)
+            elif loss_type == 'l2':
+                self.cri_pix = nn.MSELoss().to(self.device)
+            else:
+                raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type))
+            self.l_pix_w = train_opt['pixel_weight']
+
+            # optimizers
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+            for k, v in self.netG.named_parameters():  # can optimize for a part of the model
+                if v.requires_grad:
+                    optim_params.append(v)
+                else:
+                    logger.warning('Params [{:s}] will not optimize.'.format(k))
+            self.optimizer_G = torch.optim.Adam(
+                optim_params, lr=train_opt['lr_G'], weight_decay=wd_G)
+            self.optimizers.append(self.optimizer_G)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
+                        train_opt['lr_steps'], train_opt['lr_gamma']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+        # print network
+        self.print_network()
+
+    def feed_data(self, data, need_HR=True):
+        self.var_L = data['LR'].to(self.device)  # LR
+        if need_HR:
+            self.real_H = data['HR'].to(self.device)  # HR
+
+    def optimize_parameters(self, step):
+        self.optimizer_G.zero_grad()
+        self.fake_H = self.netG(self.var_L)
+        l_pix = self.l_pix_w * self.cri_pix(self.fake_H, self.real_H)
+        l_pix.backward()
+        self.optimizer_G.step()
+
+        # set log
+        self.log_dict['l_pix'] = l_pix.item()
+
+    def test(self):
+        self.netG.eval()
+        with torch.no_grad():
+            self.fake_H = self.netG(self.var_L)
+        self.netG.train()
+
+    def test_x8(self):
+        # from https://github.com/thstkdgus35/EDSR-PyTorch
+        self.netG.eval()
+        for k, v in self.netG.named_parameters():
+            v.requires_grad = False
+
+        def _transform(v, op):
+            # if self.precision != 'single': v = v.float()
+            v2np = v.data.cpu().numpy()
+            if op == 'v':
+                tfnp = v2np[:, :, :, ::-1].copy()
+            elif op == 'h':
+                tfnp = v2np[:, :, ::-1, :].copy()
+            elif op == 't':
+                tfnp = v2np.transpose((0, 1, 3, 2)).copy()
+
+            ret = torch.Tensor(tfnp).to(self.device)
+            # if self.precision == 'half': ret = ret.half()
+
+            return ret
+
+        lr_list = [self.var_L]
+        for tf in 'v', 'h', 't':
+            lr_list.extend([_transform(t, tf) for t in lr_list])
+        sr_list = [self.netG(aug) for aug in lr_list]
+        for i in range(len(sr_list)):
+            if i > 3:
+                sr_list[i] = _transform(sr_list[i], 't')
+            if i % 4 > 1:
+                sr_list[i] = _transform(sr_list[i], 'h')
+            if (i % 4) % 2 == 1:
+                sr_list[i] = _transform(sr_list[i], 'v')
+
+        output_cat = torch.cat(sr_list, dim=0)
+        self.fake_H = output_cat.mean(dim=0, keepdim=True)
+
+        for k, v in self.netG.named_parameters():
+            v.requires_grad = True
+        self.netG.train()
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self, need_HR=True):
+        out_dict = OrderedDict()
+        out_dict['LR'] = self.var_L.detach()[0].float().cpu()
+        out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+        if need_HR:
+            out_dict['HR'] = self.real_H.detach()[0].float().cpu()
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+
+        logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+        logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG)
+
+    def save(self, iter_step):
+        self.save_network(self.netG, 'G', iter_step)

esrgan_plus/codes/models/__init__.py

@@ -0,0 +1,20 @@
+import logging
+logger = logging.getLogger('base')
+
+
+def create_model(opt):
+    model = opt['model']
+
+    if model == 'sr':
+        from .SR_model import SRModel as M
+    elif model == 'srgan':
+        from .SRGAN_model import SRGANModel as M
+    elif model == 'srragan':
+        from .SRRaGAN_model import SRRaGANModel as M
+    elif model == 'sftgan':
+        from .SFTGAN_ACD_model import SFTGAN_ACD_Model as M
+    else:
+        raise NotImplementedError('Model [{:s}] not recognized.'.format(model))
+    m = M(opt)
+    logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
+    return m

esrgan_plus/codes/models/base_model.py

@@ -0,0 +1,85 @@
+import os
+import torch
+import torch.nn as nn
+
+
+class BaseModel():
+    def __init__(self, opt):
+        self.opt = opt
+        self.device = torch.device('cuda' if opt['gpu_ids'] is not None else 'cpu')
+        self.is_train = opt['is_train']
+        self.schedulers = []
+        self.optimizers = []
+
+    def feed_data(self, data):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        pass
+
+    def get_current_losses(self):
+        pass
+
+    def print_network(self):
+        pass
+
+    def save(self, label):
+        pass
+
+    def load(self):
+        pass
+
+    def update_learning_rate(self):
+        for scheduler in self.schedulers:
+            scheduler.step()
+
+    def get_current_learning_rate(self):
+        return self.schedulers[0].get_lr()[0]
+
+    def get_network_description(self, network):
+        '''Get the string and total parameters of the network'''
+        if isinstance(network, nn.DataParallel):
+            network = network.module
+        s = str(network)
+        n = sum(map(lambda x: x.numel(), network.parameters()))
+        return s, n
+
+    def save_network(self, network, network_label, iter_step):
+        save_filename = '{}_{}.pth'.format(iter_step, network_label)
+        save_path = os.path.join(self.opt['path']['models'], save_filename)
+        if isinstance(network, nn.DataParallel):
+            network = network.module
+        state_dict = network.state_dict()
+        for key, param in state_dict.items():
+            state_dict[key] = param.cpu()
+        torch.save(state_dict, save_path)
+
+    def load_network(self, load_path, network, strict=True):
+        if isinstance(network, nn.DataParallel):
+            network = network.module
+        network.load_state_dict(torch.load(load_path), strict=strict)
+
+    def save_training_state(self, epoch, iter_step):
+        '''Saves training state during training, which will be used for resuming'''
+        state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []}
+        for s in self.schedulers:
+            state['schedulers'].append(s.state_dict())
+        for o in self.optimizers:
+            state['optimizers'].append(o.state_dict())
+        save_filename = '{}.state'.format(iter_step)
+        save_path = os.path.join(self.opt['path']['training_state'], save_filename)
+        torch.save(state, save_path)
+
+    def resume_training(self, resume_state):
+        '''Resume the optimizers and schedulers for training'''
+        resume_optimizers = resume_state['optimizers']
+        resume_schedulers = resume_state['schedulers']
+        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+        for i, o in enumerate(resume_optimizers):
+            self.optimizers[i].load_state_dict(o)
+        for i, s in enumerate(resume_schedulers):
+            self.schedulers[i].load_state_dict(s)

esrgan_plus/codes/models/modules/architecture.py

@@ -0,0 +1,394 @@
+import math
+import torch
+import torch.nn as nn
+import torchvision
+from . import block as B
+from . import spectral_norm as SN
+
+####################
+# Generator
+####################
+
+
+class SRResNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, upscale=4, norm_type='batch', act_type='relu', \
+            mode='NAC', res_scale=1, upsample_mode='upconv'):
+        super(SRResNet, self).__init__()
+        n_upscale = int(math.log(upscale, 2))
+        if upscale == 3:
+            n_upscale = 1
+
+        fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
+        resnet_blocks = [B.ResNetBlock(nf, nf, nf, norm_type=norm_type, act_type=act_type,\
+            mode=mode, res_scale=res_scale) for _ in range(nb)]
+        LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
+
+        if upsample_mode == 'upconv':
+            upsample_block = B.upconv_blcok
+        elif upsample_mode == 'pixelshuffle':
+            upsample_block = B.pixelshuffle_block
+        else:
+            raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))
+        if upscale == 3:
+            upsampler = upsample_block(nf, nf, 3, act_type=act_type)
+        else:
+            upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
+        HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
+        HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
+
+        self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*resnet_blocks, LR_conv)),\
+            *upsampler, HR_conv0, HR_conv1)
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32, upscale=4, norm_type=None, \
+            act_type='leakyrelu', mode='CNA', upsample_mode='upconv'):
+        super(RRDBNet, self).__init__()
+        n_upscale = int(math.log(upscale, 2))
+        if upscale == 3:
+            n_upscale = 1
+
+        fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
+        rb_blocks = [B.RRDB(nf, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \
+            norm_type=norm_type, act_type=act_type, mode='CNA') for _ in range(nb)]
+        LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
+
+        if upsample_mode == 'upconv':
+            upsample_block = B.upconv_blcok
+        elif upsample_mode == 'pixelshuffle':
+            upsample_block = B.pixelshuffle_block
+        else:
+            raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))
+        if upscale == 3:
+            upsampler = upsample_block(nf, nf, 3, act_type=act_type)
+        else:
+            upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
+        HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
+        HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
+
+        self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),\
+            *upsampler, HR_conv0, HR_conv1)
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+
+####################
+# Discriminator
+####################
+
+
+# VGG style Discriminator with input size 128*128
+class Discriminator_VGG_128(nn.Module):
+    def __init__(self, in_nc, base_nf, norm_type='batch', act_type='leakyrelu', mode='CNA'):
+        super(Discriminator_VGG_128, self).__init__()
+        # features
+        # hxw, c
+        # 128, 64
+        conv0 = B.conv_block(in_nc, base_nf, kernel_size=3, norm_type=None, act_type=act_type, \
+            mode=mode)
+        conv1 = B.conv_block(base_nf, base_nf, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 64, 64
+        conv2 = B.conv_block(base_nf, base_nf*2, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv3 = B.conv_block(base_nf*2, base_nf*2, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 32, 128
+        conv4 = B.conv_block(base_nf*2, base_nf*4, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv5 = B.conv_block(base_nf*4, base_nf*4, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 16, 256
+        conv6 = B.conv_block(base_nf*4, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv7 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 8, 512
+        conv8 = B.conv_block(base_nf*8, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv9 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 4, 512
+        self.features = B.sequential(conv0, conv1, conv2, conv3, conv4, conv5, conv6, conv7, conv8,\
+            conv9)
+
+        # classifier
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 4 * 4, 100), nn.LeakyReLU(0.2, True), nn.Linear(100, 1))
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+
+# VGG style Discriminator with input size 128*128, Spectral Normalization
+class Discriminator_VGG_128_SN(nn.Module):
+    def __init__(self):
+        super(Discriminator_VGG_128_SN, self).__init__()
+        # features
+        # hxw, c
+        # 128, 64
+        self.lrelu = nn.LeakyReLU(0.2, True)
+
+        self.conv0 = SN.spectral_norm(nn.Conv2d(3, 64, 3, 1, 1))
+        self.conv1 = SN.spectral_norm(nn.Conv2d(64, 64, 4, 2, 1))
+        # 64, 64
+        self.conv2 = SN.spectral_norm(nn.Conv2d(64, 128, 3, 1, 1))
+        self.conv3 = SN.spectral_norm(nn.Conv2d(128, 128, 4, 2, 1))
+        # 32, 128
+        self.conv4 = SN.spectral_norm(nn.Conv2d(128, 256, 3, 1, 1))
+        self.conv5 = SN.spectral_norm(nn.Conv2d(256, 256, 4, 2, 1))
+        # 16, 256
+        self.conv6 = SN.spectral_norm(nn.Conv2d(256, 512, 3, 1, 1))
+        self.conv7 = SN.spectral_norm(nn.Conv2d(512, 512, 4, 2, 1))
+        # 8, 512
+        self.conv8 = SN.spectral_norm(nn.Conv2d(512, 512, 3, 1, 1))
+        self.conv9 = SN.spectral_norm(nn.Conv2d(512, 512, 4, 2, 1))
+        # 4, 512
+
+        # classifier
+        self.linear0 = SN.spectral_norm(nn.Linear(512 * 4 * 4, 100))
+        self.linear1 = SN.spectral_norm(nn.Linear(100, 1))
+
+    def forward(self, x):
+        x = self.lrelu(self.conv0(x))
+        x = self.lrelu(self.conv1(x))
+        x = self.lrelu(self.conv2(x))
+        x = self.lrelu(self.conv3(x))
+        x = self.lrelu(self.conv4(x))
+        x = self.lrelu(self.conv5(x))
+        x = self.lrelu(self.conv6(x))
+        x = self.lrelu(self.conv7(x))
+        x = self.lrelu(self.conv8(x))
+        x = self.lrelu(self.conv9(x))
+        x = x.view(x.size(0), -1)
+        x = self.lrelu(self.linear0(x))
+        x = self.linear1(x)
+        return x
+
+
+class Discriminator_VGG_96(nn.Module):
+    def __init__(self, in_nc, base_nf, norm_type='batch', act_type='leakyrelu', mode='CNA'):
+        super(Discriminator_VGG_96, self).__init__()
+        # features
+        # hxw, c
+        # 96, 64
+        conv0 = B.conv_block(in_nc, base_nf, kernel_size=3, norm_type=None, act_type=act_type, \
+            mode=mode)
+        conv1 = B.conv_block(base_nf, base_nf, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 48, 64
+        conv2 = B.conv_block(base_nf, base_nf*2, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv3 = B.conv_block(base_nf*2, base_nf*2, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 24, 128
+        conv4 = B.conv_block(base_nf*2, base_nf*4, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv5 = B.conv_block(base_nf*4, base_nf*4, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 12, 256
+        conv6 = B.conv_block(base_nf*4, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv7 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 6, 512
+        conv8 = B.conv_block(base_nf*8, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv9 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 3, 512
+        self.features = B.sequential(conv0, conv1, conv2, conv3, conv4, conv5, conv6, conv7, conv8,\
+            conv9)
+
+        # classifier
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 3 * 3, 100), nn.LeakyReLU(0.2, True), nn.Linear(100, 1))
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+
+class Discriminator_VGG_192(nn.Module):
+    def __init__(self, in_nc, base_nf, norm_type='batch', act_type='leakyrelu', mode='CNA'):
+        super(Discriminator_VGG_192, self).__init__()
+        # features
+        # hxw, c
+        # 192, 64
+        conv0 = B.conv_block(in_nc, base_nf, kernel_size=3, norm_type=None, act_type=act_type, \
+            mode=mode)
+        conv1 = B.conv_block(base_nf, base_nf, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 96, 64
+        conv2 = B.conv_block(base_nf, base_nf*2, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv3 = B.conv_block(base_nf*2, base_nf*2, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 48, 128
+        conv4 = B.conv_block(base_nf*2, base_nf*4, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv5 = B.conv_block(base_nf*4, base_nf*4, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 24, 256
+        conv6 = B.conv_block(base_nf*4, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv7 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 12, 512
+        conv8 = B.conv_block(base_nf*8, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv9 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 6, 512
+        conv10 = B.conv_block(base_nf*8, base_nf*8, kernel_size=3, stride=1, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        conv11 = B.conv_block(base_nf*8, base_nf*8, kernel_size=4, stride=2, norm_type=norm_type, \
+            act_type=act_type, mode=mode)
+        # 3, 512
+        self.features = B.sequential(conv0, conv1, conv2, conv3, conv4, conv5, conv6, conv7, conv8,\
+            conv9, conv10, conv11)
+
+        # classifier
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 3 * 3, 100), nn.LeakyReLU(0.2, True), nn.Linear(100, 1))
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
+
+
+####################
+# Perceptual Network
+####################
+
+
+# Assume input range is [0, 1]
+class VGGFeatureExtractor(nn.Module):
+    def __init__(self,
+                 feature_layer=34,
+                 use_bn=False,
+                 use_input_norm=True,
+                 device=torch.device('cpu')):
+        super(VGGFeatureExtractor, self).__init__()
+        if use_bn:
+            model = torchvision.models.vgg19_bn(pretrained=True)
+        else:
+            model = torchvision.models.vgg19(pretrained=True)
+        self.use_input_norm = use_input_norm
+        if self.use_input_norm:
+            mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+            # [0.485-1, 0.456-1, 0.406-1] if input in range [-1,1]
+            std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+            # [0.229*2, 0.224*2, 0.225*2] if input in range [-1,1]
+            self.register_buffer('mean', mean)
+            self.register_buffer('std', std)
+        self.features = nn.Sequential(*list(model.features.children())[:(feature_layer + 1)])
+        # No need to BP to variable
+        for k, v in self.features.named_parameters():
+            v.requires_grad = False
+
+    def forward(self, x):
+        if self.use_input_norm:
+            x = (x - self.mean) / self.std
+        output = self.features(x)
+        return output
+
+
+# Assume input range is [0, 1]
+class ResNet101FeatureExtractor(nn.Module):
+    def __init__(self, use_input_norm=True, device=torch.device('cpu')):
+        super(ResNet101FeatureExtractor, self).__init__()
+        model = torchvision.models.resnet101(pretrained=True)
+        self.use_input_norm = use_input_norm
+        if self.use_input_norm:
+            mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+            # [0.485-1, 0.456-1, 0.406-1] if input in range [-1,1]
+            std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+            # [0.229*2, 0.224*2, 0.225*2] if input in range [-1,1]
+            self.register_buffer('mean', mean)
+            self.register_buffer('std', std)
+        self.features = nn.Sequential(*list(model.children())[:8])
+        # No need to BP to variable
+        for k, v in self.features.named_parameters():
+            v.requires_grad = False
+
+    def forward(self, x):
+        if self.use_input_norm:
+            x = (x - self.mean) / self.std
+        output = self.features(x)
+        return output
+
+
+class MINCNet(nn.Module):
+    def __init__(self):
+        super(MINCNet, self).__init__()
+        self.ReLU = nn.ReLU(True)
+        self.conv11 = nn.Conv2d(3, 64, 3, 1, 1)
+        self.conv12 = nn.Conv2d(64, 64, 3, 1, 1)
+        self.maxpool1 = nn.MaxPool2d(2, stride=2, padding=0, ceil_mode=True)
+        self.conv21 = nn.Conv2d(64, 128, 3, 1, 1)
+        self.conv22 = nn.Conv2d(128, 128, 3, 1, 1)
+        self.maxpool2 = nn.MaxPool2d(2, stride=2, padding=0, ceil_mode=True)
+        self.conv31 = nn.Conv2d(128, 256, 3, 1, 1)
+        self.conv32 = nn.Conv2d(256, 256, 3, 1, 1)
+        self.conv33 = nn.Conv2d(256, 256, 3, 1, 1)
+        self.maxpool3 = nn.MaxPool2d(2, stride=2, padding=0, ceil_mode=True)
+        self.conv41 = nn.Conv2d(256, 512, 3, 1, 1)
+        self.conv42 = nn.Conv2d(512, 512, 3, 1, 1)
+        self.conv43 = nn.Conv2d(512, 512, 3, 1, 1)
+        self.maxpool4 = nn.MaxPool2d(2, stride=2, padding=0, ceil_mode=True)
+        self.conv51 = nn.Conv2d(512, 512, 3, 1, 1)
+        self.conv52 = nn.Conv2d(512, 512, 3, 1, 1)
+        self.conv53 = nn.Conv2d(512, 512, 3, 1, 1)
+
+    def forward(self, x):
+        out = self.ReLU(self.conv11(x))
+        out = self.ReLU(self.conv12(out))
+        out = self.maxpool1(out)
+        out = self.ReLU(self.conv21(out))
+        out = self.ReLU(self.conv22(out))
+        out = self.maxpool2(out)
+        out = self.ReLU(self.conv31(out))
+        out = self.ReLU(self.conv32(out))
+        out = self.ReLU(self.conv33(out))
+        out = self.maxpool3(out)
+        out = self.ReLU(self.conv41(out))
+        out = self.ReLU(self.conv42(out))
+        out = self.ReLU(self.conv43(out))
+        out = self.maxpool4(out)
+        out = self.ReLU(self.conv51(out))
+        out = self.ReLU(self.conv52(out))
+        out = self.conv53(out)
+        return out
+
+
+# Assume input range is [0, 1]
+class MINCFeatureExtractor(nn.Module):
+    def __init__(self, feature_layer=34, use_bn=False, use_input_norm=True, \
+                device=torch.device('cpu')):
+        super(MINCFeatureExtractor, self).__init__()
+
+        self.features = MINCNet()
+        self.features.load_state_dict(
+            torch.load('../experiments/pretrained_models/VGG16minc_53.pth'), strict=True)
+        self.features.eval()
+        # No need to BP to variable
+        for k, v in self.features.named_parameters():
+            v.requires_grad = False
+
+    def forward(self, x):
+        output = self.features(x)
+        return output

esrgan_plus/codes/models/modules/block.py

@@ -0,0 +1,322 @@
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import copy
+
+####################
+# Basic blocks
+####################
+
+
+def act(act_type, inplace=True, neg_slope=0.2, n_prelu=1):
+    # helper selecting activation
+    # neg_slope: for leakyrelu and init of prelu
+    # n_prelu: for p_relu num_parameters
+    act_type = act_type.lower()
+    if act_type == 'relu':
+        layer = nn.ReLU(inplace)
+    elif act_type == 'leakyrelu':
+        layer = nn.LeakyReLU(neg_slope, inplace)
+    elif act_type == 'prelu':
+        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)
+    else:
+        raise NotImplementedError('activation layer [{:s}] is not found'.format(act_type))
+    return layer
+
+
+def norm(norm_type, nc):
+    # helper selecting normalization layer
+    norm_type = norm_type.lower()
+    if norm_type == 'batch':
+        layer = nn.BatchNorm2d(nc, affine=True)
+    elif norm_type == 'instance':
+        layer = nn.InstanceNorm2d(nc, affine=False)
+    else:
+        raise NotImplementedError('normalization layer [{:s}] is not found'.format(norm_type))
+    return layer
+
+
+def pad(pad_type, padding):
+    # helper selecting padding layer
+    # if padding is 'zero', do by conv layers
+    pad_type = pad_type.lower()
+    if padding == 0:
+        return None
+    if pad_type == 'reflect':
+        layer = nn.ReflectionPad2d(padding)
+    elif pad_type == 'replicate':
+        layer = nn.ReplicationPad2d(padding)
+    else:
+        raise NotImplementedError('padding layer [{:s}] is not implemented'.format(pad_type))
+    return layer
+
+
+def get_valid_padding(kernel_size, dilation):
+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
+    padding = (kernel_size - 1) // 2
+    return padding
+
+
+class ConcatBlock(nn.Module):
+    # Concat the output of a submodule to its input
+    def __init__(self, submodule):
+        super(ConcatBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x):
+        output = torch.cat((x, self.sub(x)), dim=1)
+        return output
+
+    def __repr__(self):
+        tmpstr = 'Identity .. \n|'
+        modstr = self.sub.__repr__().replace('\n', '\n|')
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+class ShortcutBlock(nn.Module):
+    #Elementwise sum the output of a submodule to its input
+    def __init__(self, submodule):
+        super(ShortcutBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x):
+        output = x + self.sub(x)
+        return output
+
+    def __repr__(self):
+        tmpstr = 'Identity + \n|'
+        modstr = self.sub.__repr__().replace('\n', '\n|')
+        tmpstr = tmpstr + modstr
+        return tmpstr
+
+
+def sequential(*args):
+    # Flatten Sequential. It unwraps nn.Sequential.
+    if len(args) == 1:
+        if isinstance(args[0], OrderedDict):
+            raise NotImplementedError('sequential does not support OrderedDict input.')
+        return args[0]  # No sequential is needed.
+    modules = []
+    for module in args:
+        if isinstance(module, nn.Sequential):
+            for submodule in module.children():
+                modules.append(submodule)
+        elif isinstance(module, nn.Module):
+            modules.append(module)
+    return nn.Sequential(*modules)
+
+class GaussianNoise(nn.Module):
+    def __init__(self, sigma=0.1, is_relative_detach=False):
+        super().__init__()
+        self.sigma = sigma
+        self.is_relative_detach = is_relative_detach
+        self.noise = torch.tensor(0, dtype=torch.float)
+
+    def forward(self, x):
+        if self.training and self.sigma != 0:
+            scale = self.sigma * x.detach() if self.is_relative_detach else self.sigma * x
+            sampled_noise = self.noise.repeat(*x.size()).normal_() * scale
+            x = x + sampled_noise
+        return x 
+
+
+def conv_block(in_nc, out_nc, kernel_size, stride=1, dilation=1, groups=1, bias=True, \
+               pad_type='zero', norm_type=None, act_type='relu', mode='CNA'):
+    '''
+    Conv layer with padding, normalization, activation
+    mode: CNA --> Conv -> Norm -> Act
+        NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16)
+    '''
+    assert mode in ['CNA', 'NAC', 'CNAC'], 'Wong conv mode [{:s}]'.format(mode)
+    padding = get_valid_padding(kernel_size, dilation)
+    p = pad(pad_type, padding) if pad_type and pad_type != 'zero' else None
+    padding = padding if pad_type == 'zero' else 0
+
+    c = nn.Conv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding, \
+            dilation=dilation, bias=bias, groups=groups)
+    a = act(act_type) if act_type else None
+    if 'CNA' in mode:
+        n = norm(norm_type, out_nc) if norm_type else None
+        return sequential(p, c, n, a)
+    elif mode == 'NAC':
+        if norm_type is None and act_type is not None:
+            a = act(act_type, inplace=False)
+            # Important!
+            # input----ReLU(inplace)----Conv--+----output
+            #        |________________________|
+            # inplace ReLU will modify the input, therefore wrong output
+        n = norm(norm_type, in_nc) if norm_type else None
+        return sequential(n, a, p, c)
+  
+def conv1x1(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+# https://github.com/github-pengge/PyTorch-progressive_growing_of_gans/blob/master/models/base_model.py
+class minibatch_std_concat_layer(nn.Module):
+    def __init__(self, averaging='all'):
+        super(minibatch_std_concat_layer, self).__init__()
+        self.averaging = averaging.lower()
+        if 'group' in self.averaging:
+            self.n = int(self.averaging[5:])
+        else:
+            assert self.averaging in ['all', 'flat', 'spatial', 'none', 'gpool'], 'Invalid averaging mode'%self.averaging
+        self.adjusted_std = lambda x, **kwargs: torch.sqrt(torch.mean((x - torch.mean(x, **kwargs)) ** 2, **kwargs) + 1e-8)
+
+    def forward(self, x):
+        shape = list(x.size())
+        target_shape = copy.deepcopy(shape)
+        vals = self.adjusted_std(x, dim=0, keepdim=True)
+        if self.averaging == 'all':
+            target_shape[1] = 1
+            vals = torch.mean(vals, dim=1, keepdim=True)
+        elif self.averaging == 'spatial':
+            if len(shape) == 4:
+                vals = mean(vals, axis=[2,3], keepdim=True)             # torch.mean(torch.mean(vals, 2, keepdim=True), 3, keepdim=True)
+        elif self.averaging == 'none':
+            target_shape = [target_shape[0]] + [s for s in target_shape[1:]]
+        elif self.averaging == 'gpool':
+            if len(shape) == 4:
+                vals = mean(x, [0,2,3], keepdim=True)                   # torch.mean(torch.mean(torch.mean(x, 2, keepdim=True), 3, keepdim=True), 0, keepdim=True)
+        elif self.averaging == 'flat':
+            target_shape[1] = 1
+            vals = torch.FloatTensor([self.adjusted_std(x)])
+        else:                                                           # self.averaging == 'group'
+            target_shape[1] = self.n
+            vals = vals.view(self.n, self.shape[1]/self.n, self.shape[2], self.shape[3])
+            vals = mean(vals, axis=0, keepdim=True).view(1, self.n, 1, 1)
+        vals = vals.expand(*target_shape)
+        return torch.cat([x, vals], 1)
+
+
+####################
+# Useful blocks
+####################
+
+
+class ResNetBlock(nn.Module):
+    '''
+    ResNet Block, 3-3 style
+    with extra residual scaling used in EDSR
+    (Enhanced Deep Residual Networks for Single Image Super-Resolution, CVPRW 17)
+    '''
+
+    def __init__(self, in_nc, mid_nc, out_nc, kernel_size=3, stride=1, dilation=1, groups=1, \
+            bias=True, pad_type='zero', norm_type=None, act_type='relu', mode='CNA', res_scale=1):
+        super(ResNetBlock, self).__init__()
+        conv0 = conv_block(in_nc, mid_nc, kernel_size, stride, dilation, groups, bias, pad_type, \
+            norm_type, act_type, mode)
+        if mode == 'CNA':
+            act_type = None
+        if mode == 'CNAC':  # Residual path: |-CNAC-|
+            act_type = None
+            norm_type = None
+        conv1 = conv_block(mid_nc, out_nc, kernel_size, stride, dilation, groups, bias, pad_type, \
+            norm_type, act_type, mode)
+        # if in_nc != out_nc:
+        #     self.project = conv_block(in_nc, out_nc, 1, stride, dilation, 1, bias, pad_type, \
+        #         None, None)
+        #     print('Need a projecter in ResNetBlock.')
+        # else:
+        #     self.project = lambda x:x
+        self.res = sequential(conv0, conv1)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.res(x).mul(self.res_scale)
+        return x + res
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    '''
+    Residual Dense Block
+    style: 5 convs
+    The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18)
+    '''
+
+    def __init__(self, nc, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \
+            norm_type=None, act_type='leakyrelu', mode='CNA', gaussian_noise=True):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.noise = GaussianNoise() if gaussian_noise else None
+        self.conv1x1 = conv1x1(nc, gc)
+        self.conv1 = conv_block(nc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \
+            norm_type=norm_type, act_type=act_type, mode=mode)
+        self.conv2 = conv_block(nc+gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \
+            norm_type=norm_type, act_type=act_type, mode=mode)
+        self.conv3 = conv_block(nc+2*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \
+            norm_type=norm_type, act_type=act_type, mode=mode)
+        self.conv4 = conv_block(nc+3*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \
+            norm_type=norm_type, act_type=act_type, mode=mode)
+        if mode == 'CNA':
+            last_act = None
+        else:
+            last_act = act_type
+        self.conv5 = conv_block(nc+4*gc, nc, 3, stride, bias=bias, pad_type=pad_type, \
+            norm_type=norm_type, act_type=last_act, mode=mode)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x2 = x2 + self.conv1x1(x)
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x4 = x4 + x2
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return self.noise(x5.mul(0.2) + x)
+
+
+class RRDB(nn.Module):
+    '''
+    Residual in Residual Dense Block
+    (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks)
+    '''
+
+    def __init__(self, nc, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \
+            norm_type=None, act_type='leakyrelu', mode='CNA'):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \
+            norm_type, act_type, mode)
+        self.RDB2 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \
+            norm_type, act_type, mode)
+        self.RDB3 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \
+            norm_type, act_type, mode)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out.mul(0.2) + x
+
+
+####################
+# Upsampler
+####################
+
+
+def pixelshuffle_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \
+                        pad_type='zero', norm_type=None, act_type='relu'):
+    '''
+    Pixel shuffle layer
+    (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional
+    Neural Network, CVPR17)
+    '''
+    conv = conv_block(in_nc, out_nc * (upscale_factor ** 2), kernel_size, stride, bias=bias, \
+                        pad_type=pad_type, norm_type=None, act_type=None)
+    pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+    n = norm(norm_type, out_nc) if norm_type else None
+    a = act(act_type) if act_type else None
+    return sequential(conv, pixel_shuffle, n, a)
+
+
+def upconv_blcok(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \
+                pad_type='zero', norm_type=None, act_type='relu', mode='nearest'):
+    # Up conv
+    # described in https://distill.pub/2016/deconv-checkerboard/
+    upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode)
+    conv = conv_block(in_nc, out_nc, kernel_size, stride, bias=bias, \
+                        pad_type=pad_type, norm_type=norm_type, act_type=act_type)
+    return sequential(upsample, conv)

esrgan_plus/codes/models/modules/loss.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+
+
+# Define GAN loss: [vanilla | lsgan | wgan-gp]
+class GANLoss(nn.Module):
+    def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0):
+        super(GANLoss, self).__init__()
+        self.gan_type = gan_type.lower()
+        self.real_label_val = real_label_val
+        self.fake_label_val = fake_label_val
+
+        if self.gan_type == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif self.gan_type == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif self.gan_type == 'wgan-gp':
+
+            def wgan_loss(input, target):
+                # target is boolean
+                return -1 * input.mean() if target else input.mean()
+
+            self.loss = wgan_loss
+        else:
+            raise NotImplementedError('GAN type [{:s}] is not found'.format(self.gan_type))
+
+    def get_target_label(self, input, target_is_real):
+        if self.gan_type == 'wgan-gp':
+            return target_is_real
+        if target_is_real:
+            return torch.empty_like(input).fill_(self.real_label_val)
+        else:
+            return torch.empty_like(input).fill_(self.fake_label_val)
+
+    def forward(self, input, target_is_real):
+        target_label = self.get_target_label(input, target_is_real)
+        loss = self.loss(input, target_label)
+        return loss
+
+
+class GradientPenaltyLoss(nn.Module):
+    def __init__(self, device=torch.device('cpu')):
+        super(GradientPenaltyLoss, self).__init__()
+        self.register_buffer('grad_outputs', torch.Tensor())
+        self.grad_outputs = self.grad_outputs.to(device)
+
+    def get_grad_outputs(self, input):
+        if self.grad_outputs.size() != input.size():
+            self.grad_outputs.resize_(input.size()).fill_(1.0)
+        return self.grad_outputs
+
+    def forward(self, interp, interp_crit):
+        grad_outputs = self.get_grad_outputs(interp_crit)
+        grad_interp = torch.autograd.grad(outputs=interp_crit, inputs=interp, \
+            grad_outputs=grad_outputs, create_graph=True, retain_graph=True, only_inputs=True)[0]
+        grad_interp = grad_interp.view(grad_interp.size(0), -1)
+        grad_interp_norm = grad_interp.norm(2, dim=1)
+
+        loss = ((grad_interp_norm - 1)**2).mean()
+        return loss

esrgan_plus/codes/models/modules/seg_arch.py

@@ -0,0 +1,70 @@
+'''
+architecture for segmentation
+'''
+import torch.nn as nn
+from . import block as B
+
+
+class Res131(nn.Module):
+    def __init__(self, in_nc, mid_nc, out_nc, dilation=1, stride=1):
+        super(Res131, self).__init__()
+        conv0 = B.conv_block(in_nc, mid_nc, 1, 1, 1, 1, False, 'zero', 'batch')
+        conv1 = B.conv_block(mid_nc, mid_nc, 3, stride, dilation, 1, False, 'zero', 'batch')
+        conv2 = B.conv_block(mid_nc, out_nc, 1, 1, 1, 1, False, 'zero', 'batch', None)  #  No ReLU
+        self.res = B.sequential(conv0, conv1, conv2)
+        if in_nc == out_nc:
+            self.has_proj = False
+        else:
+            self.has_proj = True
+            self.proj = B.conv_block(in_nc, out_nc, 1, stride, 1, 1, False, 'zero', 'batch', None)
+            #  No ReLU
+
+    def forward(self, x):
+        res = self.res(x)
+        if self.has_proj:
+            x = self.proj(x)
+        return nn.functional.relu(x + res, inplace=True)
+
+
+class OutdoorSceneSeg(nn.Module):
+    def __init__(self):
+        super(OutdoorSceneSeg, self).__init__()
+        # conv1
+        blocks = []
+        conv1_1 = B.conv_block(3, 64, 3, 2, 1, 1, False, 'zero', 'batch')  #  /2
+        conv1_2 = B.conv_block(64, 64, 3, 1, 1, 1, False, 'zero', 'batch')
+        conv1_3 = B.conv_block(64, 128, 3, 1, 1, 1, False, 'zero', 'batch')
+        max_pool = nn.MaxPool2d(3, stride=2, padding=0, ceil_mode=True)  #  /2
+        blocks = [conv1_1, conv1_2, conv1_3, max_pool]
+        # conv2, 3 blocks
+        blocks.append(Res131(128, 64, 256))
+        for i in range(2):
+            blocks.append(Res131(256, 64, 256))
+        # conv3, 4 blocks
+        blocks.append(Res131(256, 128, 512, 1, 2))  #  /2
+        for i in range(3):
+            blocks.append(Res131(512, 128, 512))
+        # conv4, 23 blocks
+        blocks.append(Res131(512, 256, 1024, 2))
+        for i in range(22):
+            blocks.append(Res131(1024, 256, 1024, 2))
+        # conv5
+        blocks.append(Res131(1024, 512, 2048, 4))
+        blocks.append(Res131(2048, 512, 2048, 4))
+        blocks.append(Res131(2048, 512, 2048, 4))
+        blocks.append(B.conv_block(2048, 512, 3, 1, 1, 1, False, 'zero', 'batch'))
+        blocks.append(nn.Dropout(0.1))
+        # # conv6
+        blocks.append(nn.Conv2d(512, 8, 1, 1))
+
+        self.feature = B.sequential(*blocks)
+        # deconv
+        self.deconv = nn.ConvTranspose2d(8, 8, 16, 8, 4, 0, 8, False, 1)
+        # softmax
+        self.softmax = nn.Softmax(1)
+
+    def forward(self, x):
+        x = self.feature(x)
+        x = self.deconv(x)
+        x = self.softmax(x)
+        return x

esrgan_plus/codes/models/modules/sft_arch.py

@@ -0,0 +1,226 @@
+'''
+architecture for sft
+'''
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SFTLayer(nn.Module):
+    def __init__(self):
+        super(SFTLayer, self).__init__()
+        self.SFT_scale_conv0 = nn.Conv2d(32, 32, 1)
+        self.SFT_scale_conv1 = nn.Conv2d(32, 64, 1)
+        self.SFT_shift_conv0 = nn.Conv2d(32, 32, 1)
+        self.SFT_shift_conv1 = nn.Conv2d(32, 64, 1)
+
+    def forward(self, x):
+        # x[0]: fea; x[1]: cond
+        scale = self.SFT_scale_conv1(F.leaky_relu(self.SFT_scale_conv0(x[1]), 0.1, inplace=True))
+        shift = self.SFT_shift_conv1(F.leaky_relu(self.SFT_shift_conv0(x[1]), 0.1, inplace=True))
+        return x[0] * (scale + 1) + shift
+
+
+class ResBlock_SFT(nn.Module):
+    def __init__(self):
+        super(ResBlock_SFT, self).__init__()
+        self.sft0 = SFTLayer()
+        self.conv0 = nn.Conv2d(64, 64, 3, 1, 1)
+        self.sft1 = SFTLayer()
+        self.conv1 = nn.Conv2d(64, 64, 3, 1, 1)
+
+    def forward(self, x):
+        # x[0]: fea; x[1]: cond
+        fea = self.sft0(x)
+        fea = F.relu(self.conv0(fea), inplace=True)
+        fea = self.sft1((fea, x[1]))
+        fea = self.conv1(fea)
+        return (x[0] + fea, x[1])  # return a tuple containing features and conditions
+
+
+class SFT_Net(nn.Module):
+    def __init__(self):
+        super(SFT_Net, self).__init__()
+        self.conv0 = nn.Conv2d(3, 64, 3, 1, 1)
+
+        sft_branch = []
+        for i in range(16):
+            sft_branch.append(ResBlock_SFT())
+        sft_branch.append(SFTLayer())
+        sft_branch.append(nn.Conv2d(64, 64, 3, 1, 1))
+        self.sft_branch = nn.Sequential(*sft_branch)
+
+        self.HR_branch = nn.Sequential(
+            nn.Conv2d(64, 256, 3, 1, 1),
+            nn.PixelShuffle(2),
+            nn.ReLU(True),
+            nn.Conv2d(64, 256, 3, 1, 1),
+            nn.PixelShuffle(2),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 3, 3, 1, 1)
+        )
+
+        self.CondNet = nn.Sequential(
+            nn.Conv2d(8, 128, 4, 4),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 32, 1)
+        )
+
+    def forward(self, x):
+        # x[0]: img; x[1]: seg
+        cond = self.CondNet(x[1])
+        fea = self.conv0(x[0])
+        res = self.sft_branch((fea, cond))
+        fea = fea + res
+        out = self.HR_branch(fea)
+        return out
+
+
+# Auxiliary Classifier Discriminator
+class ACD_VGG_BN_96(nn.Module):
+    def __init__(self):
+        super(ACD_VGG_BN_96, self).__init__()
+
+        self.feature = nn.Sequential(
+            nn.Conv2d(3, 64, 3, 1, 1),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(64, 64, 4, 2, 1),
+            nn.BatchNorm2d(64, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(64, 128, 3, 1, 1),
+            nn.BatchNorm2d(128, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(128, 128, 4, 2, 1),
+            nn.BatchNorm2d(128, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(128, 256, 3, 1, 1),
+            nn.BatchNorm2d(256, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(256, 256, 4, 2, 1),
+            nn.BatchNorm2d(256, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(256, 512, 3, 1, 1),
+            nn.BatchNorm2d(512, affine=True),
+            nn.LeakyReLU(0.1, True),
+
+            nn.Conv2d(512, 512, 4, 2, 1),
+            nn.BatchNorm2d(512, affine=True),
+            nn.LeakyReLU(0.1, True),
+        )
+
+        # gan
+        self.gan = nn.Sequential(
+            nn.Linear(512*6*6, 100),
+            nn.LeakyReLU(0.1, True),
+            nn.Linear(100, 1)
+        )
+
+        self.cls = nn.Sequential(
+            nn.Linear(512*6*6, 100),
+            nn.LeakyReLU(0.1, True),
+            nn.Linear(100, 8)
+        )
+
+    def forward(self, x):
+        fea = self.feature(x)
+        fea = fea.view(fea.size(0), -1)
+        gan = self.gan(fea)
+        cls = self.cls(fea)
+        return [gan, cls]
+
+
+#############################################
+# below is the sft arch for the torch version
+#############################################
+
+
+class SFTLayer_torch(nn.Module):
+    def __init__(self):
+        super(SFTLayer_torch, self).__init__()
+        self.SFT_scale_conv0 = nn.Conv2d(32, 32, 1)
+        self.SFT_scale_conv1 = nn.Conv2d(32, 64, 1)
+        self.SFT_shift_conv0 = nn.Conv2d(32, 32, 1)
+        self.SFT_shift_conv1 = nn.Conv2d(32, 64, 1)
+
+    def forward(self, x):
+        # x[0]: fea; x[1]: cond
+        scale = self.SFT_scale_conv1(F.leaky_relu(self.SFT_scale_conv0(x[1]), 0.01, inplace=True))
+        shift = self.SFT_shift_conv1(F.leaky_relu(self.SFT_shift_conv0(x[1]), 0.01, inplace=True))
+        return x[0] * scale + shift
+
+
+class ResBlock_SFT_torch(nn.Module):
+    def __init__(self):
+        super(ResBlock_SFT_torch, self).__init__()
+        self.sft0 = SFTLayer_torch()
+        self.conv0 = nn.Conv2d(64, 64, 3, 1, 1)
+        self.sft1 = SFTLayer_torch()
+        self.conv1 = nn.Conv2d(64, 64, 3, 1, 1)
+
+    def forward(self, x):
+        # x[0]: fea; x[1]: cond
+        fea = F.relu(self.sft0(x), inplace=True)
+        fea = self.conv0(fea)
+        fea = F.relu(self.sft1((fea, x[1])), inplace=True)
+        fea = self.conv1(fea)
+        return (x[0] + fea, x[1])  # return a tuple containing features and conditions
+
+
+class SFT_Net_torch(nn.Module):
+    def __init__(self):
+        super(SFT_Net_torch, self).__init__()
+        self.conv0 = nn.Conv2d(3, 64, 3, 1, 1)
+
+        sft_branch = []
+        for i in range(16):
+            sft_branch.append(ResBlock_SFT_torch())
+        sft_branch.append(SFTLayer_torch())
+        sft_branch.append(nn.Conv2d(64, 64, 3, 1, 1))
+        self.sft_branch = nn.Sequential(*sft_branch)
+
+        self.HR_branch = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='nearest'),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Upsample(scale_factor=2, mode='nearest'),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(64, 3, 3, 1, 1)
+        )
+
+        # Condtion network
+        self.CondNet = nn.Sequential(
+            nn.Conv2d(8, 128, 4, 4),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 128, 1),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(128, 32, 1)
+        )
+
+    def forward(self, x):
+        # x[0]: img; x[1]: seg
+        cond = self.CondNet(x[1])
+        fea = self.conv0(x[0])
+        res = self.sft_branch((fea, cond))
+        fea = fea + res
+        out = self.HR_branch(fea)
+        return out

esrgan_plus/codes/models/modules/spectral_norm.py

@@ -0,0 +1,149 @@
+'''
+Copy from pytorch github repo
+Spectral Normalization from https://arxiv.org/abs/1802.05957
+'''
+import torch
+from torch.nn.functional import normalize
+from torch.nn.parameter import Parameter
+
+
+class SpectralNorm(object):
+    def __init__(self, name='weight', n_power_iterations=1, dim=0, eps=1e-12):
+        self.name = name
+        self.dim = dim
+        if n_power_iterations <= 0:
+            raise ValueError('Expected n_power_iterations to be positive, but '
+                             'got n_power_iterations={}'.format(n_power_iterations))
+        self.n_power_iterations = n_power_iterations
+        self.eps = eps
+
+    def compute_weight(self, module):
+        weight = getattr(module, self.name + '_orig')
+        u = getattr(module, self.name + '_u')
+        weight_mat = weight
+        if self.dim != 0:
+            # permute dim to front
+            weight_mat = weight_mat.permute(self.dim,
+                                            *[d for d in range(weight_mat.dim()) if d != self.dim])
+        height = weight_mat.size(0)
+        weight_mat = weight_mat.reshape(height, -1)
+        with torch.no_grad():
+            for _ in range(self.n_power_iterations):
+                # Spectral norm of weight equals to `u^T W v`, where `u` and `v`
+                # are the first left and right singular vectors.
+                # This power iteration produces approximations of `u` and `v`.
+                v = normalize(torch.matmul(weight_mat.t(), u), dim=0, eps=self.eps)
+                u = normalize(torch.matmul(weight_mat, v), dim=0, eps=self.eps)
+
+        sigma = torch.dot(u, torch.matmul(weight_mat, v))
+        weight = weight / sigma
+        return weight, u
+
+    def remove(self, module):
+        weight = getattr(module, self.name)
+        delattr(module, self.name)
+        delattr(module, self.name + '_u')
+        delattr(module, self.name + '_orig')
+        module.register_parameter(self.name, torch.nn.Parameter(weight))
+
+    def __call__(self, module, inputs):
+        if module.training:
+            weight, u = self.compute_weight(module)
+            setattr(module, self.name, weight)
+            setattr(module, self.name + '_u', u)
+        else:
+            r_g = getattr(module, self.name + '_orig').requires_grad
+            getattr(module, self.name).detach_().requires_grad_(r_g)
+
+    @staticmethod
+    def apply(module, name, n_power_iterations, dim, eps):
+        fn = SpectralNorm(name, n_power_iterations, dim, eps)
+        weight = module._parameters[name]
+        height = weight.size(dim)
+
+        u = normalize(weight.new_empty(height).normal_(0, 1), dim=0, eps=fn.eps)
+        delattr(module, fn.name)
+        module.register_parameter(fn.name + "_orig", weight)
+        # We still need to assign weight back as fn.name because all sorts of
+        # things may assume that it exists, e.g., when initializing weights.
+        # However, we can't directly assign as it could be an nn.Parameter and
+        # gets added as a parameter. Instead, we register weight.data as a
+        # buffer, which will cause weight to be included in the state dict
+        # and also supports nn.init due to shared storage.
+        module.register_buffer(fn.name, weight.data)
+        module.register_buffer(fn.name + "_u", u)
+
+        module.register_forward_pre_hook(fn)
+        return fn
+
+
+def spectral_norm(module, name='weight', n_power_iterations=1, eps=1e-12, dim=None):
+    r"""Applies spectral normalization to a parameter in the given module.
+
+    .. math::
+         \mathbf{W} &= \dfrac{\mathbf{W}}{\sigma(\mathbf{W})} \\
+         \sigma(\mathbf{W}) &= \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2}
+
+    Spectral normalization stabilizes the training of discriminators (critics)
+    in Generaive Adversarial Networks (GANs) by rescaling the weight tensor
+    with spectral norm :math:`\sigma` of the weight matrix calculated using
+    power iteration method. If the dimension of the weight tensor is greater
+    than 2, it is reshaped to 2D in power iteration method to get spectral
+    norm. This is implemented via a hook that calculates spectral norm and
+    rescales weight before every :meth:`~Module.forward` call.
+
+    See `Spectral Normalization for Generative Adversarial Networks`_ .
+
+    .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957
+
+    Args:
+        module (nn.Module): containing module
+        name (str, optional): name of weight parameter
+        n_power_iterations (int, optional): number of power iterations to
+            calculate spectal norm
+        eps (float, optional): epsilon for numerical stability in
+            calculating norms
+        dim (int, optional): dimension corresponding to number of outputs,
+            the default is 0, except for modules that are instances of
+            ConvTranspose1/2/3d, when it is 1
+
+    Returns:
+        The original module with the spectal norm hook
+
+    Example::
+
+        >>> m = spectral_norm(nn.Linear(20, 40))
+        Linear (20 -> 40)
+        >>> m.weight_u.size()
+        torch.Size([20])
+
+    """
+    if dim is None:
+        if isinstance(
+                module,
+            (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):
+            dim = 1
+        else:
+            dim = 0
+    SpectralNorm.apply(module, name, n_power_iterations, dim, eps)
+    return module
+
+
+def remove_spectral_norm(module, name='weight'):
+    r"""Removes the spectral normalization reparameterization from a module.
+
+    Args:
+        module (nn.Module): containing module
+        name (str, optional): name of weight parameter
+
+    Example:
+        >>> m = spectral_norm(nn.Linear(40, 10))
+        >>> remove_spectral_norm(m)
+    """
+    for k, hook in module._forward_pre_hooks.items():
+        if isinstance(hook, SpectralNorm) and hook.name == name:
+            hook.remove(module)
+            del module._forward_pre_hooks[k]
+            return module
+
+    raise ValueError("spectral_norm of '{}' not found in {}".format(name, module))

esrgan_plus/codes/models/networks.py

@@ -0,0 +1,155 @@
+import functools
+import logging
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+import models.modules.architecture as arch
+import models.modules.sft_arch as sft_arch
+logger = logging.getLogger('base')
+####################
+# initialize
+####################
+
+
+def weights_init_normal(m, std=0.02):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.normal_(m.weight.data, 0.0, std)
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('Linear') != -1:
+        init.normal_(m.weight.data, 0.0, std)
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('BatchNorm2d') != -1:
+        init.normal_(m.weight.data, 1.0, std)  # BN also uses norm
+        init.constant_(m.bias.data, 0.0)
+
+
+def weights_init_kaiming(m, scale=1):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+        m.weight.data *= scale
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('Linear') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+        m.weight.data *= scale
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('BatchNorm2d') != -1:
+        init.constant_(m.weight.data, 1.0)
+        init.constant_(m.bias.data, 0.0)
+
+
+def weights_init_orthogonal(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        init.orthogonal_(m.weight.data, gain=1)
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('Linear') != -1:
+        init.orthogonal_(m.weight.data, gain=1)
+        if m.bias is not None:
+            m.bias.data.zero_()
+    elif classname.find('BatchNorm2d') != -1:
+        init.constant_(m.weight.data, 1.0)
+        init.constant_(m.bias.data, 0.0)
+
+
+def init_weights(net, init_type='kaiming', scale=1, std=0.02):
+    # scale for 'kaiming', std for 'normal'.
+    logger.info('Initialization method [{:s}]'.format(init_type))
+    if init_type == 'normal':
+        weights_init_normal_ = functools.partial(weights_init_normal, std=std)
+        net.apply(weights_init_normal_)
+    elif init_type == 'kaiming':
+        weights_init_kaiming_ = functools.partial(weights_init_kaiming, scale=scale)
+        net.apply(weights_init_kaiming_)
+    elif init_type == 'orthogonal':
+        net.apply(weights_init_orthogonal)
+    else:
+        raise NotImplementedError('initialization method [{:s}] not implemented'.format(init_type))
+
+
+####################
+# define network
+####################
+
+
+# Generator
+def define_G(opt):
+    gpu_ids = opt['gpu_ids']
+    opt_net = opt['network_G']
+    which_model = opt_net['which_model_G']
+
+    if which_model == 'sr_resnet':  # SRResNet
+        netG = arch.SRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
+            nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
+            act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
+
+    elif which_model == 'sft_arch':  # SFT-GAN
+        netG = sft_arch.SFT_Net()
+
+    elif which_model == 'RRDB_net':  # RRDB
+        netG = arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'],
+            nb=opt_net['nb'], gc=opt_net['gc'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'],
+            act_type='leakyrelu', mode=opt_net['mode'], upsample_mode='upconv')
+    else:
+        raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))
+
+    if opt['is_train']:
+        init_weights(netG, init_type='kaiming', scale=0.1)
+    if gpu_ids:
+        assert torch.cuda.is_available()
+        netG = nn.DataParallel(netG)
+    return netG
+
+
+# Discriminator
+def define_D(opt):
+    gpu_ids = opt['gpu_ids']
+    opt_net = opt['network_D']
+    which_model = opt_net['which_model_D']
+
+    if which_model == 'discriminator_vgg_128':
+        netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
+            norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
+
+    elif which_model == 'dis_acd':  # sft-gan, Auxiliary Classifier Discriminator
+        netD = sft_arch.ACD_VGG_BN_96()
+
+    elif which_model == 'discriminator_vgg_96':
+        netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
+            norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
+    elif which_model == 'discriminator_vgg_192':
+        netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
+            norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
+    elif which_model == 'discriminator_vgg_128_SN':
+        netD = arch.Discriminator_VGG_128_SN()
+    else:
+        raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
+
+    init_weights(netD, init_type='kaiming', scale=1)
+    if gpu_ids:
+        netD = nn.DataParallel(netD)
+    return netD
+
+
+def define_F(opt, use_bn=False):
+    gpu_ids = opt['gpu_ids']
+    device = torch.device('cuda' if gpu_ids else 'cpu')
+    # pytorch pretrained VGG19-54, before ReLU.
+    if use_bn:
+        feature_layer = 49
+    else:
+        feature_layer = 34
+    netF = arch.VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn, \
+        use_input_norm=True, device=device)
+    # netF = arch.ResNet101FeatureExtractor(use_input_norm=True, device=device)
+    if gpu_ids:
+        netF = nn.DataParallel(netF)
+    netF.eval()  # No need to train
+    return netF

esrgan_plus/codes/options/options.py

@@ -0,0 +1,120 @@
+import os
+import os.path as osp
+import logging
+from collections import OrderedDict
+import json
+
+
+def parse(opt_path, is_train=True):
+    # remove comments starting with '//'
+    json_str = ''
+    with open(opt_path, 'r') as f:
+        for line in f:
+            line = line.split('//')[0] + '\n'
+            json_str += line
+    opt = json.loads(json_str, object_pairs_hook=OrderedDict)
+
+    opt['is_train'] = is_train
+    scale = opt['scale']
+
+    # datasets
+    for phase, dataset in opt['datasets'].items():
+        phase = phase.split('_')[0]
+        dataset['phase'] = phase
+        dataset['scale'] = scale
+        is_lmdb = False
+        if 'dataroot_HR' in dataset and dataset['dataroot_HR'] is not None:
+            dataset['dataroot_HR'] = os.path.expanduser(dataset['dataroot_HR'])
+            if dataset['dataroot_HR'].endswith('lmdb'):
+                is_lmdb = True
+        if 'dataroot_HR_bg' in dataset and dataset['dataroot_HR_bg'] is not None:
+            dataset['dataroot_HR_bg'] = os.path.expanduser(dataset['dataroot_HR_bg'])
+        if 'dataroot_LR' in dataset and dataset['dataroot_LR'] is not None:
+            dataset['dataroot_LR'] = os.path.expanduser(dataset['dataroot_LR'])
+            if dataset['dataroot_LR'].endswith('lmdb'):
+                is_lmdb = True
+        dataset['data_type'] = 'lmdb' if is_lmdb else 'img'
+
+        if phase == 'train' and 'subset_file' in dataset and dataset['subset_file'] is not None:
+            dataset['subset_file'] = os.path.expanduser(dataset['subset_file'])
+
+    # path
+    for key, path in opt['path'].items():
+        if path and key in opt['path']:
+            opt['path'][key] = os.path.expanduser(path)
+    if is_train:
+        experiments_root = os.path.join(opt['path']['root'], 'experiments', opt['name'])
+        opt['path']['experiments_root'] = experiments_root
+        opt['path']['models'] = os.path.join(experiments_root, 'models')
+        opt['path']['training_state'] = os.path.join(experiments_root, 'training_state')
+        opt['path']['log'] = experiments_root
+        opt['path']['val_images'] = os.path.join(experiments_root, 'val_images')
+
+        # change some options for debug mode
+        if 'debug' in opt['name']:
+            opt['train']['val_freq'] = 8
+            opt['logger']['print_freq'] = 2
+            opt['logger']['save_checkpoint_freq'] = 8
+            opt['train']['lr_decay_iter'] = 10
+    else:  # test
+        results_root = os.path.join(opt['path']['root'], 'results', opt['name'])
+        opt['path']['results_root'] = results_root
+        opt['path']['log'] = results_root
+
+    # network
+    opt['network_G']['scale'] = scale
+
+    # export CUDA_VISIBLE_DEVICES
+    gpu_list = ','.join(str(x) for x in opt['gpu_ids'])
+    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+    print('export CUDA_VISIBLE_DEVICES=' + gpu_list)
+
+    return opt
+
+
+class NoneDict(dict):
+    def __missing__(self, key):
+        return None
+
+
+# convert to NoneDict, which return None for missing key.
+def dict_to_nonedict(opt):
+    if isinstance(opt, dict):
+        new_opt = dict()
+        for key, sub_opt in opt.items():
+            new_opt[key] = dict_to_nonedict(sub_opt)
+        return NoneDict(**new_opt)
+    elif isinstance(opt, list):
+        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+    else:
+        return opt
+
+
+def dict2str(opt, indent_l=1):
+    '''dict to string for logger'''
+    msg = ''
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_l * 2) + k + ':[\n'
+            msg += dict2str(v, indent_l + 1)
+            msg += ' ' * (indent_l * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+
+def check_resume(opt):
+    '''Check resume states and pretrain_model paths'''
+    logger = logging.getLogger('base')
+    if opt['path']['resume_state']:
+        if opt['path']['pretrain_model_G'] or opt['path']['pretrain_model_D']:
+            logger.warning('pretrain_model path will be ignored when resuming training.')
+
+        state_idx = osp.basename(opt['path']['resume_state']).split('.')[0]
+        opt['path']['pretrain_model_G'] = osp.join(opt['path']['models'],
+                                                   '{}_G.pth'.format(state_idx))
+        logger.info('Set [pretrain_model_G] to ' + opt['path']['pretrain_model_G'])
+        if 'gan' in opt['model']:
+            opt['path']['pretrain_model_D'] = osp.join(opt['path']['models'],
+                                                       '{}_D.pth'.format(state_idx))
+            logger.info('Set [pretrain_model_D] to ' + opt['path']['pretrain_model_D'])

esrgan_plus/codes/options/test/test_ESRGANplus.json

@@ -0,0 +1,40 @@
+{
+  "name": "nESRGAN+_x4"
+  , "suffix": "_ESRGAN"
+  , "model": "srragan"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "test_1": { // the 1st test dataset
+      "name": "set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+    , "test_2": { // the 2nd test dataset
+      "name": "set14"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set14/Set14"
+      , "dataroot_LR": "/home/carraz/datasets/val_set14/Set14_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    , "pretrain_model_G": "../experiments/pretrained_models/RRDB_ESRGAN_x4.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "RRDB_net" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 23
+    , "in_nc": 3
+    , "out_nc": 3
+
+    , "gc": 32
+    , "group": 1
+  }
+}

esrgan_plus/codes/options/test/test_SRGAN.json

@@ -0,0 +1,37 @@
+{
+  "name": "SRGAN"
+  , "suffix": "_SRGAN"
+  , "model": "srgan"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "test_1": { // the 1st test dataset
+      "name": "set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+    , "test_2": { // the 2nd test dataset
+      "name": "set14"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set14/Set14"
+      , "dataroot_LR": "/home/carraz/datasets/val_set14/Set14_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    , "pretrain_model_G": "../experiments/pretrained_models/SRGAN_bicx4_303_505.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "sr_resnet"
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 16
+    , "in_nc": 3
+    , "out_nc": 3
+  }
+}

esrgan_plus/codes/options/test/test_SRResNet.json

@@ -0,0 +1,40 @@
+{
+  "name": "SRResNet_bicx4_in3nf64nb16"
+  , "suffix": null
+  , "model": "sr"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "test_1": { // the 1st test dataset
+      "name": "set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+    , "test_2": { // the 2nd test dataset
+      "name": "set14"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set14/Set14"
+      , "dataroot_LR": "/home/carraz/datasets/val_set14/Set14_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    , "pretrain_model_G": "../experiments/pretrained_models/SRResNet_bicx4_in3nf64nb16.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "sr_resnet" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 16
+    , "in_nc": 3
+    , "out_nc": 3
+
+    , "gc": 32
+    , "group": 1
+  }
+}

esrgan_plus/codes/options/test/test_sr.json

@@ -0,0 +1,40 @@
+{
+  "name": "RRDB_PSNR_x4"
+  , "suffix": null
+  , "model": "sr"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "test_1": { // the 1st test dataset
+      "name": "set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+    , "test_2": { // the 2nd test dataset
+      "name": "set14"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set14/Set14"
+      , "dataroot_LR": "/home/carraz/datasets/val_set14/Set14_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    , "pretrain_model_G": "../experiments/pretrained_models/RRDB_PSNR_x4.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "RRDB_net" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 23
+    , "in_nc": 3
+    , "out_nc": 3
+
+    , "gc": 32
+    , "group": 1
+  }
+}

esrgan_plus/codes/options/train/train_ESRGANplus.json

@@ -0,0 +1,83 @@
+{
+    "name": "nESRGANplus_x4_DIV2K",
+    "use_tb_logger": true,
+    "model": "srragan",
+    "scale": 4,
+    "gpu_ids": [
+        0
+    ],
+    "datasets": {
+        "train": {
+            "name": "DIV2K",
+            "mode": "LRHR",
+            "dataroot_HR": "/content/gdrive/My Drive/DIV2K_train_HR_sub",
+            "dataroot_LR": "/content/gdrive/My Drive/DIV2K_train_LR_sub_bicLRx4",
+            "subset_file": null,
+            "use_shuffle": true,
+            "n_workers": 8,
+            "batch_size": 16,
+            "HR_size": 128,
+            "use_flip": true,
+            "use_rot": true
+        },
+        "val": {
+            "name": "val_set14_part",
+            "mode": "LRHR",
+            "dataroot_HR": "/content/gdrive/My Drive/ESRGAN/Set14",
+            "dataroot_LR": "/content/gdrive/My Drive/ESRGAN/Set14_LR_sub_bicLRx4"
+        }
+    },
+    "path": {
+        "root": "/content/gdrive/My Drive/ESRGAN/BasicSR",
+        "resume_state": "/content/gdrive/My Drive/ESRGAN/BasicSR/experiments/002_RRDB_ESRGAN_x4_DIV2K/training_state/495000.state",
+        "pretrain_model_G": "/content/gdrive/My Drive/ESRGAN/RRDB_PSNR_x4.pth"
+    },
+    "network_G": {
+        "which_model_G": "RRDB_net",
+        "norm_type": null,
+        "mode": "CNA",
+        "nf": 64,
+        "nb": 23,
+        "in_nc": 3,
+        "out_nc": 3,
+        "gc": 32,
+        "group": 1
+    },
+    "network_D": {
+        "which_model_D": "discriminator_vgg_128",
+        "norm_type": "batch",
+        "act_type": "leakyrelu",
+        "mode": "CNA",
+        "nf": 64,
+        "in_nc": 3
+    },
+    "train": {
+        "lr_G": 0.0001,
+        "weight_decay_G": 0,
+        "beta1_G": 0.9,
+        "lr_D": 0.0001,
+        "weight_decay_D": 0,
+        "beta1_D": 0.9,
+        "lr_scheme": "MultiStepLR",
+        "lr_steps": [
+            50000,
+            100000,
+            200000,
+            300000
+        ],
+        "lr_gamma": 0.5,
+        "pixel_criterion": "l1",
+        "pixel_weight": 0.01,
+        "feature_criterion": "l1",
+        "feature_weight": 1,
+        "gan_type": "vanilla",
+        "gan_weight": 0.005,
+        "manual_seed": 0,
+        "niter": 500000.0,
+        "val_freq": 500.0
+    },
+    "logger": {
+        "print_freq": 50,
+        "save_checkpoint_freq": 500.0
+    }
+}

esrgan_plus/codes/options/train/train_SRGAN.json

@@ -0,0 +1,87 @@
+// Not total the same as SRGAN in <Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network>
+{
+  "name": "debug_002_SRGAN_x4_DIV2K" //  please remove "debug_" during training
+  , "use_tb_logger": true
+  , "model":"srgan"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "train": {
+      "name": "DIV2K"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub.lmdb"
+      , "dataroot_LR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub_bicLRx4.lmdb"
+      , "subset_file": null
+      , "use_shuffle": true
+      , "n_workers": 8
+      , "batch_size": 16
+      , "HR_size": 128
+      , "use_flip": true
+      , "use_rot": true
+    }
+    , "val": {
+      "name": "val_set14_part"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set14_part/Set14"
+      , "dataroot_LR": "/home/carraz/datasets/val_set14_part/Set14_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    // , "resume_state": "../experiments/debug_002_SRGAN_x4_DIV2K/training_state/16.state"
+    , "pretrain_model_G": "../experiments/pretrained_models/SRResNet_bicx4_in3nf64nb16.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "sr_resnet" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 16
+    , "in_nc": 3
+    , "out_nc": 3
+  }
+  , "network_D": {
+    "which_model_D": "discriminator_vgg_128"
+    , "norm_type": "batch"
+    , "act_type": "leakyrelu"
+    , "mode": "CNA"
+    , "nf": 64
+    , "in_nc": 3
+  }
+
+  , "train": {
+    "lr_G": 1e-4
+    , "weight_decay_G": 0
+    , "beta1_G": 0.9
+    , "lr_D": 1e-4
+    , "weight_decay_D": 0
+    , "beta1_D": 0.9
+    , "lr_scheme": "MultiStepLR"
+    , "lr_steps": [50000, 100000, 200000, 300000]
+    , "lr_gamma": 0.5
+
+    , "pixel_criterion": "l1"
+    , "pixel_weight": 1e-2
+    , "feature_criterion": "l1"
+    , "feature_weight": 1
+    , "gan_type": "vanilla"
+    , "gan_weight": 5e-3
+
+    //for wgan-gp
+    // , "D_update_ratio": 1
+    // , "D_init_iters": 0
+    // , "gp_weigth": 10
+
+    , "manual_seed": 0
+    , "niter": 5e5
+    , "val_freq": 5e3
+  }
+
+  , "logger": {
+    "print_freq": 200
+    , "save_checkpoint_freq": 5e3
+  }
+}

esrgan_plus/codes/options/train/train_SRResNet.json

@@ -0,0 +1,66 @@
+// Not total the same as SRResNet in <Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network>
+// With 16 Residual blocks w/o BN
+{
+  "name": "debug_001_SRResNet_PSNR_x4_DIV2K" //  please remove "debug_" during training
+  , "use_tb_logger": true
+  , "model":"sr"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "train": {
+      "name": "DIV2K"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub.lmdb"
+      , "dataroot_LR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub_bicLRx4.lmdb"
+      , "subset_file": null
+      , "use_shuffle": true
+      , "n_workers": 8
+      , "batch_size": 16
+      , "HR_size": 128 // 128 | 192
+      , "use_flip": true
+      , "use_rot": true
+    }
+    , "val": {
+      "name": "val_set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    // , "resume_state": "../experiments/debug_001_RRDB_PSNR_x4_DIV2K/training_state/200.state"
+    , "pretrain_model_G": null
+  }
+
+  , "network_G": {
+    "which_model_G": "sr_resnet" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 16
+    , "in_nc": 3
+    , "out_nc": 3
+  }
+
+  , "train": {
+    "lr_G": 2e-4
+    , "lr_scheme": "MultiStepLR"
+    , "lr_steps": [200000, 400000, 600000, 800000]
+    , "lr_gamma": 0.5
+
+    , "pixel_criterion": "l1"
+    , "pixel_weight": 1.0
+    , "val_freq": 5e3
+
+    , "manual_seed": 0
+    , "niter": 1e6
+  }
+
+  , "logger": {
+    "print_freq": 200
+    , "save_checkpoint_freq": 5e3
+  }
+}

esrgan_plus/codes/options/train/train_sftgan.json

@@ -0,0 +1,76 @@
+{
+  "name": "debug_003_SFTGANx4_OST" //  please remove "debug_" during training
+  , "use_tb_logger": false
+  , "model": "sftgan"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "train": {
+      "name": "OST"
+      , "mode": "LRHRseg_bg"
+      , "dataroot_HR": "/home/carraz/datasets/OST/train/img"
+      , "dataroot_HR_bg": "/home/carraz/datasets/DIV2K800/DIV2K800_sub"
+      , "dataroot_LR": null
+      , "subset_file": null
+      , "use_shuffle": true
+      , "n_workers": 8
+      , "batch_size": 16
+      , "HR_size": 96
+      , "use_flip": true
+      , "use_rot": false
+    }
+    , "val": {
+      "name": "val_OST300_part"
+      , "mode": "LRHRseg_bg"
+      , "dataroot_HR": "/home/carraz/datasets/OST/val/img"
+      , "dataroot_LR": null
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    , "resume_state": null
+    , "pretrain_model_G": "../experiments/pretrained_models/sft_net_ini.pth"
+  }
+
+  , "network_G": {
+    "which_model_G": "sft_arch"
+  }
+  , "network_D": {
+    "which_model_D": "dis_acd"
+  }
+
+  , "train": {
+    "lr_G": 1e-4
+    , "weight_decay_G": 0
+    , "beta1_G": 0.9
+    , "lr_D": 1e-4
+    , "weight_decay_D": 0
+    , "beta1_D": 0.9
+    , "lr_scheme": "MultiStepLR"
+    , "lr_steps": [50000, 100000, 150000, 200000]
+    , "lr_gamma": 0.5
+
+    , "pixel_criterion": "l1"
+    , "pixel_weight": 0
+    , "feature_criterion": "l1"
+    , "feature_weight": 1
+    , "gan_type": "vanilla"
+    , "gan_weight": 5e-3
+
+    //for wgan-gp
+    // , "D_update_ratio": 1
+    // , "D_init_iters": 0
+    // , "gp_weigth": 10
+
+    , "manual_seed": 0
+    , "niter": 6e5
+    , "val_freq": 2e3
+  }
+
+  , "logger": {
+    "print_freq": 200
+    , "save_checkpoint_freq": 2e3
+  }
+}

esrgan_plus/codes/options/train/train_sr.json

@@ -0,0 +1,66 @@
+{
+  "name": "debug_001_RRDB_PSNR_x4_DIV2K" //  please remove "debug_" during training
+  , "use_tb_logger": true
+  , "model":"sr"
+  , "scale": 4
+  , "gpu_ids": [0]
+
+  , "datasets": {
+    "train": {
+      "name": "DIV2K"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub.lmdb"
+      , "dataroot_LR": "/home/carraz/datasets/DIV2K800/DIV2K800_sub_bicLRx4.lmdb"
+      , "subset_file": null
+      , "use_shuffle": true
+      , "n_workers": 8
+      , "batch_size": 16
+      , "HR_size": 128 // 128 | 192
+      , "use_flip": true
+      , "use_rot": true
+    }
+    , "val": {
+      "name": "val_set5"
+      , "mode": "LRHR"
+      , "dataroot_HR": "/home/carraz/datasets/val_set5/Set5"
+      , "dataroot_LR": "/home/carraz/datasets/val_set5/Set5_bicLRx4"
+    }
+  }
+
+  , "path": {
+    "root": "/home/carraz/nESRGANplus"
+    // , "resume_state": "../experiments/debug_001_RRDB_PSNR_x4_DIV2K/training_state/200.state"
+    , "pretrain_model_G": null
+  }
+
+  , "network_G": {
+    "which_model_G": "RRDB_net" // RRDB_net | sr_resnet
+    , "norm_type": null
+    , "mode": "CNA"
+    , "nf": 64
+    , "nb": 23
+    , "in_nc": 3
+    , "out_nc": 3
+    , "gc": 32
+    , "group": 1
+  }
+
+  , "train": {
+    "lr_G": 2e-4
+    , "lr_scheme": "MultiStepLR"
+    , "lr_steps": [200000, 400000, 600000, 800000]
+    , "lr_gamma": 0.5
+
+    , "pixel_criterion": "l1"
+    , "pixel_weight": 1.0
+    , "val_freq": 5e3
+
+    , "manual_seed": 0
+    , "niter": 1e6
+  }
+
+  , "logger": {
+    "print_freq": 200
+    , "save_checkpoint_freq": 5e3
+  }
+}

esrgan_plus/codes/scripts/README.md

@@ -0,0 +1,8 @@
+# Scripts
+We provide some useful scripts here. 
+
+## List
+
+| Name | Description |
+|:---:|:---:|
+| back projection | `Matlab` codes for back projection | 

esrgan_plus/codes/scripts/back_projection/backprojection.m

@@ -0,0 +1,20 @@
+function [im_h] = backprojection(im_h, im_l, maxIter)
+
+[row_l, col_l,~] = size(im_l);
+[row_h, col_h,~] = size(im_h);
+
+p = fspecial('gaussian', 5, 1);
+p = p.^2;
+p = p./sum(p(:));
+
+im_l = double(im_l);
+im_h = double(im_h);
+
+for ii = 1:maxIter
+    im_l_s = imresize(im_h, [row_l, col_l], 'bicubic');
+    im_diff = im_l - im_l_s;
+    im_diff = imresize(im_diff, [row_h, col_h], 'bicubic');
+    im_h(:,:,1) = im_h(:,:,1) + conv2(im_diff(:,:,1), p, 'same');
+    im_h(:,:,2) = im_h(:,:,2) + conv2(im_diff(:,:,2), p, 'same');
+    im_h(:,:,3) = im_h(:,:,3) + conv2(im_diff(:,:,3), p, 'same');
+end

esrgan_plus/codes/scripts/back_projection/main_bp.m

@@ -0,0 +1,22 @@
+clear; close all; clc;
+
+LR_folder = './LR'; % LR
+preout_folder = './results'; % pre output
+save_folder = './results_20bp';
+filepaths  =  dir(fullfile(preout_folder, '*.png'));
+max_iter = 20;
+
+if ~ exist(save_folder, 'dir')
+    mkdir(save_folder);
+end
+
+for idx_im = 1:length(filepaths)
+    fprintf([num2str(idx_im) '\n']);
+    im_name = filepaths(idx_im).name;
+    im_LR = im2double(imread(fullfile(LR_folder, im_name)));
+    im_out = im2double(imread(fullfile(preout_folder, im_name)));
+    %tic
+    im_out = backprojection(im_out, im_LR, max_iter);
+    %toc
+    imwrite(im_out, fullfile(save_folder, im_name));
+end

esrgan_plus/codes/scripts/back_projection/main_reverse_filter.m

@@ -0,0 +1,25 @@
+clear; close all; clc;
+
+LR_folder = './LR'; % LR
+preout_folder = './results'; % pre output
+save_folder = './results_20if';
+filepaths  =  dir(fullfile(preout_folder, '*.png'));
+max_iter = 20;
+
+if ~ exist(save_folder, 'dir')
+    mkdir(save_folder);
+end
+
+for idx_im = 1:length(filepaths)
+    fprintf([num2str(idx_im) '\n']);
+    im_name = filepaths(idx_im).name;
+    im_LR = im2double(imread(fullfile(LR_folder, im_name)));
+    im_out = im2double(imread(fullfile(preout_folder, im_name)));
+    J = imresize(im_LR,4,'bicubic');
+    %tic
+    for m = 1:max_iter
+        im_out = im_out + (J - imresize(imresize(im_out,1/4,'bicubic'),4,'bicubic'));
+    end
+    %toc
+    imwrite(im_out, fullfile(save_folder, im_name));
+end

esrgan_plus/codes/scripts/color2gray.py

@@ -0,0 +1,63 @@
+import os
+import os.path
+import sys
+from multiprocessing import Pool
+import cv2
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from data.util import bgr2ycbcr
+from utils.progress_bar import ProgressBar
+
+
+def main():
+    """A multi-thread tool for converting RGB images to gary/Y images."""
+
+    input_folder = '/home/carraz/datasets/DIV2K800/DIV2K800'
+    save_folder = '/home/carraz/datasets/DIV2K800/DIV2K800_gray'
+    mode = 'gray'  # 'gray' | 'y': Y channel in YCbCr space
+    compression_level = 3  # 3 is the default value in cv2
+    # CV_IMWRITE_PNG_COMPRESSION from 0 to 9. A higher value means a smaller size and longer
+    # compression time. If read raw images during training, use 0 for faster IO speed.
+    n_thread = 20  # thread number
+
+    if not os.path.exists(save_folder):
+        os.makedirs(save_folder)
+        print('mkdir [{:s}] ...'.format(save_folder))
+    else:
+        print('Folder [{:s}] already exists. Exit...'.format(save_folder))
+        sys.exit(1)
+    # print('Parent process {:d}.'.format(os.getpid()))
+
+    img_list = []
+    for root, _, file_list in sorted(os.walk(input_folder)):
+        path = [os.path.join(root, x) for x in file_list]  # assume only images in the input_folder
+        img_list.extend(path)
+
+    def update(arg):
+        pbar.update(arg)
+
+    pbar = ProgressBar(len(img_list))
+
+    pool = Pool(n_thread)
+    for path in img_list:
+        pool.apply_async(worker, args=(path, save_folder, mode, compression_level), callback=update)
+    pool.close()
+    pool.join()
+    print('All subprocesses done.')
+
+
+def worker(path, save_folder, mode, compression_level):
+    img_name = os.path.basename(path)
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGR
+    if mode == 'gray':
+        img_y = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    else:
+        img_y = bgr2ycbcr(img, only_y=True)
+    cv2.imwrite(
+        os.path.join(save_folder, img_name), img_y,
+        [cv2.IMWRITE_PNG_COMPRESSION, compression_level])
+    return 'Processing {:s} ...'.format(img_name)
+
+
+if __name__ == '__main__':
+    main()

esrgan_plus/codes/scripts/create_lmdb.py

@@ -0,0 +1,66 @@
+import sys
+import os.path
+import glob
+import pickle
+import lmdb
+import cv2
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from utils.progress_bar import ProgressBar
+
+# configurations
+img_folder = '/home/carraz/datasets/DIV2K800/DIV2K800/*'  # glob matching pattern
+lmdb_save_path = '/home/carraz/datasets/DIV2K800/DIV2K800.lmdb'  # must end with .lmdb
+mode = 1 # 1 for small data (more memory), 2 for large data (less memory)
+
+img_list = sorted(glob.glob(img_folder))
+
+print('Read images...')
+# mode 1 small data, read all imgs
+if mode == 1:
+    dataset = [cv2.imread(v, cv2.IMREAD_UNCHANGED) for v in img_list]
+    data_size = sum([img.nbytes for img in dataset])
+# mode 2 large data, read imgs later
+elif mode == 2:
+    data_size = sum(os.stat(v).st_size for v in img_list)
+else:
+    raise ValueError('mode should be 1 or 2')
+
+env = lmdb.open(lmdb_save_path, map_size=data_size * 10)
+print('Finish reading {} images.\nWrite lmdb...'.format(len(img_list)))
+
+pbar = ProgressBar(len(img_list))
+batch = 3000   # can be modified according to memory usage
+txn = env.begin(write=True) # txn is a Transaction object
+for i, v in enumerate(img_list):
+    pbar.update('Write {}'.format(v))
+    base_name = os.path.splitext(os.path.basename(v))[0]
+    key = base_name.encode('ascii')
+    data = dataset[i] if mode == 1 else cv2.imread(v, cv2.IMREAD_UNCHANGED)
+    if data.ndim == 2:
+        H, W = data.shape
+        C = 1
+    else:
+        H, W, C = data.shape
+    meta_key = (base_name + '.meta').encode('ascii')
+    meta = '{:d}, {:d}, {:d}'.format(H, W, C)
+    # The encode is only essential in Python 3
+    txn.put(key, data)
+    txn.put(meta_key, meta.encode('ascii'))
+    if mode == 2 and i % batch == batch - 1:
+        txn.commit()
+        txn = env.begin(write=True)
+
+txn.commit()
+env.close()
+
+print('Finish writing lmdb.')
+
+# create keys cache
+keys_cache_file = os.path.join(lmdb_save_path, '_keys_cache.p')
+env = lmdb.open(lmdb_save_path, readonly=True, lock=False, readahead=False, meminit=False)
+with env.begin(write=False) as txn:
+    print('Create lmdb keys cache: {}'.format(keys_cache_file))
+    keys = [key.decode('ascii') for key, _ in txn.cursor()]
+    pickle.dump(keys, open(keys_cache_file, "wb"))
+print('Finish creating lmdb keys cache.')

esrgan_plus/codes/scripts/extract_enlarge_patches.py

@@ -0,0 +1,64 @@
+import os.path
+import glob
+import cv2
+
+crt_path = os.path.dirname(os.path.realpath(__file__))
+
+# configurations
+h_start, h_len = 170, 64
+w_start, w_len = 232, 100
+enlarge_ratio = 3
+line_width = 2
+color = 'yellow'
+
+folder = os.path.join(crt_path, './ori/*')
+save_patch_folder = os.path.join(crt_path, './patch')
+save_rect_folder = os.path.join(crt_path, './rect')
+
+color_tb = {}
+color_tb['yellow'] = (0, 255, 255)
+color_tb['green'] = (0, 255, 0)
+color_tb['red'] = (0, 0, 255)
+color_tb['magenta'] = (255, 0, 255)
+color_tb['matlab_blue'] = (189, 114, 0)
+color_tb['matlab_orange'] = (25, 83, 217)
+color_tb['matlab_yellow'] = (32, 177, 237)
+color_tb['matlab_purple'] = (142, 47, 126)
+color_tb['matlab_green'] = (48, 172, 119)
+color_tb['matlab_liblue'] = (238, 190, 77)
+color_tb['matlab_brown'] = (47, 20, 162)
+color = color_tb[color]
+img_list = glob.glob(folder)
+images = []
+
+# make temp folder
+if not os.path.exists(save_patch_folder):
+    os.makedirs(save_patch_folder)
+    print('mkdir [{}] ...'.format(save_patch_folder))
+if not os.path.exists(save_rect_folder):
+    os.makedirs(save_rect_folder)
+    print('mkdir [{}] ...'.format(save_rect_folder))
+
+for i, path in enumerate(img_list):
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    base_name = os.path.splitext(os.path.basename(path))[0]
+    print(i, base_name)
+    # crop patch
+    if img.ndim == 2:
+        patch = img[h_start:h_start + h_len, w_start:w_start + w_len]
+    elif img.ndim == 3:
+        patch = img[h_start:h_start + h_len, w_start:w_start + w_len, :]
+    else:
+        raise ValueError('Wrong image dim [{:d}]'.format(img.ndim))
+
+    # enlarge patch if necessary
+    if enlarge_ratio > 1:
+        H, W, _ = patch.shape
+        patch = cv2.resize(patch, (W * enlarge_ratio, H * enlarge_ratio), \
+            interpolation=cv2.INTER_CUBIC)
+    cv2.imwrite(os.path.join(save_patch_folder, base_name + '_patch.png'), patch)
+
+    # draw rectangle
+    img_rect = cv2.rectangle(img, (w_start, h_start), (w_start + w_len, h_start + h_len),
+        color, line_width)
+    cv2.imwrite(os.path.join(save_rect_folder, base_name + '_rect.png'), img_rect)

esrgan_plus/codes/scripts/extract_subimgs_single.py

@@ -0,0 +1,88 @@
+import os
+import os.path
+import sys
+from multiprocessing import Pool
+import numpy as np
+import cv2
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from utils.progress_bar import ProgressBar
+
+
+def main():
+    """A multi-thread tool to crop sub imags."""
+    input_folder = '/home/carraz/datasets/DIV2K800/DIV2K800'
+    save_folder = '/home/carraz/datasets/DIV2K800/DIV2K800_sub'
+    n_thread = 20
+    crop_sz = 480
+    step = 240
+    thres_sz = 48
+    compression_level = 3  # 3 is the default value in cv2
+    # CV_IMWRITE_PNG_COMPRESSION from 0 to 9. A higher value means a smaller size and longer
+    # compression time. If read raw images during training, use 0 for faster IO speed.
+
+    if not os.path.exists(save_folder):
+        os.makedirs(save_folder)
+        print('mkdir [{:s}] ...'.format(save_folder))
+    else:
+        print('Folder [{:s}] already exists. Exit...'.format(save_folder))
+        sys.exit(1)
+
+    img_list = []
+    for root, _, file_list in sorted(os.walk(input_folder)):
+        path = [os.path.join(root, x) for x in file_list]  # assume only images in the input_folder
+        img_list.extend(path)
+
+    def update(arg):
+        pbar.update(arg)
+
+    pbar = ProgressBar(len(img_list))
+
+    pool = Pool(n_thread)
+    for path in img_list:
+        pool.apply_async(worker,
+            args=(path, save_folder, crop_sz, step, thres_sz, compression_level),
+            callback=update)
+    pool.close()
+    pool.join()
+    print('All subprocesses done.')
+
+
+def worker(path, save_folder, crop_sz, step, thres_sz, compression_level):
+    img_name = os.path.basename(path)
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+
+    n_channels = len(img.shape)
+    if n_channels == 2:
+        h, w = img.shape
+    elif n_channels == 3:
+        h, w, c = img.shape
+    else:
+        raise ValueError('Wrong image shape - {}'.format(n_channels))
+
+    h_space = np.arange(0, h - crop_sz + 1, step)
+    if h - (h_space[-1] + crop_sz) > thres_sz:
+        h_space = np.append(h_space, h - crop_sz)
+    w_space = np.arange(0, w - crop_sz + 1, step)
+    if w - (w_space[-1] + crop_sz) > thres_sz:
+        w_space = np.append(w_space, w - crop_sz)
+
+    index = 0
+    for x in h_space:
+        for y in w_space:
+            index += 1
+            if n_channels == 2:
+                crop_img = img[x:x + crop_sz, y:y + crop_sz]
+            else:
+                crop_img = img[x:x + crop_sz, y:y + crop_sz, :]
+            crop_img = np.ascontiguousarray(crop_img)
+            # var = np.var(crop_img / 255)
+            # if var > 0.008:
+            #     print(img_name, index_str, var)
+            cv2.imwrite(
+                os.path.join(save_folder, img_name.replace('.png', '_s{:03d}.png'.format(index))),
+                crop_img, [cv2.IMWRITE_PNG_COMPRESSION, compression_level])
+    return 'Processing {:s} ...'.format(img_name)
+
+
+if __name__ == '__main__':
+    main()

esrgan_plus/codes/scripts/generate_mod_LR_bic.m

@@ -0,0 +1,82 @@
+function generate_mod_LR_bic()
+%% matlab code to genetate mod images, bicubic-downsampled LR, bicubic_upsampled images.
+
+%% set parameters
+% comment the unnecessary line
+input_folder = '/home/carraz/datasets/DIV2K800/DIV2K800_sub';
+% save_mod_folder = '';
+save_LR_folder = '/home/carraz/datasets/DIV2K800/DIV2K800_sub_bicLRx4';
+% save_bic_folder = '';
+
+up_scale = 4;
+mod_scale = 4;
+
+if exist('save_mod_folder', 'var')
+    if exist(save_mod_folder, 'dir')
+        disp(['It will cover ', save_mod_folder]);
+    else
+        mkdir(save_mod_folder);
+    end
+end
+if exist('save_LR_folder', 'var')
+    if exist(save_LR_folder, 'dir')
+        disp(['It will cover ', save_LR_folder]);
+    else
+        mkdir(save_LR_folder);
+    end
+end
+if exist('save_bic_folder', 'var')
+    if exist(save_bic_folder, 'dir')
+        disp(['It will cover ', save_bic_folder]);
+    else
+        mkdir(save_bic_folder);
+    end
+end
+
+idx = 0;
+filepaths = dir(fullfile(input_folder,'*.*'));
+for i = 1 : length(filepaths)
+    [paths,imname,ext] = fileparts(filepaths(i).name);
+    if isempty(imname)
+        disp('Ignore . folder.');
+    elseif strcmp(imname, '.')
+        disp('Ignore .. folder.');
+    else
+        idx = idx + 1;
+        str_rlt = sprintf('%d\t%s.\n', idx, imname);
+        fprintf(str_rlt);
+        % read image
+        img = imread(fullfile(input_folder, [imname, ext]));
+        img = im2double(img);
+        % modcrop
+        img = modcrop(img, mod_scale);
+        if exist('save_mod_folder', 'var')
+            imwrite(img, fullfile(save_mod_folder, [imname, '.png']));
+        end
+        % LR
+        im_LR = imresize(img, 1/up_scale, 'bicubic');
+        if exist('save_LR_folder', 'var')
+            imwrite(im_LR, fullfile(save_LR_folder, [imname, '_bicLRx4.png']));
+        end
+        % Bicubic
+        if exist('save_bic_folder', 'var')
+            im_B = imresize(im_LR, up_scale, 'bicubic');
+            imwrite(im_B, fullfile(save_bic_folder, [imname, '_bicx4.png']));
+        end
+    end
+end
+end
+
+%% modcrop
+function img = modcrop(img, modulo)
+if size(img,3) == 1
+    sz = size(img);
+    sz = sz - mod(sz, modulo);
+    img = img(1:sz(1), 1:sz(2));
+else
+    tmpsz = size(img);
+    sz = tmpsz(1:2);
+    sz = sz - mod(sz, modulo);
+    img = img(1:sz(1), 1:sz(2),:);
+end
+end

esrgan_plus/codes/scripts/generate_mod_LR_bic.py

@@ -0,0 +1,74 @@
+import os, sys
+import cv2
+import numpy as np
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from data.util import imresize_np
+
+def generate_mod_LR_bic():
+    # set parameters
+    up_scale = 4
+    mod_scale = 4
+    # set data dir
+    sourcedir = '/data/datasets/img'
+    savedir = '/data/datasets/mod'
+
+    saveHRpath = os.path.join(savedir, 'HR', 'x'+str(mod_scale))
+    saveLRpath = os.path.join(savedir, 'LR', 'x'+str(up_scale))
+    saveBicpath = os.path.join(savedir, 'Bic', 'x'+str(up_scale))
+    
+    if not os.path.isdir(sourcedir):
+        print('Error: No source data found')
+        exit(0)
+    if not os.path.isdir(savedir):
+        os.mkdir(savedir)
+        
+    if not os.path.isdir(os.path.join(savedir, 'HR')):    
+        os.mkdir(os.path.join(savedir, 'HR'))   
+    if not os.path.isdir(os.path.join(savedir, 'LR')):
+        os.mkdir(os.path.join(savedir, 'LR'))
+    if not os.path.isdir(os.path.join(savedir, 'Bic')):
+        os.mkdir(os.path.join(savedir, 'Bic'))
+    
+    if not os.path.isdir(saveHRpath):
+        os.mkdir(saveHRpath)
+    else:
+        print('It will cover '+str(saveHRpath))
+    
+    if not os.path.isdir(saveLRpath):
+        os.mkdir(saveLRpath)
+    else:
+        print('It will cover '+str(saveLRpath))
+    
+    if not os.path.isdir(saveBicpath):
+        os.mkdir(saveBicpath)
+    else:
+        print('It will cover '+str(saveBicpath))
+    
+    filepaths = [f for f in os.listdir(sourcedir) if f.endswith('.png')]
+    num_files = len(filepaths)
+    
+    # prepare data with augementation
+    for i in range(num_files):
+        filename = filepaths[i]
+        print('No.{} -- Processing {}'.format(i, filename))
+        # read image
+        image = cv2.imread(os.path.join(sourcedir, filename))
+    
+        width = int(np.floor(image.shape[1] / mod_scale))
+        height = int(np.floor(image.shape[0] / mod_scale))
+        # modcrop
+        if len(image.shape) == 3:
+            image_HR = image[0:mod_scale*height, 0:mod_scale*width,:]
+        else:
+            image_HR = image[0:mod_scale*height, 0:mod_scale*width]
+        # LR
+        image_LR = imresize_np(image_HR, 1/up_scale, True)
+        # bic
+        image_Bic = imresize_np(image_LR, up_scale, True)
+    
+        cv2.imwrite(os.path.join(saveHRpath, filename), image_HR)
+        cv2.imwrite(os.path.join(saveLRpath, filename), image_LR)
+        cv2.imwrite(os.path.join(saveBicpath, filename), image_Bic)
+
+if __name__ == "__main__":
+    generate_mod_LR_bic()

esrgan_plus/codes/scripts/make_gif_video.py

@@ -0,0 +1,106 @@
+"""
+Add text to images, then make gif/video sequence from images.
+
+Since the created gif has low quality with color issues, use this script to generate image with
+text and then use `gifski`.
+
+Call `ffmpeg` to make video.
+"""
+
+import os.path
+import numpy as np
+import cv2
+
+
+crt_path = os.path.dirname(os.path.realpath(__file__))
+
+# configurations
+img_name_list = ['x1', 'x2', 'x3', 'x4', 'x5']
+ext = '.png'
+text_list = ['1', '2', '3', '4', '5']
+h_start, h_len = 0, 576
+w_start, w_len = 10, 352
+enlarge_ratio = 1
+txt_pos = (10, 50)  # w, h
+font_size = 1.5
+font_thickness = 4
+color = 'red'
+duration = 0.8  # second
+use_imageio = False  # use imageio to make gif
+make_video = False  # make video using ffmpeg
+
+is_crop = True
+if h_start == 0 or w_start == 0:
+    is_crop = False # do not crop
+
+img_name_list = [x + ext for x in img_name_list]
+input_folder = os.path.join(crt_path, './ori')
+save_folder = os.path.join(crt_path, './ori')
+color_tb = {}
+color_tb['yellow'] = (0, 255, 255)
+color_tb['green'] = (0, 255, 0)
+color_tb['red'] = (0, 0, 255)
+color_tb['magenta'] = (255, 0, 255)
+color_tb['matlab_blue'] = (189, 114, 0)
+color_tb['matlab_orange'] = (25, 83, 217)
+color_tb['matlab_yellow'] = (32, 177, 237)
+color_tb['matlab_purple'] = (142, 47, 126)
+color_tb['matlab_green'] = (48, 172, 119)
+color_tb['matlab_liblue'] = (238, 190, 77)
+color_tb['matlab_brown'] = (47, 20, 162)
+color = color_tb[color]
+
+img_list = []
+
+# make temp dir
+if not os.path.exists(save_folder):
+    os.makedirs(save_folder)
+    print('mkdir [{}] ...'.format(save_folder))
+if make_video:
+    # tmp folder to save images for video
+    tmp_video_folder = os.path.join(crt_path, '_tmp_video')
+    if not os.path.exists(tmp_video_folder):
+        os.makedirs(tmp_video_folder)
+
+idx = 0
+for img_name, write_txt in zip(img_name_list, text_list):
+    img = cv2.imread(os.path.join(input_folder, img_name), cv2.IMREAD_UNCHANGED)
+    base_name = os.path.splitext(img_name)[0]
+    print(base_name)
+    # crop image
+    if is_crop:
+        print('Crop image ...')
+        if img.ndim == 2:
+            img = img[h_start:h_start + h_len, w_start:w_start + w_len]
+        elif img.ndim == 3:
+            img = img[h_start:h_start + h_len, w_start:w_start + w_len, :]
+        else:
+            raise ValueError('Wrong image dim [{:d}]'.format(img.ndim))
+
+    # enlarge img if necessary
+    if enlarge_ratio > 1:
+        H, W, _ = img.shape
+        img = cv2.resize(img, (W * enlarge_ratio, H * enlarge_ratio), \
+            interpolation=cv2.INTER_CUBIC)
+
+    # add text
+    font = cv2.FONT_HERSHEY_COMPLEX
+    cv2.putText(img, write_txt, txt_pos, font, font_size, color, font_thickness, cv2.LINE_AA)
+    cv2.imwrite(os.path.join(save_folder, base_name + '_text.png'), img)
+    if make_video:
+        idx += 1
+        cv2.imwrite(os.path.join(tmp_video_folder, '{:05d}.png'.format(idx)), img)
+
+    img = np.ascontiguousarray(img[:, :, [2, 1, 0]])
+    img_list.append(img)
+
+if use_imageio:
+    import imageio
+    imageio.mimsave(os.path.join(save_folder, 'out.gif'), img_list, format='GIF', duration=duration)
+
+if make_video:
+    os.system('ffmpeg -r {:f} -i {:s}/%05d.png -vcodec mpeg4 -y {:s}/movie.mp4'.format(
+        1 / duration, tmp_video_folder, save_folder))
+
+if os.path.exists(tmp_video_folder):
+    os.system('rm -rf {}'.format(tmp_video_folder))

esrgan_plus/codes/scripts/net_interp.py

@@ -0,0 +1,20 @@
+import torch
+from collections import OrderedDict
+
+alpha = 0.8
+
+net_PSNR_path = './models/RRDB_PSNR_x4.pth'
+net_ESRGAN_path = './models/RRDB_ESRGAN_x4.pth'
+net_interp_path = './models/interp_{:02d}.pth'.format(int(alpha*10))
+
+net_PSNR = torch.load(net_PSNR_path)
+net_ESRGAN = torch.load(net_ESRGAN_path)
+net_interp = OrderedDict()
+
+print('Interpolating with alpha = ', alpha)
+
+for k, v_PSNR in net_PSNR.items():
+    v_ESRGAN = net_ESRGAN[k]
+    net_interp[k] = (1 - alpha) * v_PSNR + alpha * v_ESRGAN
+
+torch.save(net_interp, net_interp_path)

esrgan_plus/codes/scripts/rename.py

@@ -0,0 +1,25 @@
+import os
+import os.path
+import glob
+
+
+input_folder = '/home/xtwang/Projects/PIRM18/results/pirm_selfval_img06/*'  # glob matching pattern
+save_folder = '/home/xtwang/Projects/PIRM18/results/pirm_selfval_img'
+
+mode = 'cp'  # 'cp' | 'mv'
+file_list = sorted(glob.glob(input_folder))
+
+if not os.path.exists(save_folder):
+    os.makedirs(save_folder)
+    print('mkdir ... ' + save_folder)
+else:
+    print('File [{}] already exists. Exit.'.format(save_folder))
+
+for i, path in enumerate(file_list):
+    base_name = os.path.splitext(os.path.basename(path))[0]
+
+    new_name = base_name.split('_')[0]
+    new_path = os.path.join(save_folder, new_name + '.png')
+
+    os.system(mode + ' ' + path + ' ' + new_path)
+    print(i, base_name)