Update app.py

app.py

@@ -1,10 +1,12 @@
 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 io import BytesIO
+from PIL import Image
+import os
 
 # GPU 설정을 CPU로 변경
 birefnet = AutoModelForImageSegmentation.from_pretrained(
@@ -25,11 +27,14 @@ def fn(image):
     im = im.convert("RGB")
     origin = im.copy()
     processed_image = process(im)
-    # Convert processed image to JPEG for download
-    buffered = BytesIO()
-    processed_image.convert("RGB").save(buffered, format="JPEG")
-    buffered.seek(0)
-    return processed_image, buffered
+    
+    # JPG로 변환하여 저장
+    jpg_image = origin.copy()
+    jpg_image = jpg_image.convert("RGB")
+    jpg_path = "output.jpg"
+    jpg_image.save(jpg_path, format="JPEG")
+    
+    return processed_image, jpg_path
 
 def process(image):
     image_size = image.size
@@ -43,30 +48,33 @@ def process(image):
     image.putalpha(mask)
     return image
 
-slider = ImageSlider(label="Processed Image", type="pil")
-download_output = gr.File(label="Download JPG File")
+def process_file(f):
+    name_path = f.rsplit(".", 1)[0] + ".png"
+    im = load_img(f, output_type="pil")
+    im = im.convert("RGB")
+    transparent = process(im)
+    transparent.save(name_path)
+    return name_path
 
+slider1 = ImageSlider(label="Processed Image", type="pil")
 image_upload = gr.Image(label="Upload an image")
+output_download = gr.File(label="Download JPG File")
 
-# 새로운 샘플 이미지 (예: 동일 디렉토리에 위치)
-sample_images = [
-    ["1.png"],
-    ["2.jpg"],
-    ["3.png"]
-]
+# 새로운 샘플 이미지 추가
+sample_images = ["1.png", "2.jpg", "3.png"]
 
-tab = gr.Interface(
+tab1 = gr.Interface(
     fn=fn, 
     inputs=image_upload, 
-    outputs=[slider, download_output], 
+    outputs=[slider1, output_download], 
     examples=sample_images, 
     api_name="image"
 )
 
-demo = gr.TabbedInterface(
-    [tab], 
-    ["Image Upload"], 
-    title="Background Removal Tool"
+demo = gr.Interface(
+    tab1, 
+    title="Background Removal Tool", 
+    description="이미지를 업로드하면 배경이 제거된 이미지를 확인하고 JPG 파일로 다운로드할 수 있습니다."
 )
 
 if __name__ == "__main__":