add app_enhance

.gitignore

@@ -0,0 +1,3 @@
+.vscode
+.DS_Store
+__pycache__

app.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+from app_enhance import create_demo as create_demo_enhance
+
+with gr.Blocks(css="style.css") as demo:
+    with gr.Tabs():
+        with gr.Tab(label="Enhance"):
+            create_demo_enhance()
+
+demo.launch()

app_enhance.py

@@ -0,0 +1,119 @@
+import os
+import subprocess
+import spaces
+import torch
+import cv2
+import uuid
+import gradio as gr
+import numpy as np
+
+from PIL import Image
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan.utils import GFPGANer
+from realesrgan.utils import RealESRGANer
+
+def runcmd(cmd, verbose = False):
+
+    process = subprocess.Popen(
+        cmd,
+        stdout = subprocess.PIPE,
+        stderr = subprocess.PIPE,
+        text = True,
+        shell = True
+    )
+    std_out, std_err = process.communicate()
+    if verbose:
+        print(std_out.strip(), std_err)
+    pass
+
+
+if not os.path.exists('GFPGANv1.4.pth'):
+    runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+
+@spaces.GPU(duration=5)
+def enhance_image(
+    input_image: Image,
+    scale: int,
+    enhance_mode: str,
+):
+    only_face = enhance_mode == "Only Face Enhance"
+    if enhance_mode == "Only Face Enhance":
+        face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2)
+    elif enhance_mode == "Only Image Enhance":
+        face_enhancer = None
+    else:
+        face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+    
+    img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    
+    if face_enhancer is not None:
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=only_face, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=scale)
+
+    if scale != 2:
+        interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+        h, w = img.shape[0:2]
+        output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
+
+    enhanced_image = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+    tmpPrefix = "/tmp/gradio/"
+
+    extension = 'png'
+
+    targetDir = f"{tmpPrefix}output/"
+    if not os.path.exists(targetDir):
+        os.makedirs(targetDir)
+
+    enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}"
+    enhanced_image.save(enhanced_path, quality=100)
+        
+    return enhanced_image, enhanced_path
+
+
+def create_demo() -> gr.Blocks:
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                scale = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Scale")
+            with gr.Column():
+                enhance_mode = gr.Dropdown(
+                    label="Enhance Mode",
+                    choices=[
+                        "Only Face Enhance",
+                        "Only Image Enhance",
+                        "Face Enhance + Image Enhance",
+                    ],
+                    value="Face Enhance + Image Enhance",
+                )
+                g_btn = gr.Button("Enhance Image")
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+            with gr.Column():
+                output_image = gr.Image(label="Enhanced Image", type="pil", interactive=False)
+                enhance_image_path = gr.File(label="Download the Enhanced Image", interactive=False)
+                
+        
+        g_btn.click(
+            fn=enhance_image,
+            inputs=[input_image, scale, enhance_mode],
+            outputs=[output_image, enhance_image_path],
+        )
+
+    return demo

enhance_utils.py

@@ -0,0 +1,54 @@
+import os
+import torch
+import cv2
+import numpy as np
+import subprocess
+
+from PIL import Image
+from gfpgan.utils import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+def runcmd(cmd, verbose = False, *args, **kwargs):
+
+    process = subprocess.Popen(
+        cmd,
+        stdout = subprocess.PIPE,
+        stderr = subprocess.PIPE,
+        text = True,
+        shell = True
+    )
+    std_out, std_err = process.communicate()
+    if verbose:
+        print(std_out.strip(), std_err)
+    pass
+
+runcmd("pip freeze")
+if not os.path.exists('GFPGANv1.4.pth'):
+    runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2)
+
+def enhance_image(
+    pil_image: Image,
+    enhance_face: bool = True,
+):
+    img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    if enhance_face:
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=2)
+    pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+
+    return pil_output

requirements.txt

@@ -0,0 +1,7 @@
+torch
+gradio
+spaces
+gfpgan
+git+https://github.com/XPixelGroup/BasicSR@master
+facexlib
+realesrgan