Update app.py

app.py

@@ -1,11 +1,12 @@
 import os
 import gradio as gr
+from gradio_imageslider import ImageSlider
 from loadimg import load_img
 import spaces
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
-import uuid  # Thư viện uuid để tạo UUID
+import uuid  # Thêm thư viện uuid để tạo UUID
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -24,10 +25,20 @@ transform_image = transforms.Compose(
 def create_session_folder():
     """Tạo thư mục duy nhất dựa trên UUID cho mỗi phiên làm việc."""
     session_id = str(uuid.uuid4())  # Tạo UUID duy nhất
-    session_folder = os.path.join("output_images", session_id)
+    session_folder = os.path.join('output_images', session_id)
     os.makedirs(session_folder, exist_ok=True)
     return session_folder, session_id
 
+def fn(image):
+    session_folder, session_id = create_session_folder()
+    im = load_img(image, output_type="pil")
+    im = im.convert("RGB")
+    origin = im.copy()
+    image = process(im)    
+    image_path = os.path.join(session_folder, "no_bg_image.png")
+    image.save(image_path)
+    return (image, origin), image_path
+
 @spaces.GPU
 def process(image):
     image_size = image.size
@@ -38,38 +49,38 @@ def process(image):
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
-
-    # Tạo ảnh PNG trong suốt
-    transparent_mask = mask.convert("L")  # Chuyển thành ảnh grayscale
-    result = mask.convert("RGBA")
-    pixels = result.load()
-
-    for i in range(result.width):
-        for j in range(result.height):
-            # Đặt màu đen với độ trong suốt
-            alpha = transparent_mask.getpixel((i, j))
-            pixels[i, j] = (0, 0, 0, alpha)
-
-    return result
-
-def fn(image):
+    image.putalpha(mask)
+    return image
+  
+def process_file(f):
     session_folder, session_id = create_session_folder()
-    im = load_img(image, output_type="pil")
+    name_path = os.path.join(session_folder, os.path.basename(f).rsplit(".", 1)[0] + ".png")
+    im = load_img(f, output_type="pil")
     im = im.convert("RGB")
-    mask = process(im)
-    image_path = os.path.join(session_folder, "transparent_mask.png")
-    mask.save(image_path, "PNG")
-    return image_path
+    transparent = process(im)
+    transparent.save(name_path)
+    return name_path
 
+slider1 = ImageSlider(label="RMBG-2.0", type="pil")
+slider2 = ImageSlider(label="RMBG-2.0", type="pil")
 image = gr.Image(label="Upload an image")
-png_file = gr.File(label="Output mask PNG file")
+image2 = gr.Image(label="Upload an image", type="filepath")
+text = gr.Textbox(label="Paste an image URL")
+png_file = gr.File(label="output png file")
+
+chameleon = load_img("giraffe.jpg", output_type="pil")
+
+url = "http://farm9.staticflickr.com/8488/8228323072_76eeddfea3_z.jpg"
 
 tab1 = gr.Interface(
-    fn, inputs=image, outputs=png_file, examples=[], api_name="image"
+    fn, inputs=image, outputs=[slider1, gr.File(label="output png file")], examples=[chameleon], api_name="image"
 )
 
+tab2 = gr.Interface(fn, inputs=text, outputs=[slider2, gr.File(label="output png file")], examples=[url], api_name="text")
+tab3 = gr.Interface(process_file, inputs=image2, outputs=png_file, examples=["giraffe.jpg"], api_name="png")
+
 demo = gr.TabbedInterface(
-    [tab1], ["Upload Image"], title="RMBG-2.0 for Background Removal"
+    [tab1, tab2], ["input image", "input url"], title="RMBG-2.0 for background removal"
 )
 
 if __name__ == "__main__":