Update app.py

app.py

@@ -1,149 +1,129 @@
-import gradio as gr
-import argparse
-from realesrgan import RealESRGANer
-from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 import os
+import cv2
+import torch
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
-def Generate(img, model_name):
-
-    global output
-    parser = argparse.ArgumentParser()
-    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
-    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
-    parser.add_argument(
-        '-dn',
-        '--denoise_strength',
-        type=float,
-        default=0.5,
-        help=('Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. '
-              'Only used for the realesr-general-x4v3 model'))
-    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')
-    parser.add_argument(
-        '--model_path', type=str, default=None, help='[Option] Model path. Usually, you do not need to specify it')
-    parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
-    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
-    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
-    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
-    parser.add_argument('--face_enhance', action='store_true',help='Use GFPGAN to enhance face')
-    parser.add_argument(
-        '--fp32', action='store_true',default=True,help='Use fp32 precision during inference. Default: fp16 (half precision).')
-    parser.add_argument(
-        '--alpha_upsampler',
-        type=str,
-        default='realesrgan',
-        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
-    parser.add_argument(
-        '--ext',
-        type=str,
-        default='auto',
-        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')
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+import gradio as gr
 
-    args = parser.parse_args()
+# 模型配置
+MODEL_OPTIONS = {
+    "RealESRGAN_x4plus": {
+        "model": lambda: 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"
+    },
+    "RealESRNet_x4plus": {
+        "model": lambda: 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"
+    },
+    "RealESRGAN_x4plus_anime_6B": {
+        "model": lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4),
+        "netscale": 4,
+        "file_url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+    },
+    "RealESRGAN_x2plus": {
+        "model": lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2),
+        "netscale": 2,
+        "file_url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
+    },
+    "realesr-animevideov3": {
+        "model": lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'),
+        "netscale": 4,
+        "file_url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth"
+    },
+    "realesr-general-x4v3": {
+        "model": lambda: SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu'),
+        "netscale": 4,
+        "file_url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
+    }
+}
 
-    if model_name == 'RealESRGAN_x4plus':  # 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.0/RealESRGAN_x4plus.pth']
-    elif model_name == 'RealESRNet_x4plus':  # 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']
-    elif model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
-        netscale = 4
-        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
-    elif model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
-        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-        netscale = 2
-        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
-    elif model_name == 'realesr-animevideov3':  # x4 VGG-style model (XS size)
-        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
-        netscale = 4
-        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth']
-    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
-        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-        netscale = 4
-        file_url = [
-            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
-            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
-        ]
+def load_model(model_name):
+    """加载并初始化 Real-ESRGAN 模型"""
+    model_config = MODEL_OPTIONS[model_name]
+    model = model_config["model"]()
+    netscale = model_config["netscale"]
+    file_url = model_config["file_url"]
 
-    model_path = os.path.join('weights', model_name + '.pth')
-    print(model_path)
+    # 下载模型权重
+    model_path = os.path.join("weights", f"{model_name}.pth")
     if not os.path.isfile(model_path):
-        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-        for url in file_url:
-            # model_path will be updated
-            model_path = load_file_from_url(
-                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
-    dni_weight = None
-    if model_name == 'realesr-general-x4v3' and args.denoise_strength != 1:
-        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
-        model_path = [model_path, wdn_model_path]
-        dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
+        os.makedirs("weights", exist_ok=True)
+        model_path = load_file_from_url(url=file_url, model_dir="weights", progress=True, file_name=None)
 
-    # restorer
+    # 初始化 RealESRGANer
     upsampler = RealESRGANer(
         scale=netscale,
         model_path=model_path,
-        dni_weight=dni_weight,
         model=model,
-        tile=args.tile,
-        tile_pad=args.tile_pad,
-        pre_pad=args.pre_pad,
-        half=not args.fp32,
-        gpu_id=args.gpu_id)
+        tile=0,  # 默认无分块
+        tile_pad=10,
+        pre_pad=0,
+        half=True  # 默认使用 fp16
+    )
+    return upsampler
+
+def enhance_image(input_image, model_name, outscale, face_enhance):
+    """执行图像超分辨率增强"""
+    # 将 Gradio 上传的图像转换为 OpenCV 格式
+    img = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
 
-    if args.face_enhance:  # Use GFPGAN for face enhancement
+    # 确定图像模式
+    if len(img.shape) == 3 and img.shape[2] == 4:
+        img_mode = 'RGBA'
+    else:
+        img_mode = None
+
+    # 加载模型
+    upsampler = load_model(model_name)
+
+    # 是否使用人脸增强
+    if face_enhance:
         from gfpgan import GFPGANer
         face_enhancer = GFPGANer(
             model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
-            upscale=args.outscale,
+            upscale=outscale,
             arch='clean',
             channel_multiplier=2,
-            bg_upsampler=upsampler)
-    os.makedirs(args.output, exist_ok=True)
-
-    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)
-        print("生成成功")
-
-    except RuntimeError as error:
-        print('Error', error)
-        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
-        output = None
+            bg_upsampler=upsampler
+        )
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=outscale)
 
+    # 将结果转换回 RGB 格式以供 Gradio 显示
+    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
     return output
 
-with gr.Blocks() as demo:
-
-    gr.Markdown(
-    """
-    # <center> Real-ESRGAN 在线体验程序 
-    """)
-    gr.Markdown("""
-     1. **项目模型运行在CPU上，等待时间略长**
-     2. **原工程项目旨在对图片就行修复**
-     3. **项目源地址为：[Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN)**
-    """)
+# Gradio 界面
+with gr.Blocks(title="Real-ESRGAN 图像超分辨率") as app:
+    gr.Markdown("## Real-ESRGAN 图像超分辨率系统")
+    gr.Markdown("上传图像，选择模型和参数，生成高清图像！")
 
     with gr.Row():
         with gr.Column():
-            img = gr.Image(type="numpy",label = "输入图片")
-            model_name = gr.Dropdown(["RealESRGAN_x4plus","RealESRGAN_x4plus_anime_6B","RealESRGAN_x2plus",
-                                      "realesr-animevideov3","realesr-general-x4v3"],info="选择模型")
-        with gr.Column():
-            img_out = gr.Image(type="numpy",label = "输出图片")
+            input_image = gr.Image(label="输入图像", type="numpy")
+            model_dropdown = gr.Dropdown(
+                choices=list(MODEL_OPTIONS.keys()),
+                label="选择模型",
+                value="RealESRGAN_x4plus"
+            )
+            outscale = gr.Slider(minimum=1, maximum=4, step=0.5, value=4, label="放大倍数")
+            face_enhance = gr.Checkbox(label="启用人脸增强", value=False)
+            enhance_btn = gr.Button("开始增强", variant="primary")
 
-    btn = gr.Button("Generate")
+        with gr.Column():
+            output_image = gr.Image(label="增强结果", type="numpy")
 
-    btn.click(Generate, inputs=[img,model_name], outputs=[img_out])
+    enhance_btn.click(
+        fn=enhance_image,
+        inputs=[input_image, model_dropdown, outscale, face_enhance],
+        outputs=[output_image],
+        api_name="enhance"
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch(server_name="0.0.0.0", server_port=7860)