xinntao · iankonradjohnson · May 3, 2023 · Jul 6, 2023 · Jul 6, 2023 · Jul 29, 2023
diff --git a/.gitignore b/.gitignore
@@ -138,3 +138,9 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Images
+*.png
+*.jpg
+
+.idea/
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -0,0 +1,51 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": true
+        },
+        {
+            "name": "RealESRGAN training",
+            "type": "python",
+            "module": "torch.distributed.launch",
+            "request": "launch",
+            "args": [
+                "--nproc_per_node=4",
+                "--master_port=4321",
+                "realesrgan/train.py",
+                "-opt",
+                "options/train_realesrnet_x4plus.yml",
+                "--launcher",
+                "pytorch",
+                "--debug"
+            ],
+            "cwd": "${workspaceFolder}",
+            "env": {},
+            "console": "integratedTerminal",
+            "justMyCode": false
+        },
+        {
+            "name": "RealESRGAN training",
+            "type": "python",
+            "module": "realesrgan/train.py",
+            "request": "launch",
+            "args": [
+                "-opt",
+                "options/train_realesrnet_x4plus.yml",
+                "--auto_resume"
+            ],
+            "cwd": "${workspaceFolder}",
+            "env": {},
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+    ]
+}
diff --git a/README_windows.md b/README_windows.md
@@ -0,0 +1,42 @@
+# RealESRGAN
+
+We have provided the following models:
+
+1. realesr-animevideov3 (default)
+2. realesrgan-x4plus
+3. realesrgan-x4plus-anime
+
+Command:
+
+1. ./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png
+2. ./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesr-animevideov3
+3. ./realesrgan-ncnn-vulkan.exe -i input_folder -o outputfolder -n realesr-animevideov3 -s 2 -f jpg
+4. ./realesrgan-ncnn-vulkan.exe -i input_folder -o outputfolder -n realesr-animevideov3 -s 4 -f jpg
+
+
+Commands for enhancing anime videos:
+
+1. Use ffmpeg to extract frames from a video (Remember to create the folder `tmp_frames` ahead)
+
+    ffmpeg -i onepiece_demo.mp4 -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 tmp_frames/frame%08d.jpg
+
+2. Inference with Real-ESRGAN executable file (Remember to create the folder `out_frames` ahead)
+
+    ./realesrgan-ncnn-vulkan.exe -i tmp_frames -o out_frames -n realesr-animevideov3 -s 2 -f jpg
+
+3. Merge the enhanced frames back into a video
+
+    ffmpeg -i out_frames/frame%08d.jpg -i onepiece_demo.mp4 -map 0:v:0 -map 1:a:0 -c:a copy -c:v libx264 -r 23.98 -pix_fmt yuv420p output_w_audio.mp4
+
+------------------------
+
+GitHub: https://github.com/xinntao/Real-ESRGAN/
+Paper: https://arxiv.org/abs/2107.10833
+
+------------------------
+
+This executable file is **portable** and includes all the binaries and models required. No CUDA or PyTorch environment is needed.
+
+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).
diff --git a/assets/realesrgan_logo.png b/assets/realesrgan_logo.png
diff --git a/assets/realesrgan_logo_ai.png b/assets/realesrgan_logo_ai.png
diff --git a/assets/realesrgan_logo_av.png b/assets/realesrgan_logo_av.png
diff --git a/assets/realesrgan_logo_gi.png b/assets/realesrgan_logo_gi.png
diff --git a/assets/realesrgan_logo_gv.png b/assets/realesrgan_logo_gv.png
diff --git a/assets/teaser-text.png b/assets/teaser-text.png
diff --git a/assets/teaser.jpg b/assets/teaser.jpg
diff --git a/experiments/pretrained_models/README.md b/experiments/pretrained_models/README.md
diff --git a/gfpgan/weights/tmpaj5ry5cy b/gfpgan/weights/tmpaj5ry5cy
diff --git a/inference_realesrgan.py b/inference_realesrgan.py
@@ -1,4 +1,6 @@
 import argparse
+import traceback
+
 import cv2
 import glob
 import os
@@ -7,6 +9,8 @@
 
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+from utils.ImagePreprocessor import ImagePreprocessor
+import torch.autograd.profiler as profiler
 
 
 def main():
@@ -51,6 +55,13 @@ def main():
         help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
     parser.add_argument(
         '-g', '--gpu-id', type=int, default=None, help='gpu device to use (default=None) can be 0,1,2 for multi-gpu')
+    parser.add_argument(
+        '-b', '--brightness', type=float, default=None, help='Adjust the image brightness before the enhancement')
+    parser.add_argument(
+        '-c', '--contrast', type=float, default=None, help='Adjust the image contrast before the enhancement')
+    parser.add_argument(
+        '-ps', '--prescale', type=float, default=None, help='Adjust the image scale before the enhancement')
+
 
     args = parser.parse_args()
 
@@ -60,7 +71,7 @@ def main():
         model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
         netscale = 4
         file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
-    elif args.model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
+    elif args.model_name in ['RealESRNet_x4plus', 'fine-tune-400000', 'model1', 'test', 'net_g_1000000', 'brghtness-contrast-1']:  # x4 RRDBNet model
         model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
         netscale = 4
         file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
@@ -103,6 +114,12 @@ def main():
         model_path = [model_path, wdn_model_path]
         dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
 
+    # preprocessor
+    preprocessor = ImagePreprocessor(
+        brightness=args.brightness,
+        contrast=args.contrast,
+        scale=args.prescale)
+
     # restorer
     upsampler = RealESRGANer(
         scale=netscale,
@@ -131,35 +148,41 @@ def main():
         paths = sorted(glob.glob(os.path.join(args.input, '*')))
 
     for idx, path in enumerate(paths):
-        imgname, extension = os.path.splitext(os.path.basename(path))
-        print('Testing', idx, imgname)
-
-        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
-        if len(img.shape) == 3 and img.shape[2] == 4:
-            img_mode = 'RGBA'
-        else:
-            img_mode = None
-
-        try:
-            if args.face_enhance:
-                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
-            else:
-                output, _ = upsampler.enhance(img, outscale=args.outscale)
-        except RuntimeError as error:
-            print('Error', error)
-            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
-        else:
-            if args.ext == 'auto':
-                extension = extension[1:]
+        with profiler.profile(use_cuda=True, profile_memory=True) as prof:
+            imgname, extension = os.path.splitext(os.path.basename(path))
+            print('Testing', idx, imgname)
+
+            img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+            if len(img.shape) == 3 and img.shape[2] == 4:
+                img_mode = 'RGBA'
             else:
-                extension = args.ext
-            if img_mode == 'RGBA':  # RGBA images should be saved in png format
-                extension = 'png'
-            if args.suffix == '':
-                save_path = os.path.join(args.output, f'{imgname}.{extension}')
+                img_mode = None
+
+            try:
+                img = preprocessor.process(img)
+                if args.face_enhance:
+                    _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+                else:
+                    output, _ = upsampler.enhance(img, outscale=args.outscale)
+            except RuntimeError as error:
+                print('Error', error)
+                traceback.print_exc()
+                print(prof)
+                print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
             else:
-                save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
-            cv2.imwrite(save_path, output)
+                if args.ext == 'auto':
+                    extension = extension[1:]
+                else:
+                    extension = args.ext
+                if img_mode == 'RGBA':  # RGBA images should be saved in png format
+                    extension = 'png'
+                if args.suffix == '':
+                    save_path = os.path.join(args.output, f'{imgname}.{extension}')
+                else:
+                    save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
+                cv2.imwrite(save_path, output)
+
+            print(prof)
 
 
 if __name__ == '__main__':

diff --git a/inputs/00003.png b/inputs/00003.png
diff --git a/inputs/00017_gray.png b/inputs/00017_gray.png
diff --git a/inputs/0014.jpg b/inputs/0014.jpg
diff --git a/inputs/0030.jpg b/inputs/0030.jpg
diff --git a/inputs/ADE_val_00000114.jpg b/inputs/ADE_val_00000114.jpg
diff --git a/inputs/OST_009.png b/inputs/OST_009.png
diff --git a/inputs/children-alpha.png b/inputs/children-alpha.png
diff --git a/inputs/tree_alpha_16bit.png b/inputs/tree_alpha_16bit.png
diff --git a/inputs/wolf_gray.jpg b/inputs/wolf_gray.jpg
diff --git a/models/realesr-animevideov3-x2.bin b/models/realesr-animevideov3-x2.bin
diff --git a/models/realesr-animevideov3-x2.param b/models/realesr-animevideov3-x2.param
@@ -0,0 +1,43 @@
+7767517
+41 42
+Input                    input.1                  0 1 data
+Split                    splitncnn_input0         1 2 data input.1_splitncnn_0 input.1_splitncnn_1
+Convolution              Conv_0                   1 1 input.1_splitncnn_1 54 0=64 1=3 4=1 5=1 6=1728
+PReLU                    PRelu_1                  1 1 54 56 0=64
+Convolution              Conv_2                   1 1 56 57 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_3                  1 1 57 59 0=64
+Convolution              Conv_4                   1 1 59 60 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_5                  1 1 60 62 0=64
+Convolution              Conv_6                   1 1 62 63 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_7                  1 1 63 65 0=64
+Convolution              Conv_8                   1 1 65 66 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_9                  1 1 66 68 0=64
+Convolution              Conv_10                  1 1 68 69 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_11                 1 1 69 71 0=64
+Convolution              Conv_12                  1 1 71 72 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_13                 1 1 72 74 0=64
+Convolution              Conv_14                  1 1 74 75 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_15                 1 1 75 77 0=64
+Convolution              Conv_16                  1 1 77 78 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_17                 1 1 78 80 0=64
+Convolution              Conv_18                  1 1 80 81 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_19                 1 1 81 83 0=64
+Convolution              Conv_20                  1 1 83 84 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_21                 1 1 84 86 0=64
+Convolution              Conv_22                  1 1 86 87 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_23                 1 1 87 89 0=64
+Convolution              Conv_24                  1 1 89 90 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_25                 1 1 90 92 0=64
+Convolution              Conv_26                  1 1 92 93 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_27                 1 1 93 95 0=64
+Convolution              Conv_28                  1 1 95 96 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_29                 1 1 96 98 0=64
+Convolution              Conv_30                  1 1 98 99 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_31                 1 1 99 101 0=64
+Convolution              Conv_32                  1 1 101 102 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_33                 1 1 102 104 0=64
+Convolution              Conv_34                  1 1 104 105 0=48 1=3 4=1 5=1 6=27648
+PixelShuffle             DepthToSpace_35          1 1 105 106 0=4
+Interp                   Resize_37                1 1 input.1_splitncnn_0 111 0=1 1=4.000000e+00 2=4.000000e+00
+BinaryOp                 Add_38                   2 1 106 111 112
+Interp                   Resize_40                1 1 112 output 0=3 1=5.000000e-01 2=5.000000e-01
diff --git a/models/realesr-animevideov3-x3.bin b/models/realesr-animevideov3-x3.bin
diff --git a/models/realesr-animevideov3-x3.param b/models/realesr-animevideov3-x3.param
@@ -0,0 +1,43 @@
+7767517
+41 42
+Input                    input.1                  0 1 data
+Split                    splitncnn_input0         1 2 data input.1_splitncnn_0 input.1_splitncnn_1
+Convolution              Conv_0                   1 1 input.1_splitncnn_1 54 0=64 1=3 4=1 5=1 6=1728
+PReLU                    PRelu_1                  1 1 54 56 0=64
+Convolution              Conv_2                   1 1 56 57 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_3                  1 1 57 59 0=64
+Convolution              Conv_4                   1 1 59 60 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_5                  1 1 60 62 0=64
+Convolution              Conv_6                   1 1 62 63 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_7                  1 1 63 65 0=64
+Convolution              Conv_8                   1 1 65 66 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_9                  1 1 66 68 0=64
+Convolution              Conv_10                  1 1 68 69 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_11                 1 1 69 71 0=64
+Convolution              Conv_12                  1 1 71 72 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_13                 1 1 72 74 0=64
+Convolution              Conv_14                  1 1 74 75 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_15                 1 1 75 77 0=64
+Convolution              Conv_16                  1 1 77 78 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_17                 1 1 78 80 0=64
+Convolution              Conv_18                  1 1 80 81 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_19                 1 1 81 83 0=64
+Convolution              Conv_20                  1 1 83 84 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_21                 1 1 84 86 0=64
+Convolution              Conv_22                  1 1 86 87 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_23                 1 1 87 89 0=64
+Convolution              Conv_24                  1 1 89 90 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_25                 1 1 90 92 0=64
+Convolution              Conv_26                  1 1 92 93 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_27                 1 1 93 95 0=64
+Convolution              Conv_28                  1 1 95 96 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_29                 1 1 96 98 0=64
+Convolution              Conv_30                  1 1 98 99 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_31                 1 1 99 101 0=64
+Convolution              Conv_32                  1 1 101 102 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_33                 1 1 102 104 0=64
+Convolution              Conv_34                  1 1 104 105 0=48 1=3 4=1 5=1 6=27648
+PixelShuffle             DepthToSpace_35          1 1 105 106 0=4
+Interp                   Resize_37                1 1 input.1_splitncnn_0 111 0=1 1=4.000000e+00 2=4.000000e+00
+BinaryOp                 Add_38                   2 1 106 111 112
+Interp                   Resize_40                1 1 112 output 0=3 1=7.500000e-01 2=7.500000e-01
diff --git a/models/realesr-animevideov3-x4.bin b/models/realesr-animevideov3-x4.bin
diff --git a/models/realesr-animevideov3-x4.param b/models/realesr-animevideov3-x4.param
@@ -0,0 +1,42 @@
+7767517
+40 41
+Input                    input.1                  0 1 data
+Split                    splitncnn_input0         1 2 data input.1_splitncnn_0 input.1_splitncnn_1
+Convolution              Conv_0                   1 1 input.1_splitncnn_1 54 0=64 1=3 4=1 5=1 6=1728
+PReLU                    PRelu_1                  1 1 54 56 0=64
+Convolution              Conv_2                   1 1 56 57 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_3                  1 1 57 59 0=64
+Convolution              Conv_4                   1 1 59 60 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_5                  1 1 60 62 0=64
+Convolution              Conv_6                   1 1 62 63 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_7                  1 1 63 65 0=64
+Convolution              Conv_8                   1 1 65 66 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_9                  1 1 66 68 0=64
+Convolution              Conv_10                  1 1 68 69 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_11                 1 1 69 71 0=64
+Convolution              Conv_12                  1 1 71 72 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_13                 1 1 72 74 0=64
+Convolution              Conv_14                  1 1 74 75 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_15                 1 1 75 77 0=64
+Convolution              Conv_16                  1 1 77 78 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_17                 1 1 78 80 0=64
+Convolution              Conv_18                  1 1 80 81 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_19                 1 1 81 83 0=64
+Convolution              Conv_20                  1 1 83 84 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_21                 1 1 84 86 0=64
+Convolution              Conv_22                  1 1 86 87 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_23                 1 1 87 89 0=64
+Convolution              Conv_24                  1 1 89 90 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_25                 1 1 90 92 0=64
+Convolution              Conv_26                  1 1 92 93 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_27                 1 1 93 95 0=64
+Convolution              Conv_28                  1 1 95 96 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_29                 1 1 96 98 0=64
+Convolution              Conv_30                  1 1 98 99 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_31                 1 1 99 101 0=64
+Convolution              Conv_32                  1 1 101 102 0=64 1=3 4=1 5=1 6=36864
+PReLU                    PRelu_33                 1 1 102 104 0=64
+Convolution              Conv_34                  1 1 104 105 0=48 1=3 4=1 5=1 6=27648
+PixelShuffle             DepthToSpace_35          1 1 105 106 0=4
+Interp                   Resize_37                1 1 input.1_splitncnn_0 111 0=1 1=4.000000e+00 2=4.000000e+00
+BinaryOp                 Add_38                   2 1 106 111 output
diff --git a/models/realesrgan-x4plus-anime.bin b/models/realesrgan-x4plus-anime.bin