Browse Source

adapt Real-ESRGAN-anime model

Xintao 3 years ago
parent
commit
6ff747174d
3 changed files with 16 additions and 4 deletions
  1. 4 3
      README.md
  2. 11 0
      docs/ncnn_conversion.md
  3. 1 1
      scripts/pytorch2onnx.py

+ 4 - 3
README.md

@@ -14,11 +14,12 @@
 Real-ESRGAN aims at developing **Practical Algorithms for General Image Restoration**.<br>
 We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data.
 
-:art: Real-ESRGAN needs your contribution. Any contributions are welcome, such as new features/models/typo fixes/suggestions/maintenance, *etc*. See [CONTRIBUTING.md](CONTRIBUTING.md). All contributors are list [here](CONTRIBUTING.md#Contributors).
+:art: Real-ESRGAN needs your contributions. Any contributions are welcome, such as new features/models/typo fixes/suggestions/maintenance, *etc*. See [CONTRIBUTING.md](CONTRIBUTING.md). All contributors are list [here](CONTRIBUTING.md#Contributors).
 
 :question: Frequently Asked Questions can be found in [FAQ.md](FAQ.md).
 
 :triangular_flag_on_post: **Updates**
+- :white_check_mark: Add *RealESRGAN_x4plus_anime_6B.pth*, which is optimized for **anime** images with much smaller size. More details and comparisons with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) are in [anime_model.md](docs/anime_model.md)
 - :white_check_mark: Support finetuning on your own data or paired data (*i.e.*, finetuning ESRGAN). See [here](Training.md#Finetune-Real-ESRGAN-on-your-own-dataset)
 - :white_check_mark: Integrate [GFPGAN](https://github.com/TencentARC/GFPGAN) to support **face enhancement**.
 - :white_check_mark: Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN). Thanks [@AK391](https://github.com/AK391)
@@ -62,7 +63,7 @@ Here is a TODO list in the near future:
 
 - [ ] optimize for human faces
 - [ ] optimize for texts
-- [ ] optimize for anime images [in progress]
+- [x] optimize for anime images
 - [ ] support more scales
 - [ ] support controllable restoration strength
 
@@ -87,7 +88,7 @@ We have provided three models:
 
 1. realesrgan-x4plus  (default)
 2. realesrnet-x4plus
-3. esrgan-x4
+3. realesrgan-x4plus-anime (optimized for anime images, small size)
 
 You can use the `-n` argument for other models, for example, `./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesrnet-x4plus`
 

+ 11 - 0
docs/ncnn_conversion.md

@@ -0,0 +1,11 @@
+# Instructions on converting to NCNN models
+
+1. Convert to onnx model with `scripts/pytorch2onnx.py`. Remember to modify codes accordingly
+1. Convert onnx model to ncnn model
+    1. `cd ncnn-master\ncnn\build\tools\onnx`
+    1. `onnx2ncnn.exe realesrgan-x4.onnx realesrgan-x4-raw.param realesrgan-x4-raw.bin`
+1. Optimize ncnn model
+    1. fp16 mode
+        1. `cd ncnn-master\ncnn\build\tools`
+        1. `ncnnoptimize.exe realesrgan-x4-raw.param realesrgan-x4-raw.bin realesrgan-x4.param realesrgan-x4.bin 1`
+1. Modify the blob name in `realesrgan-x4.param`: `data` and `output`

+ 1 - 1
scripts/pytorch2onnx.py

@@ -3,7 +3,7 @@ import torch.onnx
 from basicsr.archs.rrdbnet_arch import RRDBNet
 
 # An instance of your model
-model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
+model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
 model.load_state_dict(torch.load('experiments/pretrained_models/RealESRGAN_x4plus.pth')['params_ema'])
 # set the train mode to false since we will only run the forward pass.
 model.train(False)