Update app.py

app.py

@@ -1,13 +1,16 @@
 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
 from PIL import Image
+import io
+import os
+import tempfile
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
-
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
@@ -20,65 +23,58 @@ transform_image = transforms.Compose(
     ]
 )
 
-
 @spaces.GPU
 def fn(image):
-    if image is None or len(image) == 0:
-        return image, None  # 원본 이미지도 반환
-    im = Image.open(image).convert("RGB")
+    if image is None:
+        return None, None
+    im = load_img(image, output_type="pil")
+    im = im.convert("RGB")
     image_size = im.size
     origin = im.copy()
-    input_images = transform_image(im).unsqueeze(0).to("cuda")
-    # Prediction
+    image = load_img(im)
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    # 예측
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
-    im.putalpha(mask)
-    return im, origin  # 변환된 이미지와 원본 이미지 반환
-
+    image.putalpha(mask)
+    return image, origin
 
 def save_image(image):
-    if image is not None:
-        image.save("output.png")
-        return "output.png"
-    return None
+    if image is None:
+        return None
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
+        image.save(temp_file, format="PNG")
+    return temp_file.name
 
+def process_and_download(input_image):
+    result, original = fn(input_image)
+    if result is None:
+        return None, None
+    result_path = save_image(result)
+    original_path = save_image(original)
+    return [result_path, original_path], result_path
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column(scale=1):
-            image = gr.Image(label="Upload an image")
-            text = gr.Textbox(label="Paste an image URL")
-            download_button = gr.Button("Download Image")
-            output_file = gr.File()
-        
-        with gr.Column(scale=2):
-            slider1 = ImageSlider(label="Processed Image", type="pil")
-            slider2 = ImageSlider(label="Original Image", type="pil")
+image = gr.Image(label="이미지 업로드")
+slider = ImageSlider(label="배경 제거 결과", type="filepath")
+png_output = gr.File(label="PNG 다운로드")
 
-    # 스페이스에 있는 예제 이미지 파일 경로
-    example_image1 = "example_images/example1.jpg"
-    example_image2 = "example_images/example2.jpg"
-    example_image3 = "example_images/example3.jpg"
-    
-    with gr.Tab("Image Upload"):
-        tab1 = gr.Interface(
-            fn, inputs=image, outputs=[slider1, output_file], 
-            examples=[example_image1, example_image2, example_image3], api_name="image"
-        )
+examples = [
+    os.path.join(os.path.dirname(__file__), "예제1.png"),
+    os.path.join(os.path.dirname(__file__), "예제2.png"),
+    os.path.join(os.path.dirname(__file__), "예제3.png")
+]
 
-    with gr.Tab("Image URL"):
-        tab2 = gr.Interface(
-            fn, inputs=text, outputs=[slider2, output_file], 
-            examples=[example_image1, example_image2, example_image3], api_name="text"
-        )
-
-    def process_download(image):
-        return save_image(image[0])
-
-    download_button.click(process_download, inputs=slider1, outputs=output_file)
+demo = gr.Interface(
+    process_and_download,
+    inputs=image,
+    outputs=[slider, png_output],
+    examples=examples,
+    title="배경 제거",
+    description="이미지를 업로드하면 BiRefNet 모델을 사용하여 배경을 제거합니다. 결과를 PNG 파일로 다운로드할 수 있습니다."
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()