add RealESRNet model, fix bug in exe file

README.md

@@ -49,7 +49,7 @@ If you have some images that Real-ESRGAN could not well restored, please also op
 
 ### Portable executable files
 
-You can download **Windows executable files** from https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN-ncnn-vulkan.zip
+You can download **Windows executable files** from https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRGAN-ncnn-vulkan-20210725-windows.zip
 
 This executable file is **portable** and includes all the binaries and models required. No CUDA or PyTorch environment is needed.<br>
 
@@ -59,6 +59,14 @@ You can simply run the following command:
 ./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png
 ```
 
+We have provided three models:
+
+1. realesrgan-x4plus  (default)
+2. realesrnet-x4plus
+3. esrgan-x4
+
+You can use the `-n` argument for other models, for example, `./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesrnet-x4plus`
+
 Note that it may introduce block inconsistency (and also generate slightly different results from the PyTorch implementation), because this executable file first crops the input image into several tiles, and then processes them separately, finally stitches together.
 
 This executable file is based on the wonderful [Tencent/ncnn](https://github.com/Tencent/ncnn) and [realsr-ncnn-vulkan](https://github.com/nihui/realsr-ncnn-vulkan) by [nihui](https://github.com/nihui).
@@ -106,6 +114,12 @@ python inference_realesrgan.py --model_path experiments/pretrained_models/RealES
 
 Results are in the `results` folder
 
+## :european_castle: Model Zoo
+
+- [RealESRGAN-x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth)
+- [RealESRNet-x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth)
+- [official ESRGAN-x4](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth)
+
 ## :computer: Training
 
 A detailed guide can be found in [Training.md](Training.md).

Training.md

@@ -34,7 +34,10 @@ DF2K_HR_sub/000001_s003.png
 
 ## Train Real-ESRNet
 
-1. Download pre-trained model [ESRGAN](https://drive.google.com/file/d/1b3_bWZTjNO3iL2js1yWkJfjZykcQgvzT/view?usp=sharing) into `experiments/pretrained_models`.
+1. Download pre-trained model [ESRGAN](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth) into `experiments/pretrained_models`.
+    ```bash
+    wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth -P experiments/pretrained_models
+    ```
 1. Modify the content in the option file `options/train_realesrnet_x4plus.yml` accordingly:
     ```yml
     train:

inference_realesrgan.py

@@ -12,6 +12,7 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/RealESRGAN_x4plus.pth')
     parser.add_argument('--scale', type=int, default=4)
+    parser.add_argument('--suffix', type=str, default='_out')
     parser.add_argument('--input', type=str, default='inputs', help='input image or folder')
     args = parser.parse_args()
 
@@ -19,7 +20,11 @@ def main():
     # set up model
     model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=args.scale)
     loadnet = torch.load(args.model_path)
-    model.load_state_dict(loadnet['params_ema'], strict=True)
+    if 'params_ema' in loadnet:
+        keyname = 'params_ema'
+    else:
+        keyname = 'params'
+    model.load_state_dict(loadnet[keyname], strict=True)
     model.eval()
     model = model.to(device)
 
@@ -59,7 +64,7 @@ def main():
             output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
             output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
             output = (output * 255.0).round().astype(np.uint8)
-            cv2.imwrite(f'results/{imgname}_RealESRGAN.png', output)
+            cv2.imwrite(f'results/{imgname}_{args.suffix}.png', output)
         except Exception as error:
             print('Error', error)