Update app.py (#5)

- Update app.py (144ebf93ca1b27c31e62ebdb367e248b70803f41)


Co-authored-by: Arhenniuss Mallikovv <[email protected]>

app.py

@@ -1,109 +1,109 @@
-import gradio as gr
-import os
-import cv2
-from rembg import new_session, remove
-from rembg.sessions import sessions_class
-
-def inference(file, mask, model, x, y):
-    im = cv2.imread(file, cv2.IMREAD_COLOR)
-    input_path = "input.png"
-    output_path = "output.png"
-    cv2.imwrite(input_path, im)
-
-    with open(input_path, 'rb') as i:
-        with open(output_path, 'wb') as o:
-            input = i.read()
-            session = new_session(model)
-
-            output = remove(
-                input,
-                session=session,
-                **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] },
-                only_mask=(mask == "Mask only")
-            )
-            o.write(output)
-
-    return output_path
-
-title = "RemBG"
-description = "Gradio demo for **[RemBG](https://github.com/danielgatis/rembg)**. To use it, simply upload your image, select a model, click Process, and wait."
-badge = """
-    <div style="position: fixed; left: 50%; text-align: center;">
-        <a href="https://github.com/danielgatis/rembg" target="_blank" style="text-decoration: none;">
-            <img src="https://img.shields.io/badge/RemBG-Github-blue" alt="RemBG Github" />
-        </a>
-    </div>
-"""
-def get_coords(evt: gr.SelectData) -> tuple:
-    return evt.index[0], evt.index[1]
-
-def show_coords(model: str):
-    visible = model == "sam"
-    return gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible)
-
-for session in sessions_class:
-    session.download_models()
-
-with gr.Blocks() as app:
-    gr.Markdown(f"# {title}")
-    gr.Markdown(description)
-
-    with gr.Row():
-        inputs = gr.Image(type="filepath", label="Input Image")
-        outputs = gr.Image(type="filepath", label="Output Image")
-        
-    with gr.Row():
-        mask_option = gr.Radio(
-            ["Default", "Mask only"],
-            value="Default",
-            label="Output Type"
-        )
-        model_selector = gr.Dropdown(
-            [
-                "u2net",
-                "u2netp",
-                "u2net_human_seg",
-                "u2net_cloth_seg",
-                "silueta",
-                "isnet-general-use",
-                "isnet-anime",
-                "sam",
-                "birefnet-general",
-                "birefnet-general-lite",
-                "birefnet-portrait",
-                "birefnet-dis",
-                "birefnet-hrsod",
-                "birefnet-cod",
-                "birefnet-massive"
-            ],
-            value="isnet-general-use",
-            label="Model Selection"
-        )
-
-    extra = gr.Markdown("## Click on the image to capture coordinates (for SAM model)", visible=False)
-
-    x = gr.Number(label="Mouse X Coordinate", visible=False)
-    y = gr.Number(label="Mouse Y Coordinate", visible=False)
-
-    model_selector.change(show_coords, inputs=model_selector, outputs=[x, y, extra])
-    inputs.select(get_coords, None, [x, y])
-
-
-    gr.Button("Process Image").click(
-        inference,
-        inputs=[inputs, mask_option, model_selector, x, y],
-        outputs=outputs
-    )
-
-    gr.Examples(
-        examples=[
-            ["lion.png", "Default", "u2net", None, None],
-            ["girl.jpg", "Default", "u2net", None, None],
-            ["anime-girl.jpg", "Default", "isnet-anime", None, None]
-        ],
-        inputs=[inputs, mask_option, model_selector, x, y],
-        outputs=outputs
-    )
-    gr.HTML(badge)
-
-app.launch()
+import gradio as gr
+import os
+import cv2
+from rembg import new_session, remove
+from rembg.sessions import sessions_class
+
+def inference(file, mask, model, x, y):
+    im = cv2.imread(file, cv2.IMREAD_COLOR)
+    input_path = "input.png"
+    output_path = "output.png"
+    cv2.imwrite(input_path, im)
+
+    with open(input_path, 'rb') as i:
+        with open(output_path, 'wb') as o:
+            input = i.read()
+            session = new_session(model)
+
+            output = remove(
+                input,
+                session=session,
+                **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] },
+                only_mask=(mask == "Mask only")
+            )
+            o.write(output)
+
+    return output_path
+
+title = "RemBG"
+description = "Gradio demo for **[RemBG](https://github.com/danielgatis/rembg)**. To use it, simply upload your image, select a model, click Process, and wait."
+badge = """
+    <div style="position: fixed; left: 50%; text-align: center;">
+        <a href="https://github.com/danielgatis/rembg" target="_blank" style="text-decoration: none;">
+            <img src="https://img.shields.io/badge/RemBG-Github-blue" alt="RemBG Github" />
+        </a>
+    </div>
+"""
+def get_coords(evt: gr.SelectData) -> tuple:
+    return evt.index[0], evt.index[1]
+
+def show_coords(model: str):
+    visible = model == "sam"
+    return gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible)
+
+for session in sessions_class:
+    session.download_models()
+
+with gr.Blocks() as app:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+
+    with gr.Row():
+        inputs = gr.Image(type="filepath", label="Input Image")
+        outputs = gr.Image(type="filepath", label="Output Image")
+        
+    with gr.Row():
+        mask_option = gr.Radio(
+            ["Default", "Mask only"],
+            value="Default",
+            label="Output Type"
+        )
+        model_selector = gr.Dropdown(
+            [
+                "u2net",
+                "u2netp",
+                "u2net_human_seg",
+                "u2net_cloth_seg",
+                "silueta",
+                "isnet-general-use",
+                "isnet-anime",
+                "sam",
+                "birefnet-general",
+                "birefnet-general-lite",
+                "birefnet-portrait",
+                "birefnet-dis",
+                "birefnet-hrsod",
+                "birefnet-cod",
+                "birefnet-massive"
+            ],
+            value="isnet-general-use",
+            label="Model Selection"
+        )
+
+    extra = gr.Markdown("## Click on the image to capture coordinates (for SAM model)", visible=False)
+
+    x = gr.Number(label="Mouse X Coordinate", visible=False)
+    y = gr.Number(label="Mouse Y Coordinate", visible=False)
+
+    model_selector.change(show_coords, inputs=model_selector, outputs=[x, y, extra])
+    inputs.select(get_coords, None, [x, y])
+
+
+    gr.Button("Process Image").click(
+        inference,
+        inputs=[inputs, mask_option, model_selector, x, y],
+        outputs=outputs
+    )
+
+    gr.Examples(
+        examples=[
+            ["lion.png", "Default", "u2net", None, None],
+            ["girl.jpg", "Default", "u2net", None, None],
+            ["anime-girl.jpg", "Default", "isnet-anime", None, None]
+        ],
+        inputs=[inputs, mask_option, model_selector, x, y],
+        outputs=outputs
+    )
+    gr.HTML(badge)
+
+app.launch(share=True)