Update app.py

app.py

@@ -1,12 +1,11 @@
 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êm thư viện uuid để tạo UUID
+import uuid  # Thư viện uuid để tạo UUID
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -25,7 +24,7 @@ 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
 
@@ -39,15 +38,27 @@ def process(image):
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
-    return mask
+
+    # 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):
     session_folder, session_id = create_session_folder()
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     mask = process(im)
-    image_path = os.path.join(session_folder, "mask.png")
-    mask.save(image_path)
+    image_path = os.path.join(session_folder, "transparent_mask.png")
+    mask.save(image_path, "PNG")
     return image_path
 
 image = gr.Image(label="Upload an image")