Update app.py

adding image slicer

app.py

@@ -25,11 +25,13 @@
 #######################################################################################
 
 # This file implements an API endpoint for DIS background image removal system.
+# [Self space] - [https://huggingface.co/spaces/leonelhs/removebg]
 #
 # Source code is based on or inspired by several projects.
 # For more details and proper attribution, please refer to the following resources:
 #
 # - [DIS] - [https://github.com/xuebinqin/DIS]
+# - [removebg] - [https://huggingface.co/spaces/gaviego/removebg]
 
 import gradio as gr
 import numpy as np
@@ -39,6 +41,7 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 from torch.autograd import Variable
 from torchvision.transforms.functional import normalize
+from itertools import islice
 
 # project imports
 from models.isnet import ISNetDIS
@@ -54,11 +57,13 @@ net.load_state_dict(torch.load(model_path, map_location=device))
 net.to(device)
 net.eval()
 
-def im_preprocess(im,size):
+def im_preprocess(im, size):
     if len(im.shape) < 3:
         im = im[:, :, np.newaxis]
     if im.shape[2] == 1:
         im = np.repeat(im, 3, axis=2)
+    if im.shape[2] == 4:
+        im = im[:, :, :3]
     im_tensor = torch.tensor(im.copy(), dtype=torch.float32)
     im_tensor = torch.transpose(torch.transpose(im_tensor,1,2),0,1)
     if len(size)<2:
@@ -80,9 +85,8 @@ def predict(image):
         Parameters:
             image (string): File path to the input image.
         Returns:
-            paths (tuple): paths for background-removed image and cutting mask.
+            image (string): paths for image cutting mask.
     """
-
     im_tensor, shapes = im_preprocess(image, [1024, 1024])
     shapes = torch.from_numpy(np.array(shapes)).unsqueeze(0)
 
@@ -101,30 +105,34 @@ def predict(image):
     prediction = (prediction - mi) / (ma - mi)  # max = 1
 
     torch.cuda.empty_cache()
-    mask = (prediction.detach().cpu().numpy() * 255).astype(np.uint8)  # it is the mask we need
+    return (prediction.detach().cpu().numpy() * 255).astype(np.uint8)  # it is the mask we need
 
+def cuts(image):
+    mask = predict(image)
     mask = Image.fromarray(mask).convert('L')
-    image_rgb = Image.fromarray(image).convert("RGB")
-    image_rgb.putalpha(mask)
-    return image_rgb, mask
+    cutted = Image.fromarray(image).convert("RGB")
+    cutted.putalpha(mask)
+    return [image, cutted], mask
 
-article = "<div><center>Unofficial demo from:<a href='https://github.com/xuebinqin/DIS'>DIS</<></center></div>"
 
 with gr.Blocks(title="DIS") as app:
+    navbar = gr.Navbar(visible=True, main_page_name="Workspace")
     gr.Markdown("## Dichotomous Image Segmentation")
     with gr.Row():
         with gr.Column(scale=1):
-            inp = gr.Image(type="numpy", label="Upload Image")
-            btn_predict = gr.Button("Remove background")
+            inp_image = gr.Image(type="numpy", label="Upload Image")
+            btn_predict = gr.Button(variant="primary", value="Remove background")
         with gr.Column(scale=2):
             with gr.Row():
-                with gr.Column(scale=1):
-                    out = gr.Image(type="filepath", label="Output image")
-                    with gr.Accordion("See intermediates", open=False):
-                        out_mask = gr.Image(type="filepath", label="Mask")
+                preview = gr.ImageSlider(type="filepath", label="Comparer")
+
+    btn_predict.click(cuts, inputs=[inp_image], outputs=[preview, inp_image])
+
+with app.route("Readme", "/readme"):
+    with open("README.md") as f:
+        for line in islice(f, 12, None):
+            gr.Markdown(line.strip())
 
-    btn_predict.click(predict, inputs=inp, outputs=[out, out_mask])
-    gr.HTML(article)
 
 app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
 app.queue()