Update app.py

app.py

@@ -1,15 +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
 
 # GPU 설정을 CPU로 변경
-# GPU 설정을 삭제하거나 "cuda"를 "cpu"로 변경
-# torch.set_float32_matmul_precision("high")는 CPU에선 필요 없음.
-
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
@@ -28,10 +25,11 @@ def fn(image):
     im = im.convert("RGB")
     origin = im.copy()
     processed_image = process(im)
-    return (processed_image, origin)
-
-# @spaces.GPU 데코레이터 제거
-# CPU 환경에서 동작하도록 설정
+    # Convert processed image to JPEG
+    buffered = BytesIO()
+    processed_image.convert("RGB").save(buffered, format="JPEG")
+    buffered.seek(0)
+    return processed_image, buffered
 
 def process(image):
     image_size = image.size
@@ -45,31 +43,41 @@ def process(image):
     image.putalpha(mask)
     return image
 
-def process_file(f):
-    name_path = f.rsplit(".", 1)[0] + ".png"
+def process_download(f):
+    name_path = f.rsplit(".", 1)[0] + ".jpg"
     im = load_img(f, output_type="pil")
     im = im.convert("RGB")
     transparent = process(im)
-    transparent.save(name_path)
-    return name_path
+    # Convert to JPEG
+    buffered = BytesIO()
+    transparent.convert("RGB").save(buffered, format="JPEG")
+    buffered.seek(0)
+    return buffered
+
+slider = ImageSlider(label="Processed Image", type="pil")
+download_output = gr.File(label="Download JPG File")
 
-slider1 = ImageSlider(label="Processed Image", type="pil")
-slider2 = ImageSlider(label="Processed Image from URL", type="pil")
 image_upload = gr.Image(label="Upload an image")
-image_file_upload = gr.Image(label="Upload an image", type="filepath")
-url_input = gr.Textbox(label="Paste an image URL")
-output_file = gr.File(label="Output PNG File")
 
-# Example images
-chameleon = load_img("butterfly.jpg", output_type="pil")
-url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
+# 새로운 샘플 이미지
+sample_images = [
+    "1.png",
+    "2.jpg",
+    "3.png"
+]
 
-tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
-tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
-tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
+tab = gr.Interface(
+    fn=fn, 
+    inputs=image_upload, 
+    outputs=[slider, download_output], 
+    examples=sample_images, 
+    api_name="image"
+)
 
 demo = gr.TabbedInterface(
-    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
+    [tab], 
+    ["Image Upload"], 
+    title="Background Removal Tool"
 )
 
 if __name__ == "__main__":