update app

app.py

@@ -1,56 +1,56 @@
-import zipfile
-from pathlib import Path
+from glob import glob
 
 import gradio as gr
 from gradio_client import Client
-from huggingface_hub import hf_hub_download
 
-clientGFPGAN = Client("Xintao/GFPGAN")
-clientSuperface = Client("leonelhs/superface")
-clientZeroScratches = Client("leonelhs/ZeroScratches")
-clientDeoldify = Client("leonelhs/deoldify")
-clientEnhanceLight = Client("leonelhs/Zero-DCE")
-clientZeroBackground = Client("leonelhs/remove-background")
-clientFaceParser = Client("leonelhs/faceparser")
+from utils import make_flatten_background
 
 REPO_ID = "leonelhs/faceshine"
 
+clients = {
+    "GFPGAN": "leonelhs/GFPGAN",
+    "ZeroScratches": "leonelhs/ZeroScratches",
+    "Deoldify": "leonelhs/deoldify",
+    "EnhanceLight": "leonelhs/Zero-DCE",
+    "ZeroBackground": "leonelhs/rembg",
+}
 
-def download(file):
-    file_zip = hf_hub_download(repo_id=REPO_ID, filename=file + ".zip")
-    with zipfile.ZipFile(file_zip, "r") as zip_ref:
-        zip_ref.extractall(file)
-    return list(Path(file).iterdir())
 
-
-def gfpgan_face(image, version, scale):
-    return clientGFPGAN.predict(image, version, scale, fn_index=0)[0]
+def load_client(space):
+    try:
+        return Client(space)
+    except ValueError as err:
+        print(err)
+        logger.value.append(f"Space: {space}, log: {err}")
+        pass
 
 
-def enhance_face(image, upsampler, face_enhancer, scale):
-    return clientSuperface.predict(image, upsampler, face_enhancer, scale, api_name="/predict")
+def gfpgan_face(image, version, scale):
+    return clients["GFPGAN"].predict(image, version, scale, fn_index=0)[0]
 
 
 def zero_scratches(image):
-    return clientZeroScratches.predict(image, api_name="/predict")
+    return clients["ZeroScratches"].predict(image, api_name="/predict")
 
 
 def colorize_photo(image):
-    return clientDeoldify.predict(image, api_name="/predict")
+    return clients["Deoldify"].predict(image, api_name="/predict")
 
 
 def enhance_light(image):
-    return clientEnhanceLight.predict(image, api_name="/predict")
+    return clients["EnhanceLight"].predict(image, api_name="/predict")
 
 
 def zero_background(image, new_bgr=None):
-    # return clientZeroBackground.predict(image, new_bgr, api_name="/predict")[0]
     # Fixme: cant find predict function by name
-    return clientZeroBackground.predict(image, new_bgr, fn_index=0)[1]
+    # return clients["ZeroBackground"].predict(image, new_bgr, fn_index=0)[1]
+    # return clients["ZeroBackground"].predict(image, fn_index=0)
+    img, mask = clients["ZeroBackground"].predict(image, "U2NET Human Seg", False, fn_index=9)
+    return make_flatten_background(img, mask)
 
 
 def parse_face(image):
-    return clientFaceParser.predict(image, api_name="/predict")
+    return clients["FaceParser"].predict(image, api_name="/predict")
 
 
 def mirror(x):
@@ -61,6 +61,10 @@ def active_first():
     return gr.Tabs.update(selected=0)
 
 
+def clear():
+    return None, None
+
+
 footer = r"""
 <center>
 <p>This App is running on a CPU, help us to upgrade a GPU or just give us a <a href='https://github.com/leonelhs/face-shine' target='_blank'>Github ⭐</a></p>
@@ -74,68 +78,60 @@ footer = r"""
 """
 
 with gr.Blocks(title="Face Shine") as app:
+    logger = gr.State(value=[])
+
+    for client, endpoint in clients.items():
+        clients[client] = load_client(endpoint)
+
     with gr.Row():
         gr.HTML("<center><h1>Face Shine</h1></center>")
 
     with gr.Tabs() as tabs:
         with gr.TabItem("Photo restorer", id=0):
-            with gr.Row():
+            with gr.Row(equal_height=False):
                 with gr.Column(scale=1):
-                    btn_hires = gr.Button(value="Enhance faces")
                     btn_eraser = gr.Button(value="Erase scratches")
                     btn_color = gr.Button(value="Colorize photo")
+                    btn_hires = gr.Button(value="Enhance face")
                     btn_light = gr.Button(value="Enhance light")
-                    with gr.Accordion("New background", open=False):
-                        img_newbgr = gr.Image(label="Left empty for a transparent background", type="filepath")
-                        btn_newbgr = gr.Button(value="Clear background")
+                    btn_clear = gr.Button(value="Flatten background")
 
-                with gr.Column(scale=4):
+                with gr.Column(scale=2):
                     with gr.Row():
                         img_input = gr.Image(label="Input", type="filepath")
-                        img_output = gr.Image(label="Result", type="filepath", interactive=False)
                     with gr.Row():
-                        gr.ClearButton([img_input, img_output, img_newbgr], variant="stop")
-                        btn_swap = gr.Button(value="Swap images", variant="primary")
+                        btn_reset = gr.Button(value="Reset", variant="stop")
+                        btn_swap = gr.Button(value="Ok", variant="primary")
+
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        img_output = gr.Image(label="Result", type="filepath", interactive=False)
 
         with gr.TabItem("Examples", id=1):
-            gr.Examples(examples=download("lowres"), inputs=[img_input], label="Low resolution")
-            gr.Examples(examples=download("gray"), inputs=[img_input], label="Gray scale")
-            gr.Examples(examples=download("scratch"), inputs=[img_input], label="Scratched")
-            gr.Examples(examples=download("backs"), inputs=[img_newbgr], label="Backgrounds")
+            gr.Examples(examples=glob("lowres/*"), inputs=[img_input], label="Low resolution")
+            gr.Examples(examples=glob("gray/*"), inputs=[img_input], label="Gray scale")
+            gr.Examples(examples=glob("scratch/*"), inputs=[img_input], label="Scratched")
             gr.Button(value="Ok", variant="primary").click(active_first, None, tabs)
 
         with gr.TabItem("Settings", id=2):
             with gr.Accordion("Image restoration settings", open=False):
-                enhancer = gr.Dropdown([
-                    'v1.2',
-                    'v1.3',
-                    'v1.4',
-                    'RestoreFormer'],
-                    type="value", value='RestoreFormer', label='GFPGAN face restoration algorithm',
-                    info="version")
-                restorer = gr.Dropdown([
-                    'RealESRGAN_x2plus',
-                    'RealESRGAN_x4plus',
-                    'RealESRNet_x4plus',
-                    'AI-Forever_x2plus',
-                    'AI-Forever_x4plus',
-                    'RealESRGAN_x4plus_anime_6B',
-                    'realesr-animevideov3',
-                    'realesr-general-x4v3'],
-                    type="value", value='RealESRGAN_x4plus',
-                    label='Real-ESRGAN image restoration algorithm', info="version", visible=False)
-                rescale = gr.Dropdown(["1", "2", "3", "4"], type="value", value="2", label="Rescaling factor")
+                enhancer = gr.Dropdown(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer'],
+                                       label='GFPGAN face restoration algorithm',
+                                       type="value", value='RestoreFormer',
+                                       info="version")
+                rescale = gr.Dropdown(["1", "2", "3", "4"],
+                                      type="value", value="2", label="Rescaling factor")
             with gr.Accordion("Logs info", open=False):
                 text_logger = gr.Textbox(label="login", lines=5, show_label=False)
             gr.Button("Save settings")
 
-    # btn_hires.click(enhance_face, inputs=[img_input, restorer, enhancer, rescale], outputs=[img_output, text_logger])
     btn_hires.click(gfpgan_face, inputs=[img_input, enhancer, rescale], outputs=[img_output])
     btn_eraser.click(zero_scratches, inputs=[img_input], outputs=[img_output])
     btn_color.click(colorize_photo, inputs=[img_input], outputs=[img_output])
     btn_light.click(enhance_light, inputs=[img_input], outputs=[img_output])
-    btn_newbgr.click(zero_background, inputs=[img_input, img_newbgr], outputs=[img_output])
+    btn_clear.click(zero_background, inputs=[img_input], outputs=[img_output])
     btn_swap.click(mirror, inputs=[img_output], outputs=[img_input])
+    btn_reset.click(clear, outputs=[img_input, img_output])
 
     with gr.Row():
         gr.HTML(footer)